From 5036cc11b14478077aead6e6c328148361b70caa Mon Sep 17 00:00:00 2001 From: Fantix King Date: Thu, 9 Jun 2022 19:06:05 -0400 Subject: [PATCH 01/11] Add with_globals() and friends through state over the protocol --- edgedb/abstract.py | 19 ++ edgedb/base_client.py | 5 + edgedb/datatypes/datatypes.h | 48 ++++ edgedb/datatypes/datatypes.pxd | 22 ++ edgedb/datatypes/datatypes.pyx | 32 +++ edgedb/datatypes/input_shape.c | 288 ++++++++++++++++++++++ edgedb/datatypes/internal.h | 3 + edgedb/datatypes/repr.c | 74 ++++++ edgedb/datatypes/sparse_object.c | 301 +++++++++++++++++++++++ edgedb/options.py | 106 +++++++- edgedb/protocol/codecs/array.pyx | 8 +- edgedb/protocol/codecs/codecs.pxd | 1 + edgedb/protocol/codecs/codecs.pyx | 28 +++ edgedb/protocol/codecs/sparse_object.pxd | 24 ++ edgedb/protocol/codecs/sparse_object.pyx | 126 ++++++++++ edgedb/protocol/protocol.pxd | 5 + edgedb/protocol/protocol.pyx | 36 +++ edgedb/transaction.py | 5 + setup.py | 2 + tests/test_async_query.py | 1 + tests/test_proto.py | 1 + 21 files changed, 1132 insertions(+), 3 deletions(-) create mode 100644 edgedb/datatypes/input_shape.c create mode 100644 edgedb/datatypes/sparse_object.c create mode 100644 edgedb/protocol/codecs/sparse_object.pxd create mode 100644 edgedb/protocol/codecs/sparse_object.pyx diff --git a/edgedb/abstract.py b/edgedb/abstract.py index fe89f1a9..44dc4498 100644 --- a/edgedb/abstract.py +++ b/edgedb/abstract.py @@ -39,11 +39,13 @@ class QueryContext(typing.NamedTuple): cache: QueryCache query_options: QueryOptions retry_options: typing.Optional[options.RetryOptions] + session: typing.Optional[options.Session] class ScriptContext(typing.NamedTuple): query: QueryWithArgs cache: QueryCache + session: typing.Optional[options.Session] _query_opts = QueryOptions( @@ -88,6 +90,9 @@ def _get_query_cache(self) -> QueryCache: def _get_retry_options(self) -> typing.Optional[options.RetryOptions]: return None + def _get_session(self) -> options.Session: + ... + class ReadOnlyExecutor(BaseReadOnlyExecutor): """Subclasses can execute *at least* read-only queries""" @@ -104,6 +109,7 @@ def query(self, query: str, *args, **kwargs) -> datatypes.Set: cache=self._get_query_cache(), query_options=_query_opts, retry_options=self._get_retry_options(), + session=self._get_session(), )) def query_single( @@ -114,6 +120,7 @@ def query_single( cache=self._get_query_cache(), query_options=_query_single_opts, retry_options=self._get_retry_options(), + session=self._get_session(), )) def query_required_single(self, query: str, *args, **kwargs) -> typing.Any: @@ -122,6 +129,7 @@ def query_required_single(self, query: str, *args, **kwargs) -> typing.Any: cache=self._get_query_cache(), query_options=_query_required_single_opts, retry_options=self._get_retry_options(), + session=self._get_session(), )) def query_json(self, query: str, *args, **kwargs) -> str: @@ -130,6 +138,7 @@ def query_json(self, query: str, *args, **kwargs) -> str: cache=self._get_query_cache(), query_options=_query_json_opts, retry_options=self._get_retry_options(), + session=self._get_session(), )) def query_single_json(self, query: str, *args, **kwargs) -> str: @@ -138,6 +147,7 @@ def query_single_json(self, query: str, *args, **kwargs) -> str: cache=self._get_query_cache(), query_options=_query_single_json_opts, retry_options=self._get_retry_options(), + session=self._get_session(), )) def query_required_single_json(self, query: str, *args, **kwargs) -> str: @@ -146,6 +156,7 @@ def query_required_single_json(self, query: str, *args, **kwargs) -> str: cache=self._get_query_cache(), query_options=_query_required_single_json_opts, retry_options=self._get_retry_options(), + session=self._get_session(), )) @abc.abstractmethod @@ -156,6 +167,7 @@ def execute(self, query: str, *args, **kwargs) -> None: self._execute(ScriptContext( query=QueryWithArgs(query, args, kwargs), cache=self._get_query_cache(), + session=self._get_session(), )) @@ -180,6 +192,7 @@ async def query(self, query: str, *args, **kwargs) -> datatypes.Set: cache=self._get_query_cache(), query_options=_query_opts, retry_options=self._get_retry_options(), + session=self._get_session(), )) async def query_single(self, query: str, *args, **kwargs) -> typing.Any: @@ -188,6 +201,7 @@ async def query_single(self, query: str, *args, **kwargs) -> typing.Any: cache=self._get_query_cache(), query_options=_query_single_opts, retry_options=self._get_retry_options(), + session=self._get_session(), )) async def query_required_single( @@ -201,6 +215,7 @@ async def query_required_single( cache=self._get_query_cache(), query_options=_query_required_single_opts, retry_options=self._get_retry_options(), + session=self._get_session(), )) async def query_json(self, query: str, *args, **kwargs) -> str: @@ -209,6 +224,7 @@ async def query_json(self, query: str, *args, **kwargs) -> str: cache=self._get_query_cache(), query_options=_query_json_opts, retry_options=self._get_retry_options(), + session=self._get_session(), )) async def query_single_json(self, query: str, *args, **kwargs) -> str: @@ -217,6 +233,7 @@ async def query_single_json(self, query: str, *args, **kwargs) -> str: cache=self._get_query_cache(), query_options=_query_single_json_opts, retry_options=self._get_retry_options(), + session=self._get_session(), )) async def query_required_single_json( @@ -230,6 +247,7 @@ async def query_required_single_json( cache=self._get_query_cache(), query_options=_query_required_single_json_opts, retry_options=self._get_retry_options(), + session=self._get_session(), )) @abc.abstractmethod @@ -240,6 +258,7 @@ async def execute(self, query: str, *args, **kwargs) -> None: await self._execute(ScriptContext( query=QueryWithArgs(query, args, kwargs), cache=self._get_query_cache(), + session=self._get_session(), )) diff --git a/edgedb/base_client.py b/edgedb/base_client.py index 4562f623..cfbd23f5 100644 --- a/edgedb/base_client.py +++ b/edgedb/base_client.py @@ -213,6 +213,7 @@ async def raw_query(self, query_context: abstract.QueryContext): execute = self._protocol.legacy_execute_anonymous else: execute = self._protocol.query + self._protocol.set_state(query_context.session) return await execute( query=query_context.query.query, args=query_context.query.args, @@ -264,6 +265,7 @@ async def _execute(self, script: abstract.ScriptContext) -> None: script.query.query, enums.Capability.EXECUTE ) else: + self._protocol.set_state(script.session) await self._protocol.execute( query=script.query.query, args=script.query.args, @@ -697,6 +699,9 @@ def _get_query_cache(self) -> abstract.QueryCache: def _get_retry_options(self) -> typing.Optional[_options.RetryOptions]: return self._options.retry_options + def _get_session(self) -> _options.Session: + return self._options.session + @property def max_concurrency(self) -> int: """Max number of connections in the pool.""" diff --git a/edgedb/datatypes/datatypes.h b/edgedb/datatypes/datatypes.h index 65fc086a..5af9289a 100644 --- a/edgedb/datatypes/datatypes.h +++ b/edgedb/datatypes/datatypes.h @@ -85,6 +85,29 @@ PyObject * EdgeRecordDesc_List(PyObject *, uint8_t, uint8_t); +/* === edgedb.InputShape ==================================== */ + +extern PyTypeObject EdgeInputShape_Type; + +#define EdgeInputShape_Check(d) (Py_TYPE(d) == &EdgeInputShape_Type) + +typedef struct { + PyObject_HEAD + PyObject *index; + PyObject *names; + Py_ssize_t size; +} EdgeInputShapeObject; + +PyObject * EdgeInputShape_InitType(void); +PyObject * EdgeInputShape_New(PyObject *); +PyObject * EdgeInputShape_PointerName(PyObject *, Py_ssize_t); + +Py_ssize_t EdgeInputShape_GetSize(PyObject *); +edge_attr_lookup_t EdgeInputShape_Lookup(PyObject *, PyObject *, Py_ssize_t *); +PyObject * EdgeInputShape_List(PyObject *); + + + /* === edgedb.Tuple ========================================= */ #define EDGE_TUPLE_FREELIST_SIZE 500 @@ -154,6 +177,31 @@ PyObject * EdgeObject_GetItem(PyObject *, Py_ssize_t); PyObject * EdgeObject_GetID(PyObject *ob); +/* === edgedb.SparseObject ======================================== */ + +#define EDGE_SPARSE_OBJECT_FREELIST_SIZE 2000 +#define EDGE_SPARSE_OBJECT_FREELIST_MAXSAVE 20 + +extern PyTypeObject EdgeSparseObject_Type; + +#define EdgeSparseObject_Check(d) (Py_TYPE(d) == &EdgeSparseObject_Type) + +typedef struct { + PyObject_VAR_HEAD + PyObject *weakreflist; + PyObject *desc; + Py_hash_t cached_hash; + PyObject *ob_item[1]; +} EdgeSparseObject; + +PyObject * EdgeSparseObject_InitType(void); +PyObject * EdgeSparseObject_New(PyObject *); +PyObject * EdgeSparseObject_GetInputShape(PyObject *); + +int EdgeSparseObject_SetItem(PyObject *, Py_ssize_t, PyObject *); +PyObject * EdgeSparseObject_GetItem(PyObject *, Py_ssize_t); + + /* === edgedb.Set =========================================== */ extern PyTypeObject EdgeSet_Type; diff --git a/edgedb/datatypes/datatypes.pxd b/edgedb/datatypes/datatypes.pxd index 059d8933..11e1c0a0 100644 --- a/edgedb/datatypes/datatypes.pxd +++ b/edgedb/datatypes/datatypes.pxd @@ -40,12 +40,24 @@ cdef extern from "datatypes.h": MANY AT_LEAST_ONE + ctypedef enum EdgeAttrLookup "edge_attr_lookup_t": + L_ERROR + L_NOT_FOUND + L_LINKPROP + L_PROPERTY + L_LINK + object EdgeRecordDesc_InitType() object EdgeRecordDesc_New(object, object, object) object EdgeRecordDesc_PointerName(object, Py_ssize_t pos) EdgeFieldCardinality EdgeRecordDesc_PointerCardinality( object, Py_ssize_t pos) + object EdgeInputShape_InitType() + object EdgeInputShape_New(object) + object EdgeInputShape_PointerName(object, Py_ssize_t pos) + EdgeAttrLookup EdgeInputShape_Lookup(object, object, Py_ssize_t* pos) + object EdgeTuple_InitType() object EdgeTuple_New(Py_ssize_t) int EdgeTuple_SetItem(object, Py_ssize_t, object) except -1 @@ -59,6 +71,11 @@ cdef extern from "datatypes.h": int EdgeObject_SetItem(object, Py_ssize_t, object) except -1 object EdgeObject_GetRecordDesc(object) + object EdgeSparseObject_InitType() + object EdgeSparseObject_New(object); + int EdgeSparseObject_SetItem(object, Py_ssize_t, object) except -1 + object EdgeSparseObject_GetInputShape(object) + bint EdgeSet_Check(object) object EdgeSet_InitType() object EdgeSet_New(Py_ssize_t); @@ -77,12 +94,17 @@ cdef extern from "datatypes.h": cdef record_desc_new(object names, object flags, object cards) cdef record_desc_pointer_name(object desc, Py_ssize_t pos) cdef record_desc_pointer_card(object desc, Py_ssize_t pos) +cdef input_shape_new(object names) +cdef input_shape_pointer_name(object desc, Py_ssize_t pos) +cdef Py_ssize_t input_shape_get_pos(object desc, object key) except -1 cdef tuple_new(Py_ssize_t size) cdef tuple_set(object tuple, Py_ssize_t pos, object elem) cdef namedtuple_new(object desc) cdef namedtuple_set(object tuple, Py_ssize_t pos, object elem) cdef object_new(object desc) cdef object_set(object tuple, Py_ssize_t pos, object elem) +cdef sparse_object_new(object desc) +cdef sparse_object_set(object tuple, Py_ssize_t pos, object elem) cdef bint set_check(object set) cdef set_new(Py_ssize_t size) cdef set_set(object set, Py_ssize_t pos, object elem) diff --git a/edgedb/datatypes/datatypes.pyx b/edgedb/datatypes/datatypes.pyx index 848d2c5e..7b5a4c03 100644 --- a/edgedb/datatypes/datatypes.pyx +++ b/edgedb/datatypes/datatypes.pyx @@ -19,6 +19,7 @@ cimport cython cimport cpython +from libc cimport stdlib include "./relative_duration.pyx" include "./enum.pyx" @@ -96,6 +97,29 @@ cdef record_desc_pointer_card(object desc, Py_ssize_t pos): return EdgeRecordDesc_PointerCardinality(desc, pos) +cdef input_shape_new(object names): + return EdgeInputShape_New(names) + + +cdef input_shape_pointer_name(object desc, Py_ssize_t pos): + return EdgeInputShape_PointerName(desc, pos) + + +cdef Py_ssize_t input_shape_get_pos(object desc, object key) except -1: + cdef: + Py_ssize_t pos + EdgeAttrLookup res = EdgeInputShape_Lookup(desc, key, &pos) + + if res == L_ERROR: + return -1 + elif res == L_NOT_FOUND: + raise LookupError(key) + elif res == L_PROPERTY: + return pos + else: + stdlib.abort() + + cdef tuple_new(Py_ssize_t size): return EdgeTuple_New(size) @@ -120,6 +144,14 @@ cdef object_set(object obj, Py_ssize_t pos, object elem): EdgeObject_SetItem(obj, pos, elem) +cdef sparse_object_new(object desc): + return EdgeSparseObject_New(desc) + + +cdef sparse_object_set(object obj, Py_ssize_t pos, object elem): + EdgeSparseObject_SetItem(obj, pos, elem) + + cdef bint set_check(object set): return EdgeSet_Check(set) diff --git a/edgedb/datatypes/input_shape.c b/edgedb/datatypes/input_shape.c new file mode 100644 index 00000000..20c23b87 --- /dev/null +++ b/edgedb/datatypes/input_shape.c @@ -0,0 +1,288 @@ +/* +* This source file is part of the EdgeDB open source project. +* +* Copyright 2022-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 "internal.h" + + +static int init_type_called = 0; + + +static void +input_shape_dealloc(EdgeInputShapeObject *o) +{ + PyObject_GC_UnTrack(o); + Py_CLEAR(o->index); + Py_CLEAR(o->names); + PyObject_GC_Del(o); +} + + +static int +input_shape_traverse(EdgeInputShapeObject *o, visitproc visit, void *arg) +{ + Py_VISIT(o->index); + Py_VISIT(o->names); + return 0; +} + +static PyObject * +input_shape_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds) +{ + if (args == NULL || + PyTuple_Size(args) != 1 || + (kwds != NULL && PyDict_Size(kwds))) + { + PyErr_SetString( + PyExc_TypeError, + "InputShape accepts exactly one positional argument"); + return NULL; + } + + return EdgeInputShape_New(PyTuple_GET_ITEM(args, 0)); +} + + +static PyObject * +input_shape_get_pos(EdgeInputShapeObject *o, PyObject *arg) { + Py_ssize_t pos; + edge_attr_lookup_t ret = EdgeInputShape_Lookup((PyObject *)o, arg, &pos); + switch (ret) { + case L_ERROR: + return NULL; + + case L_NOT_FOUND: + PyErr_SetObject(PyExc_LookupError, arg); + return NULL; + + case L_PROPERTY: + return PyLong_FromLong((long)pos); + + default: + abort(); + } +} + + +static PyObject * +input_shape_dir(EdgeInputShapeObject *o, PyObject *args) +{ + PyObject *names = o->names; + Py_INCREF(names); + return names; +} + + +static PyMethodDef input_shape_methods[] = { + {"get_pos", (PyCFunction)input_shape_get_pos, METH_O, NULL}, + {"__dir__", (PyCFunction)input_shape_dir, METH_NOARGS, NULL}, + {NULL, NULL} +}; + + +PyTypeObject EdgeInputShape_Type = { + PyVarObject_HEAD_INIT(NULL, 0) + .tp_name = "edgedb.InputShape", + .tp_basicsize = sizeof(EdgeInputShapeObject), + .tp_dealloc = (destructor)input_shape_dealloc, + .tp_getattro = PyObject_GenericGetAttr, + .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, + .tp_traverse = (traverseproc)input_shape_traverse, + .tp_new = input_shape_tp_new, + .tp_methods = input_shape_methods, +}; + + +PyObject * +EdgeInputShape_New(PyObject *names) +{ + EdgeInputShapeObject *o; + + assert(init_type_called); + + if (!names || !PyTuple_CheckExact(names)) { + PyErr_SetString( + PyExc_TypeError, + "InputShape requires a tuple as its first argument"); + return NULL; + } + + if (Py_SIZE(names) > EDGE_MAX_TUPLE_SIZE) { + PyErr_Format( + PyExc_ValueError, + "EdgeDB does not supports tuples with more than %d elements", + EDGE_MAX_TUPLE_SIZE); + return NULL; + } + + Py_ssize_t size = Py_SIZE(names); + + PyObject *index = PyDict_New(); + if (index == NULL) { + return NULL; + } + + for (Py_ssize_t i = 0; i < size; i++) { + PyObject *key = PyTuple_GET_ITEM(names, i); /* borrowed */ + if (!PyUnicode_CheckExact(key)) { + PyErr_SetString( + PyExc_ValueError, + "InputShape received a non-str key"); + return NULL; + } + + PyObject *num = PyLong_FromLong(i); + if (num == NULL) { + Py_DECREF(index); + return NULL; + } + + if (PyDict_SetItem(index, key, num)) { + Py_DECREF(index); + Py_DECREF(num); + return NULL; + } + + Py_DECREF(num); + } + + o = PyObject_GC_New(EdgeInputShapeObject, &EdgeInputShape_Type); + if (o == NULL) { + Py_DECREF(index); + return NULL; + } + + o->index = index; + + Py_INCREF(names); + o->names = names; + + o->size = size; + + PyObject_GC_Track(o); + return (PyObject *)o; +} + + +edge_attr_lookup_t +EdgeInputShape_Lookup(PyObject *ob, PyObject *key, Py_ssize_t *pos) +{ + if (!EdgeInputShape_Check(ob)) { + PyErr_BadInternalCall(); + return L_ERROR; + } + + EdgeInputShapeObject *d = (EdgeInputShapeObject *)ob; + + PyObject *res = PyDict_GetItem(d->index, key); /* borrowed */ + if (res == NULL) { + if (PyErr_Occurred()) { + return L_ERROR; + } + else { + return L_NOT_FOUND; + } + } + + assert(PyLong_CheckExact(res)); + long res_long = PyLong_AsLong(res); + if (res_long < 0) { + assert(PyErr_Occurred()); + return L_ERROR; + } + assert(res_long < d->size); + *pos = res_long; + + return L_PROPERTY; +} + + +PyObject * +EdgeInputShape_PointerName(PyObject *ob, Py_ssize_t pos) +{ + if (!EdgeInputShape_Check(ob)) { + PyErr_BadInternalCall(); + return NULL; + } + EdgeInputShapeObject *o = (EdgeInputShapeObject *)ob; + PyObject * key = PyTuple_GetItem(o->names, pos); + if (key == NULL) { + return NULL; + } + Py_INCREF(key); + return key; +} + + +Py_ssize_t +EdgeInputShape_GetSize(PyObject *ob) +{ + assert(ob != NULL); + if (!EdgeInputShape_Check(ob)) { + PyErr_BadInternalCall(); + return -1; + } + EdgeInputShapeObject *o = (EdgeInputShapeObject *)ob; + return o->size; +} + + +PyObject * +EdgeInputShape_List(PyObject *ob) +{ + if (!EdgeInputShape_Check(ob)) { + PyErr_BadInternalCall(); + return NULL; + } + + EdgeInputShapeObject *o = (EdgeInputShapeObject *)ob; + + PyObject *ret = PyList_New(o->size); + if (ret == NULL) { + return NULL; + } + + for (Py_ssize_t i = 0; i < o->size; i++) { + PyObject *name = PyTuple_GetItem(o->names, i); + if (name == NULL) { + Py_DECREF(ret); + return NULL; + } + + Py_INCREF(name); + if (PyList_SetItem(ret, i, name)) { + Py_DECREF(ret); + return NULL; + } + } + + return ret; +} + + +PyObject * +EdgeInputShape_InitType(void) +{ + if (PyType_Ready(&EdgeInputShape_Type) < 0) { + return NULL; + } + + init_type_called = 1; + return (PyObject *)&EdgeInputShape_Type; +} diff --git a/edgedb/datatypes/internal.h b/edgedb/datatypes/internal.h index 67963b41..cddc2578 100644 --- a/edgedb/datatypes/internal.h +++ b/edgedb/datatypes/internal.h @@ -50,6 +50,9 @@ int _EdgeGeneric_RenderValues( int _EdgeGeneric_RenderItems(_PyUnicodeWriter *, PyObject *, PyObject *, PyObject **, Py_ssize_t, int, int); +int _EdgeGeneric_RenderSparseItems(_PyUnicodeWriter *writer, + PyObject *host, PyObject *desc, + PyObject **items, Py_ssize_t len); PyObject * _EdgeGeneric_RichCompareValues(PyObject **, Py_ssize_t, PyObject **, Py_ssize_t, diff --git a/edgedb/datatypes/repr.c b/edgedb/datatypes/repr.c index bfe0da4b..4af524f5 100644 --- a/edgedb/datatypes/repr.c +++ b/edgedb/datatypes/repr.c @@ -169,3 +169,77 @@ _EdgeGeneric_RenderItems(_PyUnicodeWriter *writer, Py_ReprLeave((PyObject *)host); return -1; } + + +int +_EdgeGeneric_RenderSparseItems(_PyUnicodeWriter *writer, + PyObject *host, PyObject *desc, + PyObject **items, Py_ssize_t len) +{ + assert(EdgeInputShape_GetSize(desc) == len); + + PyObject *item_repr = NULL; + PyObject *item_name = NULL; + int first = 1; + + int res = Py_ReprEnter(host); + if (res != 0) { + if (res > 0) { + if (_PyUnicodeWriter_WriteASCIIString(writer, "...", 3) < 0) { + return -1; + } + return 0; + } + else { + return -1; + } + } + + for (Py_ssize_t i = 0; i < len; i++) { + if (items[i] == Py_None) { + continue; + } + + item_repr = _EdgeGeneric_RenderObject(items[i]); + if (item_repr == NULL) { + goto error; + } + + item_name = EdgeInputShape_PointerName(desc, i); + if (item_name == NULL) { + goto error; + } + assert(PyUnicode_CheckExact(item_name)); + + if (first) { + first = 0; + } else { + if (_PyUnicodeWriter_WriteASCIIString(writer, ", ", 2) < 0) { + goto error; + } + } + + if (_PyUnicodeWriter_WriteStr(writer, item_name) < 0) { + goto error; + } + Py_CLEAR(item_name); + + if (_PyUnicodeWriter_WriteASCIIString(writer, " := ", 4) < 0) { + goto error; + } + + if (_PyUnicodeWriter_WriteStr(writer, item_repr) < 0) { + goto error; + } + Py_CLEAR(item_repr); + } + + Py_ReprLeave((PyObject *)host); + return 0; + +error: + Py_CLEAR(item_repr); + Py_CLEAR(item_name); + Py_ReprLeave((PyObject *)host); + return -1; +} diff --git a/edgedb/datatypes/sparse_object.c b/edgedb/datatypes/sparse_object.c new file mode 100644 index 00000000..8ded5d7c --- /dev/null +++ b/edgedb/datatypes/sparse_object.c @@ -0,0 +1,301 @@ +/* +* This source file is part of the EdgeDB open source project. +* +* Copyright 2022-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; +static Py_hash_t base_hash = -1; + + +EDGE_SETUP_FREELIST( + EDGE_SPARSE_OBJECT, + EdgeSparseObject, + EDGE_SPARSE_OBJECT_FREELIST_MAXSAVE, + EDGE_SPARSE_OBJECT_FREELIST_SIZE) + + +#define EdgeSparseObject_GET_ITEM(op, i) \ + (((EdgeSparseObject *)(op))->ob_item[i]) +#define EdgeSparseObject_SET_ITEM(op, i, v) \ + (((EdgeSparseObject *)(op))->ob_item[i] = v) + + +PyObject * +EdgeSparseObject_New(PyObject *desc) +{ + assert(init_type_called); + + if (desc == NULL || !EdgeInputShape_Check(desc)) { + PyErr_BadInternalCall(); + return NULL; + } + + Py_ssize_t size = EdgeInputShape_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; + } + + EdgeSparseObject *o = NULL; + EDGE_NEW_WITH_FREELIST(EDGE_SPARSE_OBJECT, EdgeSparseObject, + &EdgeSparseObject_Type, o, size); + assert(o != NULL); + assert(Py_SIZE(o) == size); + assert(EdgeSparseObject_Check(o)); + + o->weakreflist = NULL; + + Py_INCREF(desc); + o->desc = desc; + + o->cached_hash = -1; + + PyObject_GC_Track(o); + return (PyObject *)o; +} + + +PyObject * +EdgeSparseObject_GetInputShape(PyObject *o) +{ + if (!EdgeSparseObject_Check(o)) { + PyErr_Format( + PyExc_TypeError, + "an instance of edgedb.Object expected"); + return NULL; + } + + PyObject *desc = ((EdgeSparseObject *)o)->desc; + Py_INCREF(desc); + return desc; +} + + +int +EdgeSparseObject_SetItem(PyObject *ob, Py_ssize_t i, PyObject *el) +{ + assert(EdgeSparseObject_Check(ob)); + EdgeSparseObject *o = (EdgeSparseObject *)ob; + assert(i >= 0); + assert(i < Py_SIZE(o)); + Py_INCREF(el); + EdgeSparseObject_SET_ITEM(o, i, el); + return 0; +} + + +PyObject * +EdgeSparseObject_GetItem(PyObject *ob, Py_ssize_t i) +{ + assert(EdgeSparseObject_Check(ob)); + EdgeSparseObject *o = (EdgeSparseObject *)ob; + if (i < 0 || i >= Py_SIZE(o)) { + PyErr_BadInternalCall(); + return NULL; + } + PyObject *el = EdgeSparseObject_GET_ITEM(o, i); + Py_INCREF(el); + return el; +} + + +static void +sparse_object_dealloc(EdgeSparseObject *o) +{ + PyObject_GC_UnTrack(o); + if (o->weakreflist != NULL) { + PyObject_ClearWeakRefs((PyObject*)o); + } + Py_CLEAR(o->desc); + o->cached_hash = -1; + Py_TRASHCAN_SAFE_BEGIN(o) + EDGE_DEALLOC_WITH_FREELIST(EDGE_SPARSE_OBJECT, EdgeSparseObject, o); + Py_TRASHCAN_SAFE_END(o) +} + + +static int +sparse_object_traverse(EdgeSparseObject *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 Py_hash_t +sparse_object_hash(EdgeSparseObject *o) +{ + if (o->cached_hash == -1) { + o->cached_hash = _EdgeGeneric_HashWithBase( + base_hash, o->ob_item, Py_SIZE(o)); + } + return o->cached_hash; +} + + +static PyObject * +sparse_object_getattr(EdgeSparseObject *o, PyObject *name) +{ + Py_ssize_t pos; + edge_attr_lookup_t ret = EdgeInputShape_Lookup( + (PyObject *)o->desc, name, &pos); + switch (ret) { + case L_ERROR: + return NULL; + + case L_NOT_FOUND: + return PyObject_GenericGetAttr((PyObject *)o, name); + + case L_PROPERTY: { + PyObject *val = EdgeSparseObject_GET_ITEM(o, pos); + Py_INCREF(val); + return val; + } + + default: + abort(); + } +} + +static PyObject * +sparse_object_getitem(EdgeSparseObject *o, PyObject *name) +{ + Py_ssize_t pos; + edge_attr_lookup_t ret = EdgeInputShape_Lookup( + (PyObject *)o->desc, name, &pos); + switch (ret) { + case L_ERROR: + return NULL; + + case L_PROPERTY: + PyErr_Format( + PyExc_TypeError, + "property %R should be accessed via dot notation", + name); + return NULL; + + case L_NOT_FOUND: + PyErr_Format( + PyExc_KeyError, + "link %R does not exist", + name); + return NULL; + + default: + abort(); + } + +} + + +static PyObject * +sparse_object_repr(EdgeSparseObject *o) +{ + _PyUnicodeWriter writer; + _PyUnicodeWriter_Init(&writer); + writer.overallocate = 1; + + if (_PyUnicodeWriter_WriteASCIIString(&writer, "SparseObject{", 13) < 0) { + goto error; + } + + if (_EdgeGeneric_RenderSparseItems(&writer, + (PyObject *)o, o->desc, + o->ob_item, Py_SIZE(o)) < 0) + { + goto error; + } + + if (_PyUnicodeWriter_WriteChar(&writer, '}') < 0) { + goto error; + } + + return _PyUnicodeWriter_Finish(&writer); + +error: + _PyUnicodeWriter_Dealloc(&writer); + return NULL; +} + + +static PyObject * +sparse_object_dir(EdgeSparseObject *o, PyObject *args) +{ + return EdgeInputShape_List(o->desc); +} + + +static PyMethodDef sparse_object_methods[] = { + {"__dir__", (PyCFunction)sparse_object_dir, METH_NOARGS, NULL}, + {NULL, NULL} +}; + + +static PyMappingMethods sparse_object_as_mapping = { + .mp_subscript = (binaryfunc)sparse_object_getitem, +}; + + +PyTypeObject EdgeSparseObject_Type = { + PyVarObject_HEAD_INIT(NULL, 0) + "edgedb.SparseObject", + .tp_basicsize = sizeof(EdgeSparseObject) - sizeof(PyObject *), + .tp_itemsize = sizeof(PyObject *), + .tp_dealloc = (destructor)sparse_object_dealloc, + .tp_hash = (hashfunc)sparse_object_hash, + .tp_methods = sparse_object_methods, + .tp_as_mapping = &sparse_object_as_mapping, + .tp_getattro = (getattrofunc)sparse_object_getattr, + .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, + .tp_traverse = (traverseproc)sparse_object_traverse, + .tp_free = PyObject_GC_Del, + .tp_repr = (reprfunc)sparse_object_repr, + .tp_weaklistoffset = offsetof(EdgeSparseObject, weakreflist), +}; + + +PyObject * +EdgeSparseObject_InitType(void) +{ + if (PyType_Ready(&EdgeSparseObject_Type) < 0) { + return NULL; + } + + base_hash = _EdgeGeneric_HashString("edgedb.SparseObject"); + if (base_hash == -1) { + return NULL; + } + + init_type_called = 1; + return (PyObject *)&EdgeSparseObject_Type; +} diff --git a/edgedb/options.py b/edgedb/options.py index f1e16924..c680e005 100644 --- a/edgedb/options.py +++ b/edgedb/options.py @@ -1,6 +1,7 @@ import abc import enum import random +import typing from collections import namedtuple from . import errors @@ -108,6 +109,67 @@ def __repr__(self): ) +class Session: + __slots__ = ['_module', '_aliases', '_configs', '_globals'] + + def __init__( + self, + module: str = 'default', + aliases: typing.Mapping[str, str] = None, + configs: typing.Mapping[str, typing.Any] = None, + globals_: typing.Mapping[str, typing.Any] = None, + ): + self._module = module + self._aliases = {} if aliases is None else dict(aliases) + self._configs = {} if configs is None else dict(configs) + self._globals = {} if globals_ is None else dict(globals_) + + @classmethod + def defaults(cls): + return cls() + + def with_aliases(self, module=None, **aliases): + new_aliases = self._aliases.copy() + new_aliases.update(aliases) + return Session( + module=self._module if module is None else module, + aliases=new_aliases, + configs=self._configs, + globals_=self._globals, + ) + + def with_configs(self, **configs): + new_configs = self._configs.copy() + new_configs.update(configs) + return Session( + module=self._module, + aliases=self._aliases, + configs=new_configs, + globals_=self._globals, + ) + + def with_globals(self, **globals_): + new_globals = self._globals.copy() + new_globals.update(globals_) + return Session( + module=self._module, + aliases=self._aliases, + configs=self._configs, + globals_=new_globals, + ) + + def as_dict(self): + return { + "module": self._module, + "aliases": list(self._aliases.items()), + "config": self._configs, + "globals": { + (k if '::' in k else f'{self._module}::{k}'): v + for k, v in self._globals.items() + }, + } + + class _OptionsMixin: def __init__(self, *args, **kwargs): self._options = _Options.defaults() @@ -153,19 +215,47 @@ def with_retry_options(self, options: RetryOptions=None): result._options = self._options.with_retry_options(options) return result + def with_session(self, session: Session): + result = self._shallow_clone() + result._options = self._options.with_session(session) + return result + + def with_aliases(self, module=None, **aliases): + result = self._shallow_clone() + result._options = self._options.with_session( + self._options.session.with_aliases(module=module, **aliases) + ) + return result + + def with_configs(self, **configs): + result = self._shallow_clone() + result._options = self._options.with_session( + self._options.session.with_configs(**configs) + ) + return result + + def with_globals(self, **globals_): + result = self._shallow_clone() + result._options = self._options.with_session( + self._options.session.with_globals(**globals_) + ) + return result + class _Options: """Internal class for storing connection options""" - __slots__ = ['_retry_options', '_transaction_options'] + __slots__ = ['_retry_options', '_transaction_options', '_session'] def __init__( self, retry_options: RetryOptions, transaction_options: TransactionOptions, + session: Session, ): self._retry_options = retry_options self._transaction_options = transaction_options + self._session = session @property def retry_options(self): @@ -175,16 +265,29 @@ def retry_options(self): def transaction_options(self): return self._transaction_options + @property + def session(self): + return self._session + def with_retry_options(self, options: RetryOptions): return _Options( options, self._transaction_options, + self._session, ) def with_transaction_options(self, options: TransactionOptions): return _Options( self._retry_options, options, + self._session, + ) + + def with_session(self, session: Session): + return _Options( + self._retry_options, + self._transaction_options, + session, ) @classmethod @@ -192,4 +295,5 @@ def defaults(cls): return cls( RetryOptions.defaults(), TransactionOptions.defaults(), + Session.defaults(), ) diff --git a/edgedb/protocol/codecs/array.pyx b/edgedb/protocol/codecs/array.pyx index e2b9aa88..e8ef3e23 100644 --- a/edgedb/protocol/codecs/array.pyx +++ b/edgedb/protocol/codecs/array.pyx @@ -44,8 +44,12 @@ cdef class BaseArrayCodec(BaseCodec): Py_ssize_t objlen Py_ssize_t i - if not isinstance(self.sub_codec, ScalarCodec): - raise TypeError('only arrays of scalars are supported') + if not isinstance(self.sub_codec, (ScalarCodec, TupleCodec)): + raise TypeError( + 'only arrays of scalars are supported (got type {!r})'.format( + type(self.sub_codec).__name__ + ) + ) if not _is_array_iterable(obj): raise TypeError( diff --git a/edgedb/protocol/codecs/codecs.pxd b/edgedb/protocol/codecs/codecs.pxd index d64ff1ff..ac9a8963 100644 --- a/edgedb/protocol/codecs/codecs.pxd +++ b/edgedb/protocol/codecs/codecs.pxd @@ -22,6 +22,7 @@ include "./scalar.pxd" include "./tuple.pxd" include "./namedtuple.pxd" include "./object.pxd" +include "./sparse_object.pxd" include "./array.pxd" include "./set.pxd" include "./enum.pxd" diff --git a/edgedb/protocol/codecs/codecs.pyx b/edgedb/protocol/codecs/codecs.pyx index 95c80a67..245ccfe3 100644 --- a/edgedb/protocol/codecs/codecs.pyx +++ b/edgedb/protocol/codecs/codecs.pyx @@ -31,6 +31,7 @@ include "./scalar.pyx" include "./tuple.pyx" include "./namedtuple.pyx" include "./object.pyx" +include "./sparse_object.pyx" include "./array.pyx" include "./set.pyx" include "./enum.pyx" @@ -44,6 +45,7 @@ DEF CTYPE_TUPLE = 4 DEF CTYPE_NAMEDTUPLE = 5 DEF CTYPE_ARRAY = 6 DEF CTYPE_ENUM = 7 +DEF CTYPE_INPUT_SHAPE = 8 DEF _CODECS_BUILD_CACHE_SIZE = 200 @@ -147,6 +149,13 @@ cdef class CodecsRegistry: str_len = hton.unpack_uint32(frb_read(spec, 4)) frb_read(spec, str_len) + elif t == CTYPE_INPUT_SHAPE: + els = hton.unpack_int16(frb_read(spec, 2)) + for i in range(els): + str_len = hton.unpack_uint32(frb_read(spec, 4)) + # read the (`str_len` bytes) and (2 bytes) + frb_read(spec, str_len + 2) + elif (t >= 0x7f and t <= 0xff): # Ignore all type annotations. str_len = hton.unpack_uint32(frb_read(spec, 4)) @@ -264,6 +273,25 @@ cdef class CodecsRegistry: sub_codec = codecs_list[pos] res = ArrayCodec.new(tid, sub_codec, dim_len) + elif t == CTYPE_INPUT_SHAPE: + 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 = SparseObjectCodec.new(tid, names, codecs) + else: raise NotImplementedError( f'no codec implementation for EdgeDB data class {t}') diff --git a/edgedb/protocol/codecs/sparse_object.pxd b/edgedb/protocol/codecs/sparse_object.pxd new file mode 100644 index 00000000..54068fc5 --- /dev/null +++ b/edgedb/protocol/codecs/sparse_object.pxd @@ -0,0 +1,24 @@ +# +# This source file is part of the EdgeDB open source project. +# +# Copyright 2022-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 SparseObjectCodec(BaseNamedRecordCodec): + + @staticmethod + cdef BaseCodec new(bytes tid, tuple names, tuple codecs) diff --git a/edgedb/protocol/codecs/sparse_object.pyx b/edgedb/protocol/codecs/sparse_object.pyx new file mode 100644 index 00000000..635a1a3d --- /dev/null +++ b/edgedb/protocol/codecs/sparse_object.pyx @@ -0,0 +1,126 @@ +# +# This source file is part of the EdgeDB open source project. +# +# Copyright 2022-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 SparseObjectCodec(BaseNamedRecordCodec): + + cdef encode(self, WriteBuffer buf, object obj): + cdef: + WriteBuffer elem_data + Py_ssize_t objlen = 0 + Py_ssize_t i + BaseCodec sub_codec + descriptor = (self).descriptor + + elem_data = WriteBuffer.new() + for name, arg in obj.items(): + if arg is not None: + try: + i = datatypes.input_shape_get_pos(descriptor, name) + except LookupError: + raise self._make_unknown_args_error_message(obj) from None + objlen += 1 + elem_data.write_int32(i) + + sub_codec = (self.fields_codecs[i]) + try: + sub_codec.encode(elem_data, arg) + except (TypeError, ValueError) as e: + value_repr = repr(arg) + if len(value_repr) > 40: + value_repr = value_repr[:40] + '...' + raise errors.InvalidArgumentError( + 'invalid input for session argument ' + f' {name} := {value_repr} ({e})') from e + + buf.write_int32(4 + elem_data.len()) # buffer length + buf.write_int32(objlen) + buf.write_buffer(elem_data) + + def _make_unknown_args_error_message(self, args): + cdef descriptor = (self).descriptor + + acceptable_args = set() + + for i in range(len(self.fields_codecs)): + name = datatypes.input_shape_pointer_name(descriptor, i) + acceptable_args.add(name) + + passed_args = set(args.keys()) + error_message = f'acceptable {acceptable_args} arguments' + + passed_args_repr = repr(passed_args) if passed_args else 'nothing' + error_message += f', got {passed_args_repr}' + + extra_args = set(passed_args) - set(acceptable_args) + if extra_args: + error_message += f', extra {extra_args}' + + return errors.QueryArgumentError(error_message) + + 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)) + + result = datatypes.sparse_object_new(descriptor) + + for i in range(len(fields_codecs)): + datatypes.sparse_object_set(result, i, None) + + for _ in range(elem_count): + i = hton.unpack_int32(frb_read(buf, 4)) + elem_len = hton.unpack_int32(frb_read(buf, 4)) + + if elem_len == -1: + continue + + 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'object element decoding: {frb_get_len(&elem_buf)}') + cpython.Py_DECREF(None) + datatypes.sparse_object_set(result, i, elem) + + return result + + @staticmethod + cdef BaseCodec new(bytes tid, tuple names, tuple codecs): + cdef: + SparseObjectCodec codec + + codec = SparseObjectCodec.__new__(SparseObjectCodec) + + codec.tid = tid + codec.name = 'SparseObject' + codec.descriptor = datatypes.input_shape_new(names) + codec.fields_codecs = codecs + + return codec diff --git a/edgedb/protocol/protocol.pxd b/edgedb/protocol/protocol.pxd index 69ed875e..282caef9 100644 --- a/edgedb/protocol/protocol.pxd +++ b/edgedb/protocol/protocol.pxd @@ -103,6 +103,11 @@ cdef class SansIOProtocol: readonly bint is_legacy + bytes state_type_id + BaseCodec state_codec + WriteBuffer state + object user_state + cdef encode_args(self, BaseCodec in_dc, WriteBuffer buf, args, kwargs) cdef parse_data_messages(self, BaseCodec out_dc, result) diff --git a/edgedb/protocol/protocol.pyx b/edgedb/protocol/protocol.pyx index 8adb833b..f9e0ec04 100644 --- a/edgedb/protocol/protocol.pyx +++ b/edgedb/protocol/protocol.pyx @@ -150,6 +150,11 @@ cdef class SansIOProtocol: self.reset_status() self.protocol_version = (PROTO_VER_MAJOR, 0) + self.state_type_id = b'\0' * 16 + self.state_codec = None + self.state = None + self.user_state = None + cdef reset_status(self): self.last_status = None self.last_details = None @@ -193,6 +198,17 @@ cdef class SansIOProtocol: self.buffer.read_len_prefixed_bytes() # value num_fields -= 1 + def set_state(self, user_state): + cdef WriteBuffer buf = WriteBuffer.new() + if self.user_state is user_state: + return + if user_state is None: + self.state = None + else: + self.state_codec.encode(buf, user_state.as_dict()) + self.state = buf + self.user_state = user_state + cdef ensure_connected(self): if self.cancelled: raise errors.ClientConnectionClosedError( @@ -356,6 +372,13 @@ cdef class SansIOProtocol: buf.write_buffer(params) + buf.write_bytes(self.state_type_id) + buf.write_int16(1) + if self.state is None: + buf.write_int32(0) + else: + buf.write_buffer(self.state) + buf.write_bytes(in_dc.get_tid()) buf.write_bytes(out_dc.get_tid()) @@ -982,6 +1005,16 @@ cdef class SansIOProtocol: data = buf.read_len_prefixed_bytes() self.server_settings[name] = self.parse_system_config(codec, data) + elif name == 'session_state_description': + self.state_type_id = typedesc_id = val[:16] + typedesc = val[16 + 4:] + + if self.internal_reg.has_codec(typedesc_id): + self.state_codec = self.internal_reg.get_codec(typedesc_id) + else: + self.state_codec = self.internal_reg.build_codec( + typedesc, self.protocol_version + ) else: self.server_settings[name] = val @@ -1149,6 +1182,9 @@ cdef class SansIOProtocol: self.ignore_headers() self.last_capabilities = enums.Capability(self.buffer.read_int64()) self.last_status = self.buffer.read_len_prefixed_bytes() + self.buffer.read_bytes(16) # state type id + assert self.buffer.read_int16() == 1 + self.buffer.read_len_prefixed_bytes() # state self.buffer.finish_message() cdef parse_sync_message(self): diff --git a/edgedb/transaction.py b/edgedb/transaction.py index aeffa889..a305440e 100644 --- a/edgedb/transaction.py +++ b/edgedb/transaction.py @@ -21,6 +21,7 @@ from . import abstract from . import errors +from . import options class TransactionState(enum.Enum): @@ -178,6 +179,9 @@ async def _exit(self, extype, ex): def _get_query_cache(self) -> abstract.QueryCache: return self._client._get_query_cache() + def _get_session(self) -> options.Session: + return self._client._get_session() + async def _query(self, query_context: abstract.QueryContext): await self._ensure_transaction() return await self._connection.raw_query(query_context) @@ -190,6 +194,7 @@ async def _privileged_execute(self, query: str) -> None: await self._connection.privileged_execute(abstract.ScriptContext( query=abstract.QueryWithArgs(query, (), {}), cache=self._get_query_cache(), + session=self._get_session(), )) diff --git a/setup.py b/setup.py index a345d7f4..cf95bf4a 100644 --- a/setup.py +++ b/setup.py @@ -302,9 +302,11 @@ def finalize_options(self): "edgedb.datatypes.datatypes", ["edgedb/datatypes/args.c", "edgedb/datatypes/record_desc.c", + "edgedb/datatypes/input_shape.c", "edgedb/datatypes/tuple.c", "edgedb/datatypes/namedtuple.c", "edgedb/datatypes/object.c", + "edgedb/datatypes/sparse_object.c", "edgedb/datatypes/set.c", "edgedb/datatypes/hash.c", "edgedb/datatypes/array.c", diff --git a/tests/test_async_query.py b/tests/test_async_query.py index c5680694..c865fca2 100644 --- a/tests/test_async_query.py +++ b/tests/test_async_query.py @@ -917,6 +917,7 @@ async def test_json_elements(self): required_one=False, ), retry_options=None, + session=None, ) ) self.assertEqual( diff --git a/tests/test_proto.py b/tests/test_proto.py index 702cbede..aa66b1a7 100644 --- a/tests/test_proto.py +++ b/tests/test_proto.py @@ -47,6 +47,7 @@ def test_json_elements(self): required_one=False, ), retry_options=None, + session=None, ) ) ) From 08ab98bb3a424c8bf3447ca91968da944b8ba86a Mon Sep 17 00:00:00 2001 From: Fantix King Date: Mon, 13 Jun 2022 11:57:35 -0400 Subject: [PATCH 02/11] CRF: add back cardinality and flags This actually drops the custom implementation of SparseObject because the Python binding doesn't need its decoding --- edgedb/datatypes/datatypes.h | 48 ---- edgedb/datatypes/datatypes.pxd | 22 -- edgedb/datatypes/datatypes.pyx | 32 --- edgedb/datatypes/input_shape.c | 288 ---------------------- edgedb/datatypes/internal.h | 3 - edgedb/datatypes/repr.c | 74 ------ edgedb/datatypes/sparse_object.c | 301 ----------------------- edgedb/protocol/codecs/codecs.pxd | 1 - edgedb/protocol/codecs/codecs.pyx | 35 +-- edgedb/protocol/codecs/object.pxd | 3 +- edgedb/protocol/codecs/object.pyx | 59 ++++- edgedb/protocol/codecs/sparse_object.pxd | 24 -- edgedb/protocol/codecs/sparse_object.pyx | 126 ---------- setup.py | 2 - 14 files changed, 59 insertions(+), 959 deletions(-) delete mode 100644 edgedb/datatypes/input_shape.c delete mode 100644 edgedb/datatypes/sparse_object.c delete mode 100644 edgedb/protocol/codecs/sparse_object.pxd delete mode 100644 edgedb/protocol/codecs/sparse_object.pyx diff --git a/edgedb/datatypes/datatypes.h b/edgedb/datatypes/datatypes.h index 5af9289a..65fc086a 100644 --- a/edgedb/datatypes/datatypes.h +++ b/edgedb/datatypes/datatypes.h @@ -85,29 +85,6 @@ PyObject * EdgeRecordDesc_List(PyObject *, uint8_t, uint8_t); -/* === edgedb.InputShape ==================================== */ - -extern PyTypeObject EdgeInputShape_Type; - -#define EdgeInputShape_Check(d) (Py_TYPE(d) == &EdgeInputShape_Type) - -typedef struct { - PyObject_HEAD - PyObject *index; - PyObject *names; - Py_ssize_t size; -} EdgeInputShapeObject; - -PyObject * EdgeInputShape_InitType(void); -PyObject * EdgeInputShape_New(PyObject *); -PyObject * EdgeInputShape_PointerName(PyObject *, Py_ssize_t); - -Py_ssize_t EdgeInputShape_GetSize(PyObject *); -edge_attr_lookup_t EdgeInputShape_Lookup(PyObject *, PyObject *, Py_ssize_t *); -PyObject * EdgeInputShape_List(PyObject *); - - - /* === edgedb.Tuple ========================================= */ #define EDGE_TUPLE_FREELIST_SIZE 500 @@ -177,31 +154,6 @@ PyObject * EdgeObject_GetItem(PyObject *, Py_ssize_t); PyObject * EdgeObject_GetID(PyObject *ob); -/* === edgedb.SparseObject ======================================== */ - -#define EDGE_SPARSE_OBJECT_FREELIST_SIZE 2000 -#define EDGE_SPARSE_OBJECT_FREELIST_MAXSAVE 20 - -extern PyTypeObject EdgeSparseObject_Type; - -#define EdgeSparseObject_Check(d) (Py_TYPE(d) == &EdgeSparseObject_Type) - -typedef struct { - PyObject_VAR_HEAD - PyObject *weakreflist; - PyObject *desc; - Py_hash_t cached_hash; - PyObject *ob_item[1]; -} EdgeSparseObject; - -PyObject * EdgeSparseObject_InitType(void); -PyObject * EdgeSparseObject_New(PyObject *); -PyObject * EdgeSparseObject_GetInputShape(PyObject *); - -int EdgeSparseObject_SetItem(PyObject *, Py_ssize_t, PyObject *); -PyObject * EdgeSparseObject_GetItem(PyObject *, Py_ssize_t); - - /* === edgedb.Set =========================================== */ extern PyTypeObject EdgeSet_Type; diff --git a/edgedb/datatypes/datatypes.pxd b/edgedb/datatypes/datatypes.pxd index 11e1c0a0..059d8933 100644 --- a/edgedb/datatypes/datatypes.pxd +++ b/edgedb/datatypes/datatypes.pxd @@ -40,24 +40,12 @@ cdef extern from "datatypes.h": MANY AT_LEAST_ONE - ctypedef enum EdgeAttrLookup "edge_attr_lookup_t": - L_ERROR - L_NOT_FOUND - L_LINKPROP - L_PROPERTY - L_LINK - object EdgeRecordDesc_InitType() object EdgeRecordDesc_New(object, object, object) object EdgeRecordDesc_PointerName(object, Py_ssize_t pos) EdgeFieldCardinality EdgeRecordDesc_PointerCardinality( object, Py_ssize_t pos) - object EdgeInputShape_InitType() - object EdgeInputShape_New(object) - object EdgeInputShape_PointerName(object, Py_ssize_t pos) - EdgeAttrLookup EdgeInputShape_Lookup(object, object, Py_ssize_t* pos) - object EdgeTuple_InitType() object EdgeTuple_New(Py_ssize_t) int EdgeTuple_SetItem(object, Py_ssize_t, object) except -1 @@ -71,11 +59,6 @@ cdef extern from "datatypes.h": int EdgeObject_SetItem(object, Py_ssize_t, object) except -1 object EdgeObject_GetRecordDesc(object) - object EdgeSparseObject_InitType() - object EdgeSparseObject_New(object); - int EdgeSparseObject_SetItem(object, Py_ssize_t, object) except -1 - object EdgeSparseObject_GetInputShape(object) - bint EdgeSet_Check(object) object EdgeSet_InitType() object EdgeSet_New(Py_ssize_t); @@ -94,17 +77,12 @@ cdef extern from "datatypes.h": cdef record_desc_new(object names, object flags, object cards) cdef record_desc_pointer_name(object desc, Py_ssize_t pos) cdef record_desc_pointer_card(object desc, Py_ssize_t pos) -cdef input_shape_new(object names) -cdef input_shape_pointer_name(object desc, Py_ssize_t pos) -cdef Py_ssize_t input_shape_get_pos(object desc, object key) except -1 cdef tuple_new(Py_ssize_t size) cdef tuple_set(object tuple, Py_ssize_t pos, object elem) cdef namedtuple_new(object desc) cdef namedtuple_set(object tuple, Py_ssize_t pos, object elem) cdef object_new(object desc) cdef object_set(object tuple, Py_ssize_t pos, object elem) -cdef sparse_object_new(object desc) -cdef sparse_object_set(object tuple, Py_ssize_t pos, object elem) cdef bint set_check(object set) cdef set_new(Py_ssize_t size) cdef set_set(object set, Py_ssize_t pos, object elem) diff --git a/edgedb/datatypes/datatypes.pyx b/edgedb/datatypes/datatypes.pyx index 7b5a4c03..848d2c5e 100644 --- a/edgedb/datatypes/datatypes.pyx +++ b/edgedb/datatypes/datatypes.pyx @@ -19,7 +19,6 @@ cimport cython cimport cpython -from libc cimport stdlib include "./relative_duration.pyx" include "./enum.pyx" @@ -97,29 +96,6 @@ cdef record_desc_pointer_card(object desc, Py_ssize_t pos): return EdgeRecordDesc_PointerCardinality(desc, pos) -cdef input_shape_new(object names): - return EdgeInputShape_New(names) - - -cdef input_shape_pointer_name(object desc, Py_ssize_t pos): - return EdgeInputShape_PointerName(desc, pos) - - -cdef Py_ssize_t input_shape_get_pos(object desc, object key) except -1: - cdef: - Py_ssize_t pos - EdgeAttrLookup res = EdgeInputShape_Lookup(desc, key, &pos) - - if res == L_ERROR: - return -1 - elif res == L_NOT_FOUND: - raise LookupError(key) - elif res == L_PROPERTY: - return pos - else: - stdlib.abort() - - cdef tuple_new(Py_ssize_t size): return EdgeTuple_New(size) @@ -144,14 +120,6 @@ cdef object_set(object obj, Py_ssize_t pos, object elem): EdgeObject_SetItem(obj, pos, elem) -cdef sparse_object_new(object desc): - return EdgeSparseObject_New(desc) - - -cdef sparse_object_set(object obj, Py_ssize_t pos, object elem): - EdgeSparseObject_SetItem(obj, pos, elem) - - cdef bint set_check(object set): return EdgeSet_Check(set) diff --git a/edgedb/datatypes/input_shape.c b/edgedb/datatypes/input_shape.c deleted file mode 100644 index 20c23b87..00000000 --- a/edgedb/datatypes/input_shape.c +++ /dev/null @@ -1,288 +0,0 @@ -/* -* This source file is part of the EdgeDB open source project. -* -* Copyright 2022-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 "internal.h" - - -static int init_type_called = 0; - - -static void -input_shape_dealloc(EdgeInputShapeObject *o) -{ - PyObject_GC_UnTrack(o); - Py_CLEAR(o->index); - Py_CLEAR(o->names); - PyObject_GC_Del(o); -} - - -static int -input_shape_traverse(EdgeInputShapeObject *o, visitproc visit, void *arg) -{ - Py_VISIT(o->index); - Py_VISIT(o->names); - return 0; -} - -static PyObject * -input_shape_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds) -{ - if (args == NULL || - PyTuple_Size(args) != 1 || - (kwds != NULL && PyDict_Size(kwds))) - { - PyErr_SetString( - PyExc_TypeError, - "InputShape accepts exactly one positional argument"); - return NULL; - } - - return EdgeInputShape_New(PyTuple_GET_ITEM(args, 0)); -} - - -static PyObject * -input_shape_get_pos(EdgeInputShapeObject *o, PyObject *arg) { - Py_ssize_t pos; - edge_attr_lookup_t ret = EdgeInputShape_Lookup((PyObject *)o, arg, &pos); - switch (ret) { - case L_ERROR: - return NULL; - - case L_NOT_FOUND: - PyErr_SetObject(PyExc_LookupError, arg); - return NULL; - - case L_PROPERTY: - return PyLong_FromLong((long)pos); - - default: - abort(); - } -} - - -static PyObject * -input_shape_dir(EdgeInputShapeObject *o, PyObject *args) -{ - PyObject *names = o->names; - Py_INCREF(names); - return names; -} - - -static PyMethodDef input_shape_methods[] = { - {"get_pos", (PyCFunction)input_shape_get_pos, METH_O, NULL}, - {"__dir__", (PyCFunction)input_shape_dir, METH_NOARGS, NULL}, - {NULL, NULL} -}; - - -PyTypeObject EdgeInputShape_Type = { - PyVarObject_HEAD_INIT(NULL, 0) - .tp_name = "edgedb.InputShape", - .tp_basicsize = sizeof(EdgeInputShapeObject), - .tp_dealloc = (destructor)input_shape_dealloc, - .tp_getattro = PyObject_GenericGetAttr, - .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, - .tp_traverse = (traverseproc)input_shape_traverse, - .tp_new = input_shape_tp_new, - .tp_methods = input_shape_methods, -}; - - -PyObject * -EdgeInputShape_New(PyObject *names) -{ - EdgeInputShapeObject *o; - - assert(init_type_called); - - if (!names || !PyTuple_CheckExact(names)) { - PyErr_SetString( - PyExc_TypeError, - "InputShape requires a tuple as its first argument"); - return NULL; - } - - if (Py_SIZE(names) > EDGE_MAX_TUPLE_SIZE) { - PyErr_Format( - PyExc_ValueError, - "EdgeDB does not supports tuples with more than %d elements", - EDGE_MAX_TUPLE_SIZE); - return NULL; - } - - Py_ssize_t size = Py_SIZE(names); - - PyObject *index = PyDict_New(); - if (index == NULL) { - return NULL; - } - - for (Py_ssize_t i = 0; i < size; i++) { - PyObject *key = PyTuple_GET_ITEM(names, i); /* borrowed */ - if (!PyUnicode_CheckExact(key)) { - PyErr_SetString( - PyExc_ValueError, - "InputShape received a non-str key"); - return NULL; - } - - PyObject *num = PyLong_FromLong(i); - if (num == NULL) { - Py_DECREF(index); - return NULL; - } - - if (PyDict_SetItem(index, key, num)) { - Py_DECREF(index); - Py_DECREF(num); - return NULL; - } - - Py_DECREF(num); - } - - o = PyObject_GC_New(EdgeInputShapeObject, &EdgeInputShape_Type); - if (o == NULL) { - Py_DECREF(index); - return NULL; - } - - o->index = index; - - Py_INCREF(names); - o->names = names; - - o->size = size; - - PyObject_GC_Track(o); - return (PyObject *)o; -} - - -edge_attr_lookup_t -EdgeInputShape_Lookup(PyObject *ob, PyObject *key, Py_ssize_t *pos) -{ - if (!EdgeInputShape_Check(ob)) { - PyErr_BadInternalCall(); - return L_ERROR; - } - - EdgeInputShapeObject *d = (EdgeInputShapeObject *)ob; - - PyObject *res = PyDict_GetItem(d->index, key); /* borrowed */ - if (res == NULL) { - if (PyErr_Occurred()) { - return L_ERROR; - } - else { - return L_NOT_FOUND; - } - } - - assert(PyLong_CheckExact(res)); - long res_long = PyLong_AsLong(res); - if (res_long < 0) { - assert(PyErr_Occurred()); - return L_ERROR; - } - assert(res_long < d->size); - *pos = res_long; - - return L_PROPERTY; -} - - -PyObject * -EdgeInputShape_PointerName(PyObject *ob, Py_ssize_t pos) -{ - if (!EdgeInputShape_Check(ob)) { - PyErr_BadInternalCall(); - return NULL; - } - EdgeInputShapeObject *o = (EdgeInputShapeObject *)ob; - PyObject * key = PyTuple_GetItem(o->names, pos); - if (key == NULL) { - return NULL; - } - Py_INCREF(key); - return key; -} - - -Py_ssize_t -EdgeInputShape_GetSize(PyObject *ob) -{ - assert(ob != NULL); - if (!EdgeInputShape_Check(ob)) { - PyErr_BadInternalCall(); - return -1; - } - EdgeInputShapeObject *o = (EdgeInputShapeObject *)ob; - return o->size; -} - - -PyObject * -EdgeInputShape_List(PyObject *ob) -{ - if (!EdgeInputShape_Check(ob)) { - PyErr_BadInternalCall(); - return NULL; - } - - EdgeInputShapeObject *o = (EdgeInputShapeObject *)ob; - - PyObject *ret = PyList_New(o->size); - if (ret == NULL) { - return NULL; - } - - for (Py_ssize_t i = 0; i < o->size; i++) { - PyObject *name = PyTuple_GetItem(o->names, i); - if (name == NULL) { - Py_DECREF(ret); - return NULL; - } - - Py_INCREF(name); - if (PyList_SetItem(ret, i, name)) { - Py_DECREF(ret); - return NULL; - } - } - - return ret; -} - - -PyObject * -EdgeInputShape_InitType(void) -{ - if (PyType_Ready(&EdgeInputShape_Type) < 0) { - return NULL; - } - - init_type_called = 1; - return (PyObject *)&EdgeInputShape_Type; -} diff --git a/edgedb/datatypes/internal.h b/edgedb/datatypes/internal.h index cddc2578..67963b41 100644 --- a/edgedb/datatypes/internal.h +++ b/edgedb/datatypes/internal.h @@ -50,9 +50,6 @@ int _EdgeGeneric_RenderValues( int _EdgeGeneric_RenderItems(_PyUnicodeWriter *, PyObject *, PyObject *, PyObject **, Py_ssize_t, int, int); -int _EdgeGeneric_RenderSparseItems(_PyUnicodeWriter *writer, - PyObject *host, PyObject *desc, - PyObject **items, Py_ssize_t len); PyObject * _EdgeGeneric_RichCompareValues(PyObject **, Py_ssize_t, PyObject **, Py_ssize_t, diff --git a/edgedb/datatypes/repr.c b/edgedb/datatypes/repr.c index 4af524f5..bfe0da4b 100644 --- a/edgedb/datatypes/repr.c +++ b/edgedb/datatypes/repr.c @@ -169,77 +169,3 @@ _EdgeGeneric_RenderItems(_PyUnicodeWriter *writer, Py_ReprLeave((PyObject *)host); return -1; } - - -int -_EdgeGeneric_RenderSparseItems(_PyUnicodeWriter *writer, - PyObject *host, PyObject *desc, - PyObject **items, Py_ssize_t len) -{ - assert(EdgeInputShape_GetSize(desc) == len); - - PyObject *item_repr = NULL; - PyObject *item_name = NULL; - int first = 1; - - int res = Py_ReprEnter(host); - if (res != 0) { - if (res > 0) { - if (_PyUnicodeWriter_WriteASCIIString(writer, "...", 3) < 0) { - return -1; - } - return 0; - } - else { - return -1; - } - } - - for (Py_ssize_t i = 0; i < len; i++) { - if (items[i] == Py_None) { - continue; - } - - item_repr = _EdgeGeneric_RenderObject(items[i]); - if (item_repr == NULL) { - goto error; - } - - item_name = EdgeInputShape_PointerName(desc, i); - if (item_name == NULL) { - goto error; - } - assert(PyUnicode_CheckExact(item_name)); - - if (first) { - first = 0; - } else { - if (_PyUnicodeWriter_WriteASCIIString(writer, ", ", 2) < 0) { - goto error; - } - } - - if (_PyUnicodeWriter_WriteStr(writer, item_name) < 0) { - goto error; - } - Py_CLEAR(item_name); - - if (_PyUnicodeWriter_WriteASCIIString(writer, " := ", 4) < 0) { - goto error; - } - - if (_PyUnicodeWriter_WriteStr(writer, item_repr) < 0) { - goto error; - } - Py_CLEAR(item_repr); - } - - Py_ReprLeave((PyObject *)host); - return 0; - -error: - Py_CLEAR(item_repr); - Py_CLEAR(item_name); - Py_ReprLeave((PyObject *)host); - return -1; -} diff --git a/edgedb/datatypes/sparse_object.c b/edgedb/datatypes/sparse_object.c deleted file mode 100644 index 8ded5d7c..00000000 --- a/edgedb/datatypes/sparse_object.c +++ /dev/null @@ -1,301 +0,0 @@ -/* -* This source file is part of the EdgeDB open source project. -* -* Copyright 2022-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; -static Py_hash_t base_hash = -1; - - -EDGE_SETUP_FREELIST( - EDGE_SPARSE_OBJECT, - EdgeSparseObject, - EDGE_SPARSE_OBJECT_FREELIST_MAXSAVE, - EDGE_SPARSE_OBJECT_FREELIST_SIZE) - - -#define EdgeSparseObject_GET_ITEM(op, i) \ - (((EdgeSparseObject *)(op))->ob_item[i]) -#define EdgeSparseObject_SET_ITEM(op, i, v) \ - (((EdgeSparseObject *)(op))->ob_item[i] = v) - - -PyObject * -EdgeSparseObject_New(PyObject *desc) -{ - assert(init_type_called); - - if (desc == NULL || !EdgeInputShape_Check(desc)) { - PyErr_BadInternalCall(); - return NULL; - } - - Py_ssize_t size = EdgeInputShape_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; - } - - EdgeSparseObject *o = NULL; - EDGE_NEW_WITH_FREELIST(EDGE_SPARSE_OBJECT, EdgeSparseObject, - &EdgeSparseObject_Type, o, size); - assert(o != NULL); - assert(Py_SIZE(o) == size); - assert(EdgeSparseObject_Check(o)); - - o->weakreflist = NULL; - - Py_INCREF(desc); - o->desc = desc; - - o->cached_hash = -1; - - PyObject_GC_Track(o); - return (PyObject *)o; -} - - -PyObject * -EdgeSparseObject_GetInputShape(PyObject *o) -{ - if (!EdgeSparseObject_Check(o)) { - PyErr_Format( - PyExc_TypeError, - "an instance of edgedb.Object expected"); - return NULL; - } - - PyObject *desc = ((EdgeSparseObject *)o)->desc; - Py_INCREF(desc); - return desc; -} - - -int -EdgeSparseObject_SetItem(PyObject *ob, Py_ssize_t i, PyObject *el) -{ - assert(EdgeSparseObject_Check(ob)); - EdgeSparseObject *o = (EdgeSparseObject *)ob; - assert(i >= 0); - assert(i < Py_SIZE(o)); - Py_INCREF(el); - EdgeSparseObject_SET_ITEM(o, i, el); - return 0; -} - - -PyObject * -EdgeSparseObject_GetItem(PyObject *ob, Py_ssize_t i) -{ - assert(EdgeSparseObject_Check(ob)); - EdgeSparseObject *o = (EdgeSparseObject *)ob; - if (i < 0 || i >= Py_SIZE(o)) { - PyErr_BadInternalCall(); - return NULL; - } - PyObject *el = EdgeSparseObject_GET_ITEM(o, i); - Py_INCREF(el); - return el; -} - - -static void -sparse_object_dealloc(EdgeSparseObject *o) -{ - PyObject_GC_UnTrack(o); - if (o->weakreflist != NULL) { - PyObject_ClearWeakRefs((PyObject*)o); - } - Py_CLEAR(o->desc); - o->cached_hash = -1; - Py_TRASHCAN_SAFE_BEGIN(o) - EDGE_DEALLOC_WITH_FREELIST(EDGE_SPARSE_OBJECT, EdgeSparseObject, o); - Py_TRASHCAN_SAFE_END(o) -} - - -static int -sparse_object_traverse(EdgeSparseObject *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 Py_hash_t -sparse_object_hash(EdgeSparseObject *o) -{ - if (o->cached_hash == -1) { - o->cached_hash = _EdgeGeneric_HashWithBase( - base_hash, o->ob_item, Py_SIZE(o)); - } - return o->cached_hash; -} - - -static PyObject * -sparse_object_getattr(EdgeSparseObject *o, PyObject *name) -{ - Py_ssize_t pos; - edge_attr_lookup_t ret = EdgeInputShape_Lookup( - (PyObject *)o->desc, name, &pos); - switch (ret) { - case L_ERROR: - return NULL; - - case L_NOT_FOUND: - return PyObject_GenericGetAttr((PyObject *)o, name); - - case L_PROPERTY: { - PyObject *val = EdgeSparseObject_GET_ITEM(o, pos); - Py_INCREF(val); - return val; - } - - default: - abort(); - } -} - -static PyObject * -sparse_object_getitem(EdgeSparseObject *o, PyObject *name) -{ - Py_ssize_t pos; - edge_attr_lookup_t ret = EdgeInputShape_Lookup( - (PyObject *)o->desc, name, &pos); - switch (ret) { - case L_ERROR: - return NULL; - - case L_PROPERTY: - PyErr_Format( - PyExc_TypeError, - "property %R should be accessed via dot notation", - name); - return NULL; - - case L_NOT_FOUND: - PyErr_Format( - PyExc_KeyError, - "link %R does not exist", - name); - return NULL; - - default: - abort(); - } - -} - - -static PyObject * -sparse_object_repr(EdgeSparseObject *o) -{ - _PyUnicodeWriter writer; - _PyUnicodeWriter_Init(&writer); - writer.overallocate = 1; - - if (_PyUnicodeWriter_WriteASCIIString(&writer, "SparseObject{", 13) < 0) { - goto error; - } - - if (_EdgeGeneric_RenderSparseItems(&writer, - (PyObject *)o, o->desc, - o->ob_item, Py_SIZE(o)) < 0) - { - goto error; - } - - if (_PyUnicodeWriter_WriteChar(&writer, '}') < 0) { - goto error; - } - - return _PyUnicodeWriter_Finish(&writer); - -error: - _PyUnicodeWriter_Dealloc(&writer); - return NULL; -} - - -static PyObject * -sparse_object_dir(EdgeSparseObject *o, PyObject *args) -{ - return EdgeInputShape_List(o->desc); -} - - -static PyMethodDef sparse_object_methods[] = { - {"__dir__", (PyCFunction)sparse_object_dir, METH_NOARGS, NULL}, - {NULL, NULL} -}; - - -static PyMappingMethods sparse_object_as_mapping = { - .mp_subscript = (binaryfunc)sparse_object_getitem, -}; - - -PyTypeObject EdgeSparseObject_Type = { - PyVarObject_HEAD_INIT(NULL, 0) - "edgedb.SparseObject", - .tp_basicsize = sizeof(EdgeSparseObject) - sizeof(PyObject *), - .tp_itemsize = sizeof(PyObject *), - .tp_dealloc = (destructor)sparse_object_dealloc, - .tp_hash = (hashfunc)sparse_object_hash, - .tp_methods = sparse_object_methods, - .tp_as_mapping = &sparse_object_as_mapping, - .tp_getattro = (getattrofunc)sparse_object_getattr, - .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, - .tp_traverse = (traverseproc)sparse_object_traverse, - .tp_free = PyObject_GC_Del, - .tp_repr = (reprfunc)sparse_object_repr, - .tp_weaklistoffset = offsetof(EdgeSparseObject, weakreflist), -}; - - -PyObject * -EdgeSparseObject_InitType(void) -{ - if (PyType_Ready(&EdgeSparseObject_Type) < 0) { - return NULL; - } - - base_hash = _EdgeGeneric_HashString("edgedb.SparseObject"); - if (base_hash == -1) { - return NULL; - } - - init_type_called = 1; - return (PyObject *)&EdgeSparseObject_Type; -} diff --git a/edgedb/protocol/codecs/codecs.pxd b/edgedb/protocol/codecs/codecs.pxd index ac9a8963..d64ff1ff 100644 --- a/edgedb/protocol/codecs/codecs.pxd +++ b/edgedb/protocol/codecs/codecs.pxd @@ -22,7 +22,6 @@ include "./scalar.pxd" include "./tuple.pxd" include "./namedtuple.pxd" include "./object.pxd" -include "./sparse_object.pxd" include "./array.pxd" include "./set.pxd" include "./enum.pxd" diff --git a/edgedb/protocol/codecs/codecs.pyx b/edgedb/protocol/codecs/codecs.pyx index 245ccfe3..cb74ee1b 100644 --- a/edgedb/protocol/codecs/codecs.pyx +++ b/edgedb/protocol/codecs/codecs.pyx @@ -31,7 +31,6 @@ include "./scalar.pyx" include "./tuple.pyx" include "./namedtuple.pyx" include "./object.pyx" -include "./sparse_object.pyx" include "./array.pyx" include "./set.pyx" include "./enum.pyx" @@ -107,7 +106,7 @@ cdef class CodecsRegistry: if t == CTYPE_SET: frb_read(spec, 2) - elif t == CTYPE_SHAPE: + elif t == CTYPE_SHAPE or t == CTYPE_INPUT_SHAPE: els = hton.unpack_int16(frb_read(spec, 2)) for i in range(els): frb_read(spec, 4) # flags @@ -149,13 +148,6 @@ cdef class CodecsRegistry: str_len = hton.unpack_uint32(frb_read(spec, 4)) frb_read(spec, str_len) - elif t == CTYPE_INPUT_SHAPE: - els = hton.unpack_int16(frb_read(spec, 2)) - for i in range(els): - str_len = hton.unpack_uint32(frb_read(spec, 4)) - # read the (`str_len` bytes) and (2 bytes) - frb_read(spec, str_len + 2) - elif (t >= 0x7f and t <= 0xff): # Ignore all type annotations. str_len = hton.unpack_uint32(frb_read(spec, 4)) @@ -173,7 +165,7 @@ cdef class CodecsRegistry: sub_codec = codecs_list[pos] res = SetCodec.new(tid, sub_codec) - elif t == CTYPE_SHAPE: + elif t == CTYPE_SHAPE or t == CTYPE_INPUT_SHAPE: els = hton.unpack_int16(frb_read(spec, 2)) codecs = cpython.PyTuple_New(els) names = cpython.PyTuple_New(els) @@ -201,7 +193,9 @@ cdef class CodecsRegistry: cpython.Py_INCREF(cardinality) cpython.PyTuple_SetItem(cards, i, cardinality) - res = ObjectCodec.new(tid, names, flags, cards, codecs) + res = ObjectCodec.new( + tid, names, flags, cards, codecs, t == CTYPE_INPUT_SHAPE + ) elif t == CTYPE_BASE_SCALAR: if tid in self.base_codec_overrides: @@ -273,25 +267,6 @@ cdef class CodecsRegistry: sub_codec = codecs_list[pos] res = ArrayCodec.new(tid, sub_codec, dim_len) - elif t == CTYPE_INPUT_SHAPE: - 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 = SparseObjectCodec.new(tid, names, codecs) - else: raise NotImplementedError( f'no codec implementation for EdgeDB data class {t}') diff --git a/edgedb/protocol/codecs/object.pxd b/edgedb/protocol/codecs/object.pxd index 0b01f3aa..8524f55c 100644 --- a/edgedb/protocol/codecs/object.pxd +++ b/edgedb/protocol/codecs/object.pxd @@ -19,9 +19,10 @@ @cython.final cdef class ObjectCodec(BaseNamedRecordCodec): + cdef bint is_sparse cdef encode_args(self, WriteBuffer buf, dict obj) @staticmethod cdef BaseCodec new(bytes tid, tuple names, tuple flags, - tuple cards, tuple codecs) + tuple cards, tuple codecs, bint is_sparse) diff --git a/edgedb/protocol/codecs/object.pyx b/edgedb/protocol/codecs/object.pyx index 14931af5..da879ded 100644 --- a/edgedb/protocol/codecs/object.pyx +++ b/edgedb/protocol/codecs/object.pyx @@ -21,7 +21,40 @@ cdef class ObjectCodec(BaseNamedRecordCodec): cdef encode(self, WriteBuffer buf, object obj): - raise NotImplementedError + cdef: + WriteBuffer elem_data + Py_ssize_t objlen = 0 + Py_ssize_t i + BaseCodec sub_codec + descriptor = (self).descriptor + + if not self.is_sparse: + raise NotImplementedError + + elem_data = WriteBuffer.new() + for name, arg in obj.items(): + if arg is not None: + try: + i = descriptor.get_pos(name) + except LookupError: + raise self._make_missing_args_error_message(obj) from None + objlen += 1 + elem_data.write_int32(i) + + sub_codec = (self.fields_codecs[i]) + try: + sub_codec.encode(elem_data, arg) + except (TypeError, ValueError) as e: + value_repr = repr(arg) + if len(value_repr) > 40: + value_repr = value_repr[:40] + '...' + raise errors.InvalidArgumentError( + 'invalid input for session argument ' + f' {name} := {value_repr} ({e})') from e + + buf.write_int32(4 + elem_data.len()) # buffer length + buf.write_int32(objlen) + buf.write_buffer(elem_data) cdef encode_args(self, WriteBuffer buf, dict obj): cdef: @@ -31,6 +64,9 @@ cdef class ObjectCodec(BaseNamedRecordCodec): BaseCodec sub_codec descriptor = (self).descriptor + if self.is_sparse: + raise NotImplementedError + self._check_encoder() objlen = len(obj) @@ -83,15 +119,17 @@ cdef class ObjectCodec(BaseNamedRecordCodec): passed_args = set(args.keys()) missed_args = required_args - passed_args extra_args = passed_args - required_args + required = 'acceptable' if self.is_sparse else 'expected' - error_message = f'expected {required_args} arguments' + error_message = f'{required} {required_args} arguments' passed_args_repr = repr(passed_args) if passed_args else 'nothing' error_message += f', got {passed_args_repr}' - missed_args = set(required_args) - set(passed_args) - if missed_args: - error_message += f', missed {missed_args}' + if not self.is_sparse: + missed_args = set(required_args) - set(passed_args) + if missed_args: + error_message += f', missed {missed_args}' extra_args = set(passed_args) - set(required_args) if extra_args: @@ -110,6 +148,9 @@ cdef class ObjectCodec(BaseNamedRecordCodec): tuple fields_codecs = (self).fields_codecs descriptor = (self).descriptor + if self.is_sparse: + raise NotImplementedError + elem_count = hton.unpack_int32(frb_read(buf, 4)) if elem_count != len(fields_codecs): @@ -140,14 +181,18 @@ cdef class ObjectCodec(BaseNamedRecordCodec): @staticmethod cdef BaseCodec new(bytes tid, tuple names, tuple flags, tuple cards, - tuple codecs): + tuple codecs, bint is_sparse): cdef: ObjectCodec codec codec = ObjectCodec.__new__(ObjectCodec) codec.tid = tid - codec.name = 'Object' + if is_sparse: + codec.name = 'SparseObject' + else: + codec.name = 'Object' + codec.is_sparse = is_sparse codec.descriptor = datatypes.record_desc_new(names, flags, cards) codec.fields_codecs = codecs diff --git a/edgedb/protocol/codecs/sparse_object.pxd b/edgedb/protocol/codecs/sparse_object.pxd deleted file mode 100644 index 54068fc5..00000000 --- a/edgedb/protocol/codecs/sparse_object.pxd +++ /dev/null @@ -1,24 +0,0 @@ -# -# This source file is part of the EdgeDB open source project. -# -# Copyright 2022-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 SparseObjectCodec(BaseNamedRecordCodec): - - @staticmethod - cdef BaseCodec new(bytes tid, tuple names, tuple codecs) diff --git a/edgedb/protocol/codecs/sparse_object.pyx b/edgedb/protocol/codecs/sparse_object.pyx deleted file mode 100644 index 635a1a3d..00000000 --- a/edgedb/protocol/codecs/sparse_object.pyx +++ /dev/null @@ -1,126 +0,0 @@ -# -# This source file is part of the EdgeDB open source project. -# -# Copyright 2022-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 SparseObjectCodec(BaseNamedRecordCodec): - - cdef encode(self, WriteBuffer buf, object obj): - cdef: - WriteBuffer elem_data - Py_ssize_t objlen = 0 - Py_ssize_t i - BaseCodec sub_codec - descriptor = (self).descriptor - - elem_data = WriteBuffer.new() - for name, arg in obj.items(): - if arg is not None: - try: - i = datatypes.input_shape_get_pos(descriptor, name) - except LookupError: - raise self._make_unknown_args_error_message(obj) from None - objlen += 1 - elem_data.write_int32(i) - - sub_codec = (self.fields_codecs[i]) - try: - sub_codec.encode(elem_data, arg) - except (TypeError, ValueError) as e: - value_repr = repr(arg) - if len(value_repr) > 40: - value_repr = value_repr[:40] + '...' - raise errors.InvalidArgumentError( - 'invalid input for session argument ' - f' {name} := {value_repr} ({e})') from e - - buf.write_int32(4 + elem_data.len()) # buffer length - buf.write_int32(objlen) - buf.write_buffer(elem_data) - - def _make_unknown_args_error_message(self, args): - cdef descriptor = (self).descriptor - - acceptable_args = set() - - for i in range(len(self.fields_codecs)): - name = datatypes.input_shape_pointer_name(descriptor, i) - acceptable_args.add(name) - - passed_args = set(args.keys()) - error_message = f'acceptable {acceptable_args} arguments' - - passed_args_repr = repr(passed_args) if passed_args else 'nothing' - error_message += f', got {passed_args_repr}' - - extra_args = set(passed_args) - set(acceptable_args) - if extra_args: - error_message += f', extra {extra_args}' - - return errors.QueryArgumentError(error_message) - - 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)) - - result = datatypes.sparse_object_new(descriptor) - - for i in range(len(fields_codecs)): - datatypes.sparse_object_set(result, i, None) - - for _ in range(elem_count): - i = hton.unpack_int32(frb_read(buf, 4)) - elem_len = hton.unpack_int32(frb_read(buf, 4)) - - if elem_len == -1: - continue - - 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'object element decoding: {frb_get_len(&elem_buf)}') - cpython.Py_DECREF(None) - datatypes.sparse_object_set(result, i, elem) - - return result - - @staticmethod - cdef BaseCodec new(bytes tid, tuple names, tuple codecs): - cdef: - SparseObjectCodec codec - - codec = SparseObjectCodec.__new__(SparseObjectCodec) - - codec.tid = tid - codec.name = 'SparseObject' - codec.descriptor = datatypes.input_shape_new(names) - codec.fields_codecs = codecs - - return codec diff --git a/setup.py b/setup.py index cf95bf4a..a345d7f4 100644 --- a/setup.py +++ b/setup.py @@ -302,11 +302,9 @@ def finalize_options(self): "edgedb.datatypes.datatypes", ["edgedb/datatypes/args.c", "edgedb/datatypes/record_desc.c", - "edgedb/datatypes/input_shape.c", "edgedb/datatypes/tuple.c", "edgedb/datatypes/namedtuple.c", "edgedb/datatypes/object.c", - "edgedb/datatypes/sparse_object.c", "edgedb/datatypes/set.c", "edgedb/datatypes/hash.c", "edgedb/datatypes/array.c", From 307f0c88e0bd78f23b491072efc98a268b1de070 Mon Sep 17 00:00:00 2001 From: Fantix King Date: Mon, 20 Jun 2022 11:49:46 -0400 Subject: [PATCH 03/11] Fix issue sending wrong empty state --- edgedb/protocol/protocol.pyx | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/edgedb/protocol/protocol.pyx b/edgedb/protocol/protocol.pyx index f9e0ec04..d3516d34 100644 --- a/edgedb/protocol/protocol.pyx +++ b/edgedb/protocol/protocol.pyx @@ -150,7 +150,7 @@ cdef class SansIOProtocol: self.reset_status() self.protocol_version = (PROTO_VER_MAJOR, 0) - self.state_type_id = b'\0' * 16 + self.state_type_id = NULL_CODEC_ID self.state_codec = None self.state = None self.user_state = None @@ -372,11 +372,13 @@ cdef class SansIOProtocol: buf.write_buffer(params) - buf.write_bytes(self.state_type_id) - buf.write_int16(1) if self.state is None: + buf.write_bytes(NULL_CODEC_ID) + buf.write_int16(1) buf.write_int32(0) else: + buf.write_bytes(self.state_type_id) + buf.write_int16(1) buf.write_buffer(self.state) buf.write_bytes(in_dc.get_tid()) From 6ef56f3c36b4c25bc3508b84600874e8e690c98e Mon Sep 17 00:00:00 2001 From: Fantix King Date: Mon, 20 Jun 2022 11:59:45 -0400 Subject: [PATCH 04/11] Rename configs to config and minimize state size --- edgedb/options.py | 50 ++++++++++++++++++++++++++--------------------- 1 file changed, 28 insertions(+), 22 deletions(-) diff --git a/edgedb/options.py b/edgedb/options.py index c680e005..487d4bb7 100644 --- a/edgedb/options.py +++ b/edgedb/options.py @@ -110,41 +110,41 @@ def __repr__(self): class Session: - __slots__ = ['_module', '_aliases', '_configs', '_globals'] + __slots__ = ['_module', '_aliases', '_config', '_globals'] def __init__( self, - module: str = 'default', + module: typing.Optional[str] = None, aliases: typing.Mapping[str, str] = None, - configs: typing.Mapping[str, typing.Any] = None, + config: typing.Mapping[str, typing.Any] = None, globals_: typing.Mapping[str, typing.Any] = None, ): self._module = module self._aliases = {} if aliases is None else dict(aliases) - self._configs = {} if configs is None else dict(configs) + self._config = {} if config is None else dict(config) self._globals = {} if globals_ is None else dict(globals_) @classmethod def defaults(cls): return cls() - def with_aliases(self, module=None, **aliases): + def with_aliases(self, module=..., **aliases): new_aliases = self._aliases.copy() new_aliases.update(aliases) return Session( - module=self._module if module is None else module, + module=self._module if module is ... else module, aliases=new_aliases, - configs=self._configs, + config=self._config, globals_=self._globals, ) - def with_configs(self, **configs): - new_configs = self._configs.copy() - new_configs.update(configs) + def with_config(self, **config): + new_config = self._config.copy() + new_config.update(config) return Session( module=self._module, aliases=self._aliases, - configs=new_configs, + config=new_config, globals_=self._globals, ) @@ -154,20 +154,26 @@ def with_globals(self, **globals_): return Session( module=self._module, aliases=self._aliases, - configs=self._configs, + config=self._config, globals_=new_globals, ) def as_dict(self): - return { - "module": self._module, - "aliases": list(self._aliases.items()), - "config": self._configs, - "globals": { - (k if '::' in k else f'{self._module}::{k}'): v + rv = {} + if self._module is not None: + module = rv["module"] = self._module + else: + module = 'default' + if self._aliases: + rv["aliases"] = list(self._aliases.items()) + if self._config: + rv["config"] = self._config + if self._globals: + rv["globals"] = { + (k if '::' in k else f'{module}::{k}'): v for k, v in self._globals.items() - }, - } + } + return rv class _OptionsMixin: @@ -227,10 +233,10 @@ def with_aliases(self, module=None, **aliases): ) return result - def with_configs(self, **configs): + def with_config(self, **config): result = self._shallow_clone() result._options = self._options.with_session( - self._options.session.with_configs(**configs) + self._options.session.with_config(**config) ) return result From 563ba85e770b6ba67d1750dc83591664cd54ca4d Mon Sep 17 00:00:00 2001 From: Fantix King Date: Mon, 20 Jun 2022 12:13:13 -0400 Subject: [PATCH 05/11] Support using aliases in with_globals() --- edgedb/options.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/edgedb/options.py b/edgedb/options.py index 487d4bb7..0b0df1aa 100644 --- a/edgedb/options.py +++ b/edgedb/options.py @@ -169,10 +169,19 @@ def as_dict(self): if self._config: rv["config"] = self._config if self._globals: - rv["globals"] = { - (k if '::' in k else f'{module}::{k}'): v - for k, v in self._globals.items() - } + rv["globals"] = g = {} + for k, v in self._globals.items(): + parts = k.split("::") + if len(parts) == 1: + g[f"{module}::{k}"] = v + elif len(parts) == 2: + mod, glob = parts + mod = self._aliases.get(mod, mod) + g[f"{mod}::{glob}"] = v + else: + raise errors.InvalidArgumentError( + f"Illegal global name: {k}" + ) return rv From f320bc6a2ef36cb19452d3f60be91e122b218816 Mon Sep 17 00:00:00 2001 From: Fantix King Date: Mon, 20 Jun 2022 12:37:56 -0400 Subject: [PATCH 06/11] Handle CommandCompleteWithConsequence message Also sync errors --- edgedb/errors/__init__.py | 15 +++++++++++++++ edgedb/protocol/consts.pxi | 1 + edgedb/protocol/protocol.pxd | 1 + edgedb/protocol/protocol.pyx | 26 ++++++++++++++++++++++++++ 4 files changed, 43 insertions(+) diff --git a/edgedb/errors/__init__.py b/edgedb/errors/__init__.py index f2bd92ee..9320992c 100644 --- a/edgedb/errors/__init__.py +++ b/edgedb/errors/__init__.py @@ -27,6 +27,9 @@ 'CapabilityError', 'UnsupportedCapabilityError', 'DisabledCapabilityError', + 'StateError', + 'StateSerializationError', + 'StateMismatchError', 'QueryError', 'InvalidSyntaxError', 'EdgeQLSyntaxError', @@ -162,6 +165,18 @@ class DisabledCapabilityError(CapabilityError): _code = 0x_03_04_02_00 +class StateError(ProtocolError): + _code = 0x_03_05_00_00 + + +class StateSerializationError(StateError): + _code = 0x_03_05_00_01 + + +class StateMismatchError(StateError): + _code = 0x_03_05_00_02 + + class QueryError(EdgeDBError): _code = 0x_04_00_00_00 diff --git a/edgedb/protocol/consts.pxi b/edgedb/protocol/consts.pxi index 96130e0f..f158203e 100644 --- a/edgedb/protocol/consts.pxi +++ b/edgedb/protocol/consts.pxi @@ -27,6 +27,7 @@ DEF READY_FOR_COMMAND_MSG = b'Z' DEF SYNC_MSG = b'S' DEF FLUSH_MSG = b'H' DEF COMMAND_COMPLETE_MSG = b'C' +DEF COMMAND_COMPLETE_WITH_CONSEQUENCE_MSG = b'c' DEF DATA_MSG = b'D' DEF COMMAND_DATA_DESC_MSG = b'T' DEF LOG_MSG = b'L' diff --git a/edgedb/protocol/protocol.pxd b/edgedb/protocol/protocol.pxd index 282caef9..50a31496 100644 --- a/edgedb/protocol/protocol.pxd +++ b/edgedb/protocol/protocol.pxd @@ -113,6 +113,7 @@ cdef class SansIOProtocol: cdef parse_data_messages(self, BaseCodec out_dc, result) cdef parse_sync_message(self) cdef parse_command_complete_message(self) + cdef parse_command_complete_with_conseq_message(self) cdef parse_describe_type_message(self, CodecsRegistry reg) cdef parse_type_data(self, CodecsRegistry reg) cdef _amend_parse_error( diff --git a/edgedb/protocol/protocol.pyx b/edgedb/protocol/protocol.pyx index d3516d34..ff089590 100644 --- a/edgedb/protocol/protocol.pyx +++ b/edgedb/protocol/protocol.pyx @@ -444,6 +444,11 @@ cdef class SansIOProtocol: elif mtype == COMMAND_COMPLETE_MSG: self.parse_command_complete_message() + elif mtype == COMMAND_COMPLETE_WITH_CONSEQUENCE_MSG: + ex = self.parse_command_complete_with_conseq_message() + if not isinstance(ex, errors.StateSerializationError): + exc = ex + elif mtype == ERROR_RESPONSE_MSG: exc = self.parse_error_message() exc = self._amend_parse_error( @@ -1189,6 +1194,27 @@ cdef class SansIOProtocol: self.buffer.read_len_prefixed_bytes() # state self.buffer.finish_message() + cdef parse_command_complete_with_conseq_message(self): + assert ( + self.buffer.get_message_type() == + COMMAND_COMPLETE_WITH_CONSEQUENCE_MSG + ) + self.ignore_headers() + self.last_capabilities = enums.Capability(self.buffer.read_int64()) + self.last_status = self.buffer.read_len_prefixed_bytes() + + code = self.buffer.read_int32() + msg = self.buffer.read_len_prefixed_utf8() + attrs = self.parse_error_headers() + self.buffer.finish_message() + + # It's safe to always map error codes as we don't reuse them + code = OLD_ERROR_CODES.get(code, code) + + exc = errors.EdgeDBError._from_code(code, msg) + exc._attrs = attrs + return exc + cdef parse_sync_message(self): cdef char status From b74ad353dfe8ca2ba6014e699431db76c94c111e Mon Sep 17 00:00:00 2001 From: Fantix King Date: Tue, 21 Jun 2022 15:48:26 -0400 Subject: [PATCH 07/11] Apply async session_state_description for AsyncClient --- edgedb/protocol/protocol.pyx | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/edgedb/protocol/protocol.pyx b/edgedb/protocol/protocol.pyx index ff089590..6b24d2c1 100644 --- a/edgedb/protocol/protocol.pyx +++ b/edgedb/protocol/protocol.pyx @@ -205,6 +205,10 @@ cdef class SansIOProtocol: if user_state is None: self.state = None else: + # Apply async session_state_description for AsyncClient + while self.buffer.take_message(): + self.fallthrough() + self.state_codec.encode(buf, user_state.as_dict()) self.state = buf self.user_state = user_state From d35ca77b5bc0467d71c8220731aeb65f85f5f29f Mon Sep 17 00:00:00 2001 From: Fantix King Date: Thu, 23 Jun 2022 16:40:09 -0400 Subject: [PATCH 08/11] Forbid SESSION_CONFIG capability This disallows EdgeQL commands like `configure session`, `set alias`, `set module`, `set global` and their `reset` variants. --- edgedb/enums.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/edgedb/enums.py b/edgedb/enums.py index 96859bde..f3fc3f60 100644 --- a/edgedb/enums.py +++ b/edgedb/enums.py @@ -10,8 +10,8 @@ class Capability(enum.IntFlag): DDL = 1 << 3 # noqa PERSISTENT_CONFIG = 1 << 4 # noqa - ALL = 0xFFFF_FFFF_FFFF_FFFF # noqa - EXECUTE = ALL & ~TRANSACTION # noqa + ALL = 0xFFFF_FFFF_FFFF_FFFF # noqa + EXECUTE = ALL & ~TRANSACTION & ~SESSION_CONFIG # noqa class CompilationFlag(enum.IntFlag): From 81fb783b6985af3dfdf7dbd4d6614a66239e3461 Mon Sep 17 00:00:00 2001 From: Fantix King Date: Thu, 23 Jun 2022 16:56:36 -0400 Subject: [PATCH 09/11] Allow legacy protocol use SESSION_CONFIG capability --- edgedb/base_client.py | 6 ++++-- edgedb/enums.py | 1 + 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/edgedb/base_client.py b/edgedb/base_client.py index cfbd23f5..c66ff631 100644 --- a/edgedb/base_client.py +++ b/edgedb/base_client.py @@ -211,9 +211,11 @@ async def raw_query(self, query_context: abstract.QueryContext): await self.connect(single_attempt=True) if self._protocol.is_legacy: execute = self._protocol.legacy_execute_anonymous + allow_capabilities = enums.Capability.LEGACY_EXECUTE else: execute = self._protocol.query self._protocol.set_state(query_context.session) + allow_capabilities = enums.Capability.EXECUTE return await execute( query=query_context.query.query, args=query_context.query.args, @@ -223,7 +225,7 @@ async def raw_query(self, query_context: abstract.QueryContext): output_format=query_context.query_options.output_format, expect_one=query_context.query_options.expect_one, required_one=query_context.query_options.required_one, - allow_capabilities=enums.Capability.EXECUTE, + allow_capabilities=allow_capabilities, ) except errors.EdgeDBError as e: if query_context.retry_options is None: @@ -262,7 +264,7 @@ async def _execute(self, script: abstract.ScriptContext) -> None: "Legacy protocol doesn't support arguments in execute()" ) await self._protocol.legacy_simple_query( - script.query.query, enums.Capability.EXECUTE + script.query.query, enums.Capability.LEGACY_EXECUTE ) else: self._protocol.set_state(script.session) diff --git a/edgedb/enums.py b/edgedb/enums.py index f3fc3f60..f3306192 100644 --- a/edgedb/enums.py +++ b/edgedb/enums.py @@ -12,6 +12,7 @@ class Capability(enum.IntFlag): ALL = 0xFFFF_FFFF_FFFF_FFFF # noqa EXECUTE = ALL & ~TRANSACTION & ~SESSION_CONFIG # noqa + LEGACY_EXECUTE = ALL & ~TRANSACTION # noqa class CompilationFlag(enum.IntFlag): From b3081c63f6ef42e381314c15b8c4fc752cb24547 Mon Sep 17 00:00:00 2001 From: Fantix King Date: Fri, 24 Jun 2022 14:29:55 -0400 Subject: [PATCH 10/11] Drop COMMAND_COMPLETE_WITH_CONSEQUENCE_MSG and sync errors --- edgedb/errors/__init__.py | 21 ++++++--------------- edgedb/protocol/consts.pxi | 1 - edgedb/protocol/protocol.pxd | 1 - edgedb/protocol/protocol.pyx | 26 -------------------------- 4 files changed, 6 insertions(+), 43 deletions(-) diff --git a/edgedb/errors/__init__.py b/edgedb/errors/__init__.py index 9320992c..87117443 100644 --- a/edgedb/errors/__init__.py +++ b/edgedb/errors/__init__.py @@ -23,13 +23,11 @@ 'UnexpectedMessageError', 'InputDataError', 'ParameterTypeMismatchError', + 'StateMismatchError', 'ResultCardinalityMismatchError', 'CapabilityError', 'UnsupportedCapabilityError', 'DisabledCapabilityError', - 'StateError', - 'StateSerializationError', - 'StateMismatchError', 'QueryError', 'InvalidSyntaxError', 'EdgeQLSyntaxError', @@ -149,6 +147,10 @@ class ParameterTypeMismatchError(InputDataError): _code = 0x_03_02_01_00 +class StateMismatchError(InputDataError): + _code = 0x_03_02_02_00 + + class ResultCardinalityMismatchError(ProtocolError): _code = 0x_03_03_00_00 @@ -165,18 +167,6 @@ class DisabledCapabilityError(CapabilityError): _code = 0x_03_04_02_00 -class StateError(ProtocolError): - _code = 0x_03_05_00_00 - - -class StateSerializationError(StateError): - _code = 0x_03_05_00_01 - - -class StateMismatchError(StateError): - _code = 0x_03_05_00_02 - - class QueryError(EdgeDBError): _code = 0x_04_00_00_00 @@ -502,3 +492,4 @@ class NoDataError(ClientError): class InternalClientError(ClientError): _code = 0x_FF_04_00_00 + diff --git a/edgedb/protocol/consts.pxi b/edgedb/protocol/consts.pxi index f158203e..96130e0f 100644 --- a/edgedb/protocol/consts.pxi +++ b/edgedb/protocol/consts.pxi @@ -27,7 +27,6 @@ DEF READY_FOR_COMMAND_MSG = b'Z' DEF SYNC_MSG = b'S' DEF FLUSH_MSG = b'H' DEF COMMAND_COMPLETE_MSG = b'C' -DEF COMMAND_COMPLETE_WITH_CONSEQUENCE_MSG = b'c' DEF DATA_MSG = b'D' DEF COMMAND_DATA_DESC_MSG = b'T' DEF LOG_MSG = b'L' diff --git a/edgedb/protocol/protocol.pxd b/edgedb/protocol/protocol.pxd index 50a31496..282caef9 100644 --- a/edgedb/protocol/protocol.pxd +++ b/edgedb/protocol/protocol.pxd @@ -113,7 +113,6 @@ cdef class SansIOProtocol: cdef parse_data_messages(self, BaseCodec out_dc, result) cdef parse_sync_message(self) cdef parse_command_complete_message(self) - cdef parse_command_complete_with_conseq_message(self) cdef parse_describe_type_message(self, CodecsRegistry reg) cdef parse_type_data(self, CodecsRegistry reg) cdef _amend_parse_error( diff --git a/edgedb/protocol/protocol.pyx b/edgedb/protocol/protocol.pyx index 6b24d2c1..b93f9de0 100644 --- a/edgedb/protocol/protocol.pyx +++ b/edgedb/protocol/protocol.pyx @@ -448,11 +448,6 @@ cdef class SansIOProtocol: elif mtype == COMMAND_COMPLETE_MSG: self.parse_command_complete_message() - elif mtype == COMMAND_COMPLETE_WITH_CONSEQUENCE_MSG: - ex = self.parse_command_complete_with_conseq_message() - if not isinstance(ex, errors.StateSerializationError): - exc = ex - elif mtype == ERROR_RESPONSE_MSG: exc = self.parse_error_message() exc = self._amend_parse_error( @@ -1198,27 +1193,6 @@ cdef class SansIOProtocol: self.buffer.read_len_prefixed_bytes() # state self.buffer.finish_message() - cdef parse_command_complete_with_conseq_message(self): - assert ( - self.buffer.get_message_type() == - COMMAND_COMPLETE_WITH_CONSEQUENCE_MSG - ) - self.ignore_headers() - self.last_capabilities = enums.Capability(self.buffer.read_int64()) - self.last_status = self.buffer.read_len_prefixed_bytes() - - code = self.buffer.read_int32() - msg = self.buffer.read_len_prefixed_utf8() - attrs = self.parse_error_headers() - self.buffer.finish_message() - - # It's safe to always map error codes as we don't reuse them - code = OLD_ERROR_CODES.get(code, code) - - exc = errors.EdgeDBError._from_code(code, msg) - exc._attrs = attrs - return exc - cdef parse_sync_message(self): cdef char status From 0e73ddfe9ea6e9a2b29fbefb549c32f60f2f9748 Mon Sep 17 00:00:00 2001 From: Fantix King Date: Mon, 27 Jun 2022 11:10:48 -0400 Subject: [PATCH 11/11] Update protocol: drop the array in state Also renamed `session_state_description` to `state_description`. --- edgedb/protocol/protocol.pyx | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/edgedb/protocol/protocol.pyx b/edgedb/protocol/protocol.pyx index b93f9de0..1f217a0b 100644 --- a/edgedb/protocol/protocol.pyx +++ b/edgedb/protocol/protocol.pyx @@ -205,7 +205,7 @@ cdef class SansIOProtocol: if user_state is None: self.state = None else: - # Apply async session_state_description for AsyncClient + # Apply async state_description for AsyncClient while self.buffer.take_message(): self.fallthrough() @@ -376,22 +376,20 @@ cdef class SansIOProtocol: buf.write_buffer(params) - if self.state is None: - buf.write_bytes(NULL_CODEC_ID) - buf.write_int16(1) - buf.write_int32(0) - else: - buf.write_bytes(self.state_type_id) - buf.write_int16(1) - buf.write_buffer(self.state) - buf.write_bytes(in_dc.get_tid()) buf.write_bytes(out_dc.get_tid()) + buf.write_bytes( + NULL_CODEC_ID if self.state is None else self.state_type_id + ) if not isinstance(in_dc, NullCodec): self.encode_args(in_dc, buf, args, kwargs) else: buf.write_bytes(EMPTY_NULL_DATA) + if self.state is not None: + buf.write_buffer(self.state) + else: + buf.write_bytes(EMPTY_NULL_DATA) buf.end_message() @@ -1011,7 +1009,7 @@ cdef class SansIOProtocol: data = buf.read_len_prefixed_bytes() self.server_settings[name] = self.parse_system_config(codec, data) - elif name == 'session_state_description': + elif name == 'state_description': self.state_type_id = typedesc_id = val[:16] typedesc = val[16 + 4:] @@ -1189,7 +1187,6 @@ cdef class SansIOProtocol: self.last_capabilities = enums.Capability(self.buffer.read_int64()) self.last_status = self.buffer.read_len_prefixed_bytes() self.buffer.read_bytes(16) # state type id - assert self.buffer.read_int16() == 1 self.buffer.read_len_prefixed_bytes() # state self.buffer.finish_message()