Implement gel.Record type (#560)

It's returned from `client.query_sql()` and represents
SQL rows.

Currently it supports:

* integer indexing - fetch columns by position
* string indexing - fetch columns by name
* as_dict() method - render record as a dict

Author: 1st1

gel/__init__.py

@@ -24,7 +24,7 @@
 from gel.datatypes.datatypes import (
     Tuple, NamedTuple, EnumValue, RelativeDuration, DateDuration, ConfigMemory
 )
-from gel.datatypes.datatypes import Set, Object, Array
+from gel.datatypes.datatypes import Record, Set, Object, Array
 from gel.datatypes.range import Range, MultiRange
 
 from .abstract import (
@@ -68,6 +68,7 @@
     "RetryCondition",
     "RetryOptions",
     "Set",
+    "Record",
     "State",
     "TransactionOptions",
     "Tuple",

gel/abstract.py

@@ -22,6 +22,7 @@
 import dataclasses
 import typing
 
+from . import datatypes
 from . import describe
 from . import enums
 from . import options
@@ -278,7 +279,7 @@ def query_required_single_json(self, query: str, *args, **kwargs) -> str:
             annotations=self._get_annotations(),
         ))
 
-    def query_sql(self, query: str, *args, **kwargs) -> typing.Any:
+    def query_sql(self, query: str, *args, **kwargs) -> list[datatypes.Record]:
         return self._query(QueryContext(
             query=QueryWithArgs(
                 query,

gel/datatypes/datatypes.h

@@ -70,6 +70,10 @@ typedef enum {
     L_LINK
 } edge_attr_lookup_t;
 
+
+#define EdgeRecordDesc_GET_NAMES(o) (((EdgeRecordDescObject *)(o))->names)
+
+
 PyObject * EdgeRecordDesc_InitType(void);
 PyObject * EdgeRecordDesc_New(PyObject *, PyObject *, PyObject *);
 PyObject * EdgeRecordDesc_PointerName(PyObject *, Py_ssize_t);
@@ -85,6 +89,8 @@ edge_attr_lookup_t EdgeRecordDesc_Lookup(PyObject *, PyObject *, Py_ssize_t *);
 PyObject * EdgeRecordDesc_List(PyObject *, uint8_t, uint8_t);
 PyObject * EdgeRecordDesc_GetDataclassFields(PyObject *);
 
+PyObject * EdgeRecordDesc_GetNames(PyObject *);
+
 
 /* === gel.NamedTuple ==================================== */
 
@@ -122,6 +128,29 @@ PyObject * EdgeObject_GetRecordDesc(PyObject *);
 int EdgeObject_SetItem(PyObject *, Py_ssize_t, PyObject *);
 PyObject * EdgeObject_GetItem(PyObject *, Py_ssize_t);
 
-PyObject * EdgeObject_GetID(PyObject *ob);
+
+/* === edgedb.Record ======================================== */
+
+#define EDGE_RECORD_FREELIST_SIZE 2000
+#define EDGE_RECORD_FREELIST_MAXSAVE 20
+
+extern PyTypeObject EdgeRecord_Type;
+
+#define EdgeRecord_Check(d) (Py_TYPE(d) == &EdgeRecord_Type)
+
+typedef struct {
+    PyObject_VAR_HEAD
+    PyObject *weakreflist;
+    PyObject *desc;
+    Py_hash_t cached_hash;
+    PyObject *ob_item[1];
+} EdgeRecord;
+
+PyObject * EdgeRecord_InitType(void);
+PyObject * EdgeRecord_New(PyObject *);
+PyObject * EdgeRecord_GetRecordDesc(PyObject *);
+
+int EdgeRecord_SetItem(PyObject *, Py_ssize_t, PyObject *);
+PyObject * EdgeRecord_GetItem(PyObject *, Py_ssize_t);
 
 #endif

gel/datatypes/datatypes.pxd

@@ -58,6 +58,10 @@ cdef extern from "datatypes.h":
     int EdgeObject_SetItem(object, Py_ssize_t, object) except -1
     object EdgeObject_GetRecordDesc(object)
 
+    object EdgeRecord_InitType()
+    object EdgeRecord_New(object);
+    int EdgeRecord_SetItem(object, Py_ssize_t, object) except -1
+    object EdgeRecord_GetRecordDesc(object)
 
 cdef record_desc_new(object names, object flags, object cards)
 cdef record_desc_pointer_name(object desc, Py_ssize_t pos)
@@ -69,5 +73,7 @@ cdef namedtuple_new(object namedtuple_type)
 cdef namedtuple_type_new(object desc)
 cdef object_new(object desc)
 cdef object_set(object tuple, Py_ssize_t pos, object elem)
+cdef record_new(object desc)
+cdef record_set(object obj, Py_ssize_t pos, object elem)
 
 cdef extern cpython.PyObject* at_sign_ptr

gel/datatypes/datatypes.pyx

@@ -30,6 +30,7 @@ _RecordDescriptor = EdgeRecordDesc_InitType()
 Tuple = tuple
 NamedTuple = EdgeNamedTuple_InitType()
 Object = EdgeObject_InitType()
+Record = EdgeRecord_InitType()
 Set = list
 Array = list
 
@@ -132,6 +133,11 @@ cdef namedtuple_type_new(object desc):
 cdef object_new(object desc):
     return EdgeObject_New(desc)
 
-
 cdef object_set(object obj, Py_ssize_t pos, object elem):
     EdgeObject_SetItem(obj, pos, elem)
+
+cdef record_new(object desc):
+    return EdgeRecord_New(desc)
+
+cdef record_set(object obj, Py_ssize_t pos, object elem):
+    EdgeRecord_SetItem(obj, pos, elem)

gel/datatypes/internal.h

@@ -37,15 +37,21 @@ PyObject * _EdgeGeneric_RenderObject(PyObject *obj);
 int _EdgeGeneric_RenderValues(
     _PyUnicodeWriter *, PyObject *, PyObject **, Py_ssize_t);
 
-int _EdgeGeneric_RenderItems(_PyUnicodeWriter *,
-                             PyObject *, PyObject *,
-                             PyObject **, Py_ssize_t, int, int);
-
 PyObject * _EdgeGeneric_RichCompareValues(PyObject **, Py_ssize_t,
                                           PyObject **, Py_ssize_t,
                                           int);
 
 
+#define EDGE_RENDER_NAMES       0x1
+#define EDGE_RENDER_LINK_PROPS  0x2
+#define EDGE_RENDER_IMPLICIT    0x4
+#define EDGE_RENDER_DEFAULT     0
+
+int _EdgeGeneric_RenderItems(_PyUnicodeWriter *,
+                             PyObject *, PyObject *,
+                             PyObject **, Py_ssize_t,
+                             int);
+
 #ifndef _PyList_CAST
 #  define _PyList_CAST(op) (assert(PyList_Check(op)), (PyListObject *)(op))
 #endif

gel/datatypes/namedtuple.c

@@ -43,6 +43,7 @@ PyObject *
 EdgeNamedTuple_New(PyObject *type)
 {
     assert(init_type_called);
+    assert(PyType_Check(type));
 
     PyObject *desc = EdgeNamedTuple_Type_DESC(type);
     if (desc == NULL || !EdgeRecordDesc_Check(desc)) {
@@ -77,7 +78,7 @@ EdgeNamedTuple_New(PyObject *type)
             _EDGE_NAMED_TUPLE_FL_NUM_FREE[size]--;
             _Py_NewReference((PyObject *)nt);
             Py_INCREF(type);
-            Py_SET_TYPE(nt, type);
+            Py_SET_TYPE(nt, (PyTypeObject*)type);
         }
     } else {
         if (
@@ -88,7 +89,7 @@ EdgeNamedTuple_New(PyObject *type)
             PyErr_NoMemory();
             return NULL;
         }
-        nt = PyObject_GC_NewVar(PyTupleObject, type, size);
+        nt = PyObject_GC_NewVar(PyTupleObject, (PyTypeObject*)type, size);
         if (nt == NULL) {
             return NULL;
         }
@@ -197,9 +198,9 @@ namedtuple_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) {
         goto fail;
     }
 
-    type = EdgeNamedTuple_Type_New(desc);
-    o = (PyTupleObject *)EdgeNamedTuple_New(type);
-    Py_CLEAR(type);  // the type is now referenced by the object
+    PyObject *new_type = EdgeNamedTuple_Type_New(desc);
+    o = (PyTupleObject *)EdgeNamedTuple_New(new_type);
+    Py_CLEAR(new_type);  // the type is now referenced by the object
 
     if (o == NULL) {
         goto fail;
@@ -236,7 +237,7 @@ namedtuple_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) {
 
 static PyObject *
 namedtuple_derived_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) {
-    PyTupleObject *o = (PyTupleObject *)EdgeNamedTuple_New(type);
+    PyTupleObject *o = (PyTupleObject *)EdgeNamedTuple_New((PyObject*)type);
     if (o == NULL) {
         goto fail;
     }
@@ -356,7 +357,8 @@ namedtuple_repr(PyTupleObject *o)
     if (_EdgeGeneric_RenderItems(&writer,
                                  (PyObject *)o,
                                  EdgeNamedTuple_Type_DESC(Py_TYPE(o)),
-                                 o->ob_item, Py_SIZE(o), 0, 0) < 0)
+                                 o->ob_item, Py_SIZE(o),
+                                 EDGE_RENDER_NAMES) < 0)
     {
         goto error;
     }
@@ -458,7 +460,7 @@ EdgeNamedTuple_Type_New(PyObject *desc)
 
     // store `_fields` for collections.namedtuple duck-typing
     Py_ssize_t size = EdgeRecordDesc_GetSize(desc);
-    PyTupleObject *fields = PyTuple_New(size);
+    PyTupleObject *fields = (PyTupleObject *)PyTuple_New(size);
     if (fields == NULL) {
         goto fail;
     }
@@ -470,11 +472,11 @@ EdgeNamedTuple_Type_New(PyObject *desc)
         }
         PyTuple_SET_ITEM(fields, i, name);
     }
-    if (PyDict_SetItemString(rv->tp_dict, "_fields", fields) < 0) {
+    if (PyDict_SetItemString(rv->tp_dict, "_fields", (PyObject*)fields) < 0) {
         goto fail;
     }
 
-    return rv;
+    return (PyObject*)rv;
 
 fail:
     Py_DECREF(rv);

gel/datatypes/object.c

@@ -124,25 +124,6 @@ EdgeObject_GetItem(PyObject *ob, Py_ssize_t i)
 }
 
 
-PyObject *
-EdgeObject_GetID(PyObject *ob)
-{
-    assert(EdgeObject_Check(ob));
-    EdgeObject *o = (EdgeObject *)ob;
-    Py_ssize_t i = EdgeRecordDesc_IDPos(o->desc);
-    if (i < 0) {
-        Py_RETURN_NONE;
-    }
-    if (i >= Py_SIZE(o)) {
-        PyErr_BadInternalCall();
-        return NULL;
-    }
-    PyObject *el = EdgeObject_GET_ITEM(o, i);
-    Py_INCREF(el);
-    return el;
-}
-
-
 static void
 object_dealloc(EdgeObject *o)
 {
@@ -282,9 +263,11 @@ object_repr(EdgeObject *o)
         goto error;
     }
 
-    if (_EdgeGeneric_RenderItems(&writer,
-                                 (PyObject *)o, o->desc,
-                                 o->ob_item, Py_SIZE(o), 1, 0) < 0)
+    if (_EdgeGeneric_RenderItems(
+            &writer,
+            (PyObject *)o, o->desc,
+            o->ob_item, Py_SIZE(o),
+            EDGE_RENDER_NAMES | EDGE_RENDER_LINK_PROPS) < 0)
     {
         goto error;
     }