Skip to content

Commit

Permalink
Implement gel.Record type (#560)
Browse files Browse the repository at this point in the history
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
  • Loading branch information
1st1 authored Dec 12, 2024
1 parent 1bc75c4 commit ff3aad2
Show file tree
Hide file tree
Showing 16 changed files with 527 additions and 55 deletions.
3 changes: 2 additions & 1 deletion gel/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -68,6 +68,7 @@
"RetryCondition",
"RetryOptions",
"Set",
"Record",
"State",
"TransactionOptions",
"Tuple",
Expand Down
3 changes: 2 additions & 1 deletion gel/abstract.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import dataclasses
import typing

from . import datatypes
from . import describe
from . import enums
from . import options
Expand Down Expand Up @@ -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,
Expand Down
31 changes: 30 additions & 1 deletion gel/datatypes/datatypes.h
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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 ==================================== */

Expand Down Expand Up @@ -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
6 changes: 6 additions & 0 deletions gel/datatypes/datatypes.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
8 changes: 7 additions & 1 deletion gel/datatypes/datatypes.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ _RecordDescriptor = EdgeRecordDesc_InitType()
Tuple = tuple
NamedTuple = EdgeNamedTuple_InitType()
Object = EdgeObject_InitType()
Record = EdgeRecord_InitType()
Set = list
Array = list

Expand Down Expand Up @@ -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)
14 changes: 10 additions & 4 deletions gel/datatypes/internal.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
22 changes: 12 additions & 10 deletions gel/datatypes/namedtuple.c
Original file line number Diff line number Diff line change
Expand Up @@ -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)) {
Expand Down Expand Up @@ -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 (
Expand All @@ -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;
}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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;
}
Expand All @@ -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);
Expand Down
27 changes: 5 additions & 22 deletions gel/datatypes/object.c
Original file line number Diff line number Diff line change
Expand Up @@ -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)
{
Expand Down Expand Up @@ -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;
}
Expand Down
Loading

0 comments on commit ff3aad2

Please sign in to comment.