From ff3aad258fbcfe8d2a082da3ea610d5e7e82b01c Mon Sep 17 00:00:00 2001 From: Yury Selivanov Date: Thu, 12 Dec 2024 16:56:00 -0600 Subject: [PATCH] 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 --- gel/__init__.py | 3 +- gel/abstract.py | 3 +- gel/datatypes/datatypes.h | 31 +++- gel/datatypes/datatypes.pxd | 6 + gel/datatypes/datatypes.pyx | 8 +- gel/datatypes/internal.h | 14 +- gel/datatypes/namedtuple.c | 22 +-- gel/datatypes/object.c | 27 +-- gel/datatypes/record.c | 295 +++++++++++++++++++++++++++++++++ gel/datatypes/repr.c | 31 ++-- gel/protocol/codecs/codecs.pxd | 1 + gel/protocol/codecs/codecs.pyx | 27 +++ gel/protocol/codecs/record.pxd | 23 +++ gel/protocol/codecs/record.pyx | 79 +++++++++ setup.py | 1 + tests/test_async_query.py | 11 ++ 16 files changed, 527 insertions(+), 55 deletions(-) create mode 100644 gel/datatypes/record.c create mode 100644 gel/protocol/codecs/record.pxd create mode 100644 gel/protocol/codecs/record.pyx diff --git a/gel/__init__.py b/gel/__init__.py index b9b9dc83..80f953ff 100644 --- a/gel/__init__.py +++ b/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", diff --git a/gel/abstract.py b/gel/abstract.py index abffbd23..22264405 100644 --- a/gel/abstract.py +++ b/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, diff --git a/gel/datatypes/datatypes.h b/gel/datatypes/datatypes.h index d1077477..b85c8056 100644 --- a/gel/datatypes/datatypes.h +++ b/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 diff --git a/gel/datatypes/datatypes.pxd b/gel/datatypes/datatypes.pxd index ec83bf39..3e2298bd 100644 --- a/gel/datatypes/datatypes.pxd +++ b/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 diff --git a/gel/datatypes/datatypes.pyx b/gel/datatypes/datatypes.pyx index e453b945..751c1ff2 100644 --- a/gel/datatypes/datatypes.pyx +++ b/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) diff --git a/gel/datatypes/internal.h b/gel/datatypes/internal.h index c6754f60..01c91b03 100644 --- a/gel/datatypes/internal.h +++ b/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 diff --git a/gel/datatypes/namedtuple.c b/gel/datatypes/namedtuple.c index 251a0e92..08395713 100644 --- a/gel/datatypes/namedtuple.c +++ b/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); diff --git a/gel/datatypes/object.c b/gel/datatypes/object.c index a5965bc1..955680e6 100644 --- a/gel/datatypes/object.c +++ b/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; } diff --git a/gel/datatypes/record.c b/gel/datatypes/record.c new file mode 100644 index 00000000..d2480754 --- /dev/null +++ b/gel/datatypes/record.c @@ -0,0 +1,295 @@ +#include "pythoncapi_compat.h" + +/* +* This source file is part of the EdgeDB open source project. +* +* Copyright 2024-present MagicStack Inc. and the EdgeDB authors. +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ + + +#include "datatypes.h" +#include "freelist.h" +#include "internal.h" + + +static int init_type_called = 0; + +EDGE_SETUP_FREELIST( + EDGE_RECORD, + EdgeRecord, + EDGE_RECORD_FREELIST_MAXSAVE, + EDGE_RECORD_FREELIST_SIZE) + + +#define EdgeRecord_GET_ITEM(op, i) \ + (((EdgeRecord *)(op))->ob_item[i]) +#define EdgeRecord_SET_ITEM(op, i, v) \ + (((EdgeRecord *)(op))->ob_item[i] = v) + + +PyObject * +EdgeRecord_New(PyObject *desc) +{ + assert(init_type_called); + + if (desc == NULL || !EdgeRecordDesc_Check(desc)) { + PyErr_BadInternalCall(); + return NULL; + } + + Py_ssize_t size = EdgeRecordDesc_GetSize(desc); + + if (size > EDGE_MAX_TUPLE_SIZE) { + PyErr_Format( + PyExc_ValueError, + "Cannot create Object with more than %d elements", + EDGE_MAX_TUPLE_SIZE); + return NULL; + } + + EdgeRecord *o = NULL; + EDGE_NEW_WITH_FREELIST(EDGE_RECORD, EdgeRecord, + &EdgeRecord_Type, o, size); + assert(o != NULL); + assert(Py_SIZE(o) == size); + assert(EdgeRecord_Check(o)); + + o->weakreflist = NULL; + + Py_INCREF(desc); + o->desc = desc; + + o->cached_hash = -1; + + PyObject_GC_Track(o); + return (PyObject *)o; +} + + +PyObject * +EdgeRecord_GetRecordDesc(PyObject *o) +{ + if (!EdgeRecord_Check(o)) { + PyErr_Format( + PyExc_TypeError, + "an instance of edgedb.Object expected"); + return NULL; + } + + PyObject *desc = ((EdgeRecord *)o)->desc; + Py_INCREF(desc); + return desc; +} + + +int +EdgeRecord_SetItem(PyObject *ob, Py_ssize_t i, PyObject *el) +{ + assert(EdgeRecord_Check(ob)); + EdgeRecord *o = (EdgeRecord *)ob; + assert(i >= 0); + assert(i < Py_SIZE(o)); + Py_INCREF(el); + EdgeRecord_SET_ITEM(o, i, el); + return 0; +} + + +PyObject * +EdgeRecord_GetItem(PyObject *ob, Py_ssize_t i) +{ + assert(EdgeRecord_Check(ob)); + EdgeRecord *o = (EdgeRecord *)ob; + if (i < 0 || i >= Py_SIZE(o)) { + PyErr_SetString(PyExc_IndexError, "record index out of range"); + return NULL; + } + PyObject *el = EdgeRecord_GET_ITEM(o, i); + Py_INCREF(el); + return el; +} + + +static void +record_dealloc(EdgeRecord *o) +{ + PyObject_GC_UnTrack(o); + if (o->weakreflist != NULL) { + PyObject_ClearWeakRefs((PyObject*)o); + } + Py_CLEAR(o->desc); + o->cached_hash = -1; + Py_TRASHCAN_BEGIN(o, record_dealloc); + EDGE_DEALLOC_WITH_FREELIST(EDGE_RECORD, EdgeRecord, o); + Py_TRASHCAN_END(o); +} + + +static int +record_traverse(EdgeRecord *o, visitproc visit, void *arg) +{ + Py_VISIT(o->desc); + + Py_ssize_t i; + for (i = Py_SIZE(o); --i >= 0;) { + if (o->ob_item[i] != NULL) { + Py_VISIT(o->ob_item[i]); + } + } + return 0; +} + + +static PyObject * +record_getitem(EdgeRecord *o, PyObject *name) +{ + Py_ssize_t pos; + + if (PyLong_Check(name)) { + pos = PyLong_AsSsize_t(name); + if (pos == -1 && PyErr_Occurred()) { + return NULL; + } + return EdgeRecord_GetItem((PyObject*)o, pos); + } else { + edge_attr_lookup_t ret = EdgeRecordDesc_Lookup(o->desc, name, &pos); + + switch (ret) { + case L_ERROR: + return NULL; + + case L_NOT_FOUND: + PyErr_SetObject(PyExc_KeyError, name); + return NULL; + + case L_LINK: + case L_LINKPROP: + /* isn't possible for records */ + PyErr_BadInternalCall(); + return NULL; + + case L_PROPERTY: { + PyObject *val = EdgeRecord_GET_ITEM(o, pos); + Py_INCREF(val); + return val; + } + + default: + abort(); + } + } +} + + +static PyObject * +record_repr(EdgeRecord *o) +{ + _PyUnicodeWriter writer; + _PyUnicodeWriter_Init(&writer); + writer.overallocate = 1; + + if (_PyUnicodeWriter_WriteASCIIString(&writer, "Record{", 4) < 0) { + goto error; + } + + if (_EdgeGeneric_RenderItems(&writer, + (PyObject *)o, o->desc, + o->ob_item, Py_SIZE(o), + EDGE_RENDER_DEFAULT) < 0) + { + goto error; + } + + if (_PyUnicodeWriter_WriteChar(&writer, '}') < 0) { + goto error; + } + + return _PyUnicodeWriter_Finish(&writer); + +error: + _PyUnicodeWriter_Dealloc(&writer); + return NULL; +} + + +static PyObject * +record_as_dict(EdgeRecord *o, PyObject *args) +{ + if (!EdgeRecord_Check(o)) { + PyErr_BadInternalCall(); + return NULL; + } + + PyObject *ret = PyDict_New(); + if (ret == NULL) { + return NULL; + } + + assert(EdgeRecordDesc_Check(o->desc)); + PyObject *names = EdgeRecordDesc_GET_NAMES(o->desc); /* borrowed */ + + for (Py_ssize_t i = 0; i < PyTuple_GET_SIZE(names); i++) { + PyObject *name = PyTuple_GET_ITEM(names, i); /* borrowed */ + PyObject *value = EdgeRecord_GET_ITEM(o, i); /* borrowed */ + if (PyDict_SetItem(ret, name, value)) { + goto err; + } + } + + return ret; + +err: + Py_DECREF(ret); + return NULL; +} + + +static PyMappingMethods record_as_mapping = { + .mp_subscript = (binaryfunc)record_getitem, +}; + + +static PyMethodDef record_methods[] = { + {"as_dict", (PyCFunction)record_as_dict, METH_NOARGS, NULL}, + {NULL, NULL} +}; + + +PyTypeObject EdgeRecord_Type = { + PyVarObject_HEAD_INIT(NULL, 0) + "edgedb.Record", + .tp_basicsize = sizeof(EdgeRecord) - sizeof(PyObject *), + .tp_itemsize = sizeof(PyObject *), + .tp_dealloc = (destructor)record_dealloc, + .tp_as_mapping = &record_as_mapping, + .tp_methods = record_methods, + .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, + .tp_traverse = (traverseproc)record_traverse, + .tp_free = PyObject_GC_Del, + .tp_repr = (reprfunc)record_repr, + .tp_weaklistoffset = offsetof(EdgeRecord, weakreflist), +}; + + +PyObject * +EdgeRecord_InitType(void) +{ + if (PyType_Ready(&EdgeRecord_Type) < 0) { + return NULL; + } + + init_type_called = 1; + return (PyObject *)&EdgeRecord_Type; +} diff --git a/gel/datatypes/repr.c b/gel/datatypes/repr.c index 752a5fa0..42a010e0 100644 --- a/gel/datatypes/repr.c +++ b/gel/datatypes/repr.c @@ -81,8 +81,7 @@ int _EdgeGeneric_RenderItems(_PyUnicodeWriter *writer, PyObject *host, PyObject *desc, PyObject **items, Py_ssize_t len, - int include_link_props, - int include_implicit) + int render_flags) { assert(EdgeRecordDesc_GetSize(desc) == len); @@ -113,12 +112,12 @@ _EdgeGeneric_RenderItems(_PyUnicodeWriter *writer, goto error; } - if (is_implicit && !include_implicit) { + if (is_implicit && !(render_flags & EDGE_RENDER_IMPLICIT)) { continue; } if (is_linkprop) { - if (!include_link_props) { + if (!(render_flags & EDGE_RENDER_LINK_PROPS)) { continue; } } @@ -128,19 +127,21 @@ _EdgeGeneric_RenderItems(_PyUnicodeWriter *writer, goto error; } - item_name = EdgeRecordDesc_PointerName(desc, i); - if (item_name == NULL) { - goto error; - } - assert(PyUnicode_CheckExact(item_name)); + if (render_flags & EDGE_RENDER_NAMES) { + item_name = EdgeRecordDesc_PointerName(desc, i); + if (item_name == NULL) { + goto error; + } + assert(PyUnicode_CheckExact(item_name)); - if (_PyUnicodeWriter_WriteStr(writer, item_name) < 0) { - goto error; - } - Py_CLEAR(item_name); + if (_PyUnicodeWriter_WriteStr(writer, item_name) < 0) { + goto error; + } + Py_CLEAR(item_name); - if (_PyUnicodeWriter_WriteASCIIString(writer, " := ", 4) < 0) { - goto error; + if (_PyUnicodeWriter_WriteASCIIString(writer, " := ", 4) < 0) { + goto error; + } } if (_PyUnicodeWriter_WriteStr(writer, item_repr) < 0) { diff --git a/gel/protocol/codecs/codecs.pxd b/gel/protocol/codecs/codecs.pxd index ab11d5f5..53affecc 100644 --- a/gel/protocol/codecs/codecs.pxd +++ b/gel/protocol/codecs/codecs.pxd @@ -26,6 +26,7 @@ include "./array.pxd" include "./range.pxd" include "./set.pxd" include "./enum.pxd" +include "./record.pxd" cdef class CodecsRegistry: diff --git a/gel/protocol/codecs/codecs.pyx b/gel/protocol/codecs/codecs.pyx index 62a0d81e..66358ca7 100644 --- a/gel/protocol/codecs/codecs.pyx +++ b/gel/protocol/codecs/codecs.pyx @@ -40,6 +40,7 @@ include "./array.pyx" include "./range.pyx" include "./set.pyx" include "./enum.pyx" +include "./record.pyx" DEF CTYPE_SET = 0 @@ -55,6 +56,7 @@ DEF CTYPE_RANGE = 9 DEF CTYPE_OBJECT = 10 DEF CTYPE_COMPOUND = 11 DEF CTYPE_MULTIRANGE = 12 +DEF CTYPE_SQL_ROW = 13 DEF CTYPE_ANNO_TYPENAME = 255 DEF _CODECS_BUILD_CACHE_SIZE = 200 @@ -169,6 +171,12 @@ cdef class CodecsRegistry: elif t == CTYPE_MULTIRANGE: frb_read(spec, 2) + elif t == CTYPE_SQL_ROW: + els = hton.unpack_int16(frb_read(spec, 2)) + for i in range(els): + str_len = hton.unpack_uint32(frb_read(spec, 4)) + frb_read(spec, str_len + 2) + elif t == CTYPE_ENUM: els = hton.unpack_int16(frb_read(spec, 2)) for i in range(els): @@ -380,6 +388,25 @@ cdef class CodecsRegistry: res = NamedTupleCodec.new(tid, names, codecs) res.type_name = type_name + elif t == CTYPE_SQL_ROW: + els = hton.unpack_int16(frb_read(spec, 2)) + codecs = cpython.PyTuple_New(els) + names = cpython.PyTuple_New(els) + for i in range(els): + str_len = hton.unpack_uint32(frb_read(spec, 4)) + name = cpythonx.PyUnicode_FromStringAndSize( + frb_read(spec, str_len), str_len) + pos = hton.unpack_int16(frb_read(spec, 2)) + + cpython.Py_INCREF(name) + cpython.PyTuple_SetItem(names, i, name) + + sub_codec = codecs_list[pos] + cpython.Py_INCREF(sub_codec) + cpython.PyTuple_SetItem(codecs, i, sub_codec) + + res = RecordCodec.new(tid, names, codecs) + elif t == CTYPE_ENUM: if protocol_version >= (2, 0): str_len = hton.unpack_uint32(frb_read(spec, 4)) diff --git a/gel/protocol/codecs/record.pxd b/gel/protocol/codecs/record.pxd new file mode 100644 index 00000000..e00e5a45 --- /dev/null +++ b/gel/protocol/codecs/record.pxd @@ -0,0 +1,23 @@ +# +# This source file is part of the EdgeDB open source project. +# +# Copyright 2016-present MagicStack Inc. and the EdgeDB authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +@cython.final +cdef class RecordCodec(BaseNamedRecordCodec): + @staticmethod + cdef BaseCodec new(bytes tid, tuple names, tuple codecs) diff --git a/gel/protocol/codecs/record.pyx b/gel/protocol/codecs/record.pyx new file mode 100644 index 00000000..dcf6578a --- /dev/null +++ b/gel/protocol/codecs/record.pyx @@ -0,0 +1,79 @@ +# +# This source file is part of the EdgeDB open source project. +# +# Copyright 2016-present MagicStack Inc. and the EdgeDB authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +@cython.final +cdef class RecordCodec(BaseNamedRecordCodec): + + cdef encode(self, WriteBuffer buf, object obj): + raise NotImplementedError + + cdef decode(self, FRBuffer *buf): + cdef: + object result + Py_ssize_t elem_count + Py_ssize_t i + int32_t elem_len + BaseCodec elem_codec + FRBuffer elem_buf + tuple fields_codecs = (self).fields_codecs + descriptor = (self).descriptor + + elem_count = hton.unpack_int32(frb_read(buf, 4)) + + if elem_count != len(fields_codecs): + raise RuntimeError( + f'cannot decode Record: expected {len(fields_codecs)} ' + f'elements, got {elem_count}') + + result = datatypes.record_new(descriptor) + + for i in range(elem_count): + frb_read(buf, 4) # reserved + elem_len = hton.unpack_int32(frb_read(buf, 4)) + + if elem_len == -1: + elem = None + else: + elem_codec = fields_codecs[i] + elem = elem_codec.decode( + frb_slice_from(&elem_buf, buf, elem_len)) + if frb_get_len(&elem_buf): + raise RuntimeError( + f'unexpected trailing data in buffer after ' + f'record element decoding: {frb_get_len(&elem_buf)}') + + datatypes.record_set(result, i, elem) + + return result + + @staticmethod + cdef BaseCodec new(bytes tid, tuple names, tuple codecs): + cdef: + RecordCodec codec + + codec = RecordCodec.__new__(RecordCodec) + + codec.tid = tid + codec.name = 'Record' + codec.descriptor = datatypes.record_desc_new( + names, NULL, NULL) + codec.fields_codecs = codecs + + return codec + diff --git a/setup.py b/setup.py index f4ea7b0c..447c3a45 100644 --- a/setup.py +++ b/setup.py @@ -325,6 +325,7 @@ def finalize_options(self): "gel/datatypes/record_desc.c", "gel/datatypes/namedtuple.c", "gel/datatypes/object.c", + "gel/datatypes/record.c", "gel/datatypes/hash.c", "gel/datatypes/repr.c", "gel/datatypes/comp.c", diff --git a/tests/test_async_query.py b/tests/test_async_query.py index 602346ad..bd71596b 100644 --- a/tests/test_async_query.py +++ b/tests/test_async_query.py @@ -1151,3 +1151,14 @@ async def test_transaction_state(self): await tx.execute(''' INSERT test::Tmp { id := $0, tmp := '' } ''', uuid.uuid4()) + + @unittest.expectedFailure + async def test_async_query_sql_01(self): + res = await self.client.query_sql("SELECT 1") + self.assertEqual(res[0].as_dict(), {'col~1': 1}) + + res = await self.client.query_sql("SELECT 1 as aa") + self.assertEqual(res[0].as_dict(), {'aa': 1}) + + res = await self.client.query_sql("SELECT FROM generate_series(0, 1)") + self.assertEqual(res[0].as_dict(), {})