From 520d7c806b99e9ffd2daf0a8b8acf0da92949c9b Mon Sep 17 00:00:00 2001 From: Elvis Pranskevichus Date: Thu, 20 Apr 2023 08:28:23 -0700 Subject: [PATCH 1/2] Implement support for new type descriptor protocol --- edgedb/protocol/codecs/codecs.pyx | 179 ++++++++++++++++++++++++++++-- edgedb/protocol/consts.pxi | 6 +- edgedb/protocol/protocol.pyx | 13 +-- 3 files changed, 176 insertions(+), 22 deletions(-) diff --git a/edgedb/protocol/codecs/codecs.pyx b/edgedb/protocol/codecs/codecs.pyx index d0cfa7b5..9e8764f0 100644 --- a/edgedb/protocol/codecs/codecs.pyx +++ b/edgedb/protocol/codecs/codecs.pyx @@ -49,6 +49,8 @@ DEF CTYPE_ARRAY = 6 DEF CTYPE_ENUM = 7 DEF CTYPE_INPUT_SHAPE = 8 DEF CTYPE_RANGE = 9 +DEF CTYPE_OBJECT = 10 +DEF CTYPE_COMPOUND = 11 DEF CTYPE_ANNO_TYPENAME = 255 DEF _CODECS_BUILD_CACHE_SIZE = 200 @@ -91,8 +93,9 @@ cdef class CodecsRegistry: cdef BaseCodec _build_codec(self, FRBuffer *spec, list codecs_list, protocol_version): cdef: - uint8_t t = (frb_read(spec, 1)[0]) - bytes tid = frb_read(spec, 16)[:16] + uint32_t desc_len = 0 + uint8_t t + bytes tid uint16_t els uint16_t i uint32_t str_len @@ -101,12 +104,21 @@ cdef class CodecsRegistry: BaseCodec res BaseCodec sub_codec + if protocol_version >= (2, 0): + desc_len = frb_get_len(spec) - 16 - 1 + + t = (frb_read(spec, 1)[0]) + tid = frb_read(spec, 16)[:16] + res = self.codecs.get(tid, None) if res is None: res = self.codecs_build_cache.get(tid, None) if res is not None: # We have a codec for this "tid"; advance the buffer # so that we can process the next codec. + if desc_len > 0: + frb_read(spec, desc_len) + return res if t == CTYPE_SET: frb_read(spec, 2) @@ -179,7 +191,50 @@ cdef class CodecsRegistry: sub_codec = codecs_list[pos] res = SetCodec.new(tid, sub_codec) - elif t == CTYPE_SHAPE or t == CTYPE_INPUT_SHAPE: + elif t == CTYPE_SHAPE: + if protocol_version >= (2, 0): + ephemeral_free_shape = frb_read(spec, 1)[0] + objtype_pos = hton.unpack_int16(frb_read(spec, 2)) + + els = hton.unpack_int16(frb_read(spec, 2)) + codecs = cpython.PyTuple_New(els) + names = cpython.PyTuple_New(els) + flags = cpython.PyTuple_New(els) + cards = cpython.PyTuple_New(els) + for i in range(els): + flag = hton.unpack_uint32(frb_read(spec, 4)) # flags + cardinality = frb_read(spec, 1)[0] + + 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)) + + if flag & datatypes._EDGE_POINTER_IS_LINKPROP: + name = "@" + name + 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) + + cpython.Py_INCREF(flag) + cpython.PyTuple_SetItem(flags, i, flag) + + cpython.Py_INCREF(cardinality) + cpython.PyTuple_SetItem(cards, i, cardinality) + + if protocol_version >= (2, 0): + source_type_pos = hton.unpack_int16( + frb_read(spec, 2)) + source_type = codecs_list[source_type_pos] + + res = ObjectCodec.new( + tid, names, flags, cards, codecs, t == CTYPE_INPUT_SHAPE + ) + + elif t == CTYPE_INPUT_SHAPE: els = hton.unpack_int16(frb_read(spec, 2)) codecs = cpython.PyTuple_New(els) names = cpython.PyTuple_New(els) @@ -220,15 +275,57 @@ cdef class CodecsRegistry: res = BASE_SCALAR_CODECS[tid] elif t == CTYPE_SCALAR: - pos = hton.unpack_int16(frb_read(spec, 2)) - codec = codecs_list[pos] - if type(codec) is not ScalarCodec: - raise RuntimeError( - f'a scalar codec expected for base scalar type, ' - f'got {type(codec).__name__}') - res = (codecs_list[pos]).derive(tid) + if protocol_version >= (2, 0): + str_len = hton.unpack_uint32(frb_read(spec, 4)) + name = cpythonx.PyUnicode_FromStringAndSize( + frb_read(spec, str_len), str_len) + schema_defined = frb_read(spec, 1)[0] + + ancestor_count = hton.unpack_int16(frb_read(spec, 2)) + ancestors = [] + for _ in range(ancestor_count): + ancestor_pos = hton.unpack_int16( + frb_read(spec, 2)) + ancestor_codec = codecs_list[ancestor_pos] + if type(ancestor_codec) is not ScalarCodec: + raise RuntimeError( + f'a scalar codec expected for base scalar type, ' + f'got {type(ancestor_codec).__name__}') + ancestors.append(ancestor_codec) + + if ancestor_count == 0: + if tid in self.base_codec_overrides: + res = self.base_codec_overrides[tid] + else: + res = BASE_SCALAR_CODECS[tid] + else: + fundamental_codec = ancestors[-1] + if type(fundamental_codec) is not ScalarCodec: + raise RuntimeError( + f'a scalar codec expected for base scalar type, ' + f'got {type(fundamental_codec).__name__}') + res = (fundamental_codec).derive(tid) + else: + fundamental_pos = hton.unpack_int16( + frb_read(spec, 2)) + fundamental_codec = codecs_list[fundamental_pos] + if type(fundamental_codec) is not ScalarCodec: + raise RuntimeError( + f'a scalar codec expected for base scalar type, ' + f'got {type(fundamental_codec).__name__}') + res = (fundamental_codec).derive(tid) elif t == CTYPE_TUPLE: + if protocol_version >= (2, 0): + str_len = hton.unpack_uint32(frb_read(spec, 4)) + name = cpythonx.PyUnicode_FromStringAndSize( + frb_read(spec, str_len), str_len) + schema_defined = frb_read(spec, 1)[0] + ancestor_count = hton.unpack_int16(frb_read(spec, 2)) + for _ in range(ancestor_count): + ancestor_pos = hton.unpack_int16( + frb_read(spec, 2)) + ancestor_codec = codecs_list[ancestor_pos] els = hton.unpack_int16(frb_read(spec, 2)) codecs = cpython.PyTuple_New(els) for i in range(els): @@ -241,6 +338,16 @@ cdef class CodecsRegistry: res = TupleCodec.new(tid, codecs) elif t == CTYPE_NAMEDTUPLE: + if protocol_version >= (2, 0): + str_len = hton.unpack_uint32(frb_read(spec, 4)) + name = cpythonx.PyUnicode_FromStringAndSize( + frb_read(spec, str_len), str_len) + schema_defined = frb_read(spec, 1)[0] + ancestor_count = hton.unpack_int16(frb_read(spec, 2)) + for _ in range(ancestor_count): + ancestor_pos = hton.unpack_int16( + frb_read(spec, 2)) + ancestor_codec = codecs_list[ancestor_pos] els = hton.unpack_int16(frb_read(spec, 2)) codecs = cpython.PyTuple_New(els) names = cpython.PyTuple_New(els) @@ -260,6 +367,16 @@ cdef class CodecsRegistry: res = NamedTupleCodec.new(tid, names, codecs) elif t == CTYPE_ENUM: + if protocol_version >= (2, 0): + str_len = hton.unpack_uint32(frb_read(spec, 4)) + name = cpythonx.PyUnicode_FromStringAndSize( + frb_read(spec, str_len), str_len) + schema_defined = frb_read(spec, 1)[0] + ancestor_count = hton.unpack_int16(frb_read(spec, 2)) + for _ in range(ancestor_count): + ancestor_pos = hton.unpack_int16( + frb_read(spec, 2)) + ancestor_codec = codecs_list[ancestor_pos] els = hton.unpack_int16(frb_read(spec, 2)) names = cpython.PyTuple_New(els) for i in range(els): @@ -273,6 +390,16 @@ cdef class CodecsRegistry: res = EnumCodec.new(tid, names) elif t == CTYPE_ARRAY: + if protocol_version >= (2, 0): + str_len = hton.unpack_uint32(frb_read(spec, 4)) + name = cpythonx.PyUnicode_FromStringAndSize( + frb_read(spec, str_len), str_len) + schema_defined = frb_read(spec, 1)[0] + ancestor_count = hton.unpack_int16(frb_read(spec, 2)) + for _ in range(ancestor_count): + ancestor_pos = hton.unpack_int16( + frb_read(spec, 2)) + ancestor_codec = codecs_list[ancestor_pos] pos = hton.unpack_int16(frb_read(spec, 2)) els = hton.unpack_int16(frb_read(spec, 2)) if els != 1: @@ -284,10 +411,30 @@ cdef class CodecsRegistry: res = ArrayCodec.new(tid, sub_codec, dim_len) elif t == CTYPE_RANGE: + if protocol_version >= (2, 0): + str_len = hton.unpack_uint32(frb_read(spec, 4)) + name = cpythonx.PyUnicode_FromStringAndSize( + frb_read(spec, str_len), str_len) + schema_defined = frb_read(spec, 1)[0] + ancestor_count = hton.unpack_int16(frb_read(spec, 2)) + for _ in range(ancestor_count): + ancestor_pos = hton.unpack_int16( + frb_read(spec, 2)) + ancestor_codec = codecs_list[ancestor_pos] pos = hton.unpack_int16(frb_read(spec, 2)) sub_codec = codecs_list[pos] res = RangeCodec.new(tid, sub_codec) + elif t == CTYPE_OBJECT and protocol_version >= (2, 0): + # Ignore + frb_read(spec, desc_len) + res = NULL_CODEC + + elif t == CTYPE_COMPOUND and protocol_version >= (2, 0): + # Ignore + frb_read(spec, desc_len) + res = NULL_CODEC + else: raise NotImplementedError( f'no codec implementation for EdgeDB data class {t}') @@ -318,6 +465,7 @@ cdef class CodecsRegistry: cdef BaseCodec build_codec(self, bytes spec, protocol_version): cdef: FRBuffer buf + FRBuffer elem_buf BaseCodec res list codecs_list @@ -328,7 +476,16 @@ cdef class CodecsRegistry: codecs_list = [] while frb_get_len(&buf): - res = self._build_codec(&buf, codecs_list, protocol_version) + if protocol_version >= (2, 0): + desc_len = hton.unpack_int32(frb_read(&buf, 4)) + frb_slice_from(&elem_buf, &buf, desc_len) + res = self._build_codec( + &elem_buf, codecs_list, protocol_version) + if frb_get_len(&elem_buf): + raise RuntimeError( + f'unexpected trailing data in type descriptor datum') + else: + res = self._build_codec(&buf, codecs_list, protocol_version) if res is None: # An annotation; ignore. continue diff --git a/edgedb/protocol/consts.pxi b/edgedb/protocol/consts.pxi index 6d93c07f..10fe4cf8 100644 --- a/edgedb/protocol/consts.pxi +++ b/edgedb/protocol/consts.pxi @@ -61,8 +61,8 @@ DEF TRANS_STATUS_IDLE = b'I' DEF TRANS_STATUS_INTRANS = b'T' DEF TRANS_STATUS_ERROR = b'E' -DEF PROTO_VER_MAJOR = 1 +DEF PROTO_VER_MAJOR = 2 DEF PROTO_VER_MINOR = 0 -DEF LEGACY_PROTO_VER_MAJOR = 0 -DEF LEGACY_PROTO_VER_MINOR_MIN = 13 +DEF MIN_PROTO_VER_MAJOR = 0 +DEF MIN_PROTO_VER_MINOR = 13 diff --git a/edgedb/protocol/protocol.pyx b/edgedb/protocol/protocol.pyx index b973b885..469d3cb4 100644 --- a/edgedb/protocol/protocol.pyx +++ b/edgedb/protocol/protocol.pyx @@ -148,7 +148,7 @@ cdef class SansIOProtocol: self.internal_reg = CodecsRegistry() self.server_settings = {} self.reset_status() - self.protocol_version = (PROTO_VER_MAJOR, 0) + self.protocol_version = (PROTO_VER_MAJOR, PROTO_VER_MINOR) self.state_type_id = NULL_CODEC_ID self.state_codec = None @@ -873,22 +873,19 @@ cdef class SansIOProtocol: minor = self.buffer.read_int16() # TODO: drop this branch when dropping protocol_v0 - if major == LEGACY_PROTO_VER_MAJOR: + if major == 0: self.is_legacy = True self.ignore_headers() self.buffer.finish_message() - if major != PROTO_VER_MAJOR and not ( - major == LEGACY_PROTO_VER_MAJOR and - minor >= LEGACY_PROTO_VER_MINOR_MIN - ): + if (major, minor) < (MIN_PROTO_VER_MAJOR, MIN_PROTO_VER_MINOR): raise errors.ClientConnectionError( f'the server requested an unsupported version of ' f'the protocol: {major}.{minor}' ) - - self.protocol_version = (major, minor) + else: + self.protocol_version = (major, minor) elif mtype == AUTH_REQUEST_MSG: # Authentication... From 8769db8406fb2ad68899b700842ca147b7c359e2 Mon Sep 17 00:00:00 2001 From: Fantix King Date: Fri, 7 Jul 2023 16:36:35 -0400 Subject: [PATCH 2/2] Apply type names from proto v2 --- edgedb/protocol/codecs/array.pyx | 2 +- edgedb/protocol/codecs/codecs.pyx | 28 +++++++++++++++++++++------ edgedb/protocol/codecs/namedtuple.pyx | 2 +- edgedb/protocol/codecs/range.pyx | 2 +- edgedb/protocol/codecs/tuple.pyx | 2 +- 5 files changed, 26 insertions(+), 10 deletions(-) diff --git a/edgedb/protocol/codecs/array.pyx b/edgedb/protocol/codecs/array.pyx index fb0f872b..ef64fadc 100644 --- a/edgedb/protocol/codecs/array.pyx +++ b/edgedb/protocol/codecs/array.pyx @@ -156,7 +156,7 @@ cdef class ArrayCodec(BaseArrayCodec): def make_type(self, describe_context): return describe.ArrayType( desc_id=uuid.UUID(bytes=self.tid), - name=None, + name=self.type_name, element_type=self.sub_codec.make_type(describe_context), ) diff --git a/edgedb/protocol/codecs/codecs.pyx b/edgedb/protocol/codecs/codecs.pyx index 70da3747..bb669350 100644 --- a/edgedb/protocol/codecs/codecs.pyx +++ b/edgedb/protocol/codecs/codecs.pyx @@ -280,7 +280,7 @@ cdef class CodecsRegistry: elif t == CTYPE_SCALAR: if protocol_version >= (2, 0): str_len = hton.unpack_uint32(frb_read(spec, 4)) - name = cpythonx.PyUnicode_FromStringAndSize( + type_name = cpythonx.PyUnicode_FromStringAndSize( frb_read(spec, str_len), str_len) schema_defined = frb_read(spec, 1)[0] @@ -308,6 +308,7 @@ cdef class CodecsRegistry: f'a scalar codec expected for base scalar type, ' f'got {type(fundamental_codec).__name__}') res = (fundamental_codec).derive(tid) + res.type_name = type_name else: fundamental_pos = hton.unpack_int16( frb_read(spec, 2)) @@ -321,7 +322,7 @@ cdef class CodecsRegistry: elif t == CTYPE_TUPLE: if protocol_version >= (2, 0): str_len = hton.unpack_uint32(frb_read(spec, 4)) - name = cpythonx.PyUnicode_FromStringAndSize( + type_name = cpythonx.PyUnicode_FromStringAndSize( frb_read(spec, str_len), str_len) schema_defined = frb_read(spec, 1)[0] ancestor_count = hton.unpack_int16(frb_read(spec, 2)) @@ -329,6 +330,8 @@ cdef class CodecsRegistry: ancestor_pos = hton.unpack_int16( frb_read(spec, 2)) ancestor_codec = codecs_list[ancestor_pos] + else: + type_name = None els = hton.unpack_int16(frb_read(spec, 2)) codecs = cpython.PyTuple_New(els) for i in range(els): @@ -339,11 +342,12 @@ cdef class CodecsRegistry: cpython.PyTuple_SetItem(codecs, i, sub_codec) res = TupleCodec.new(tid, codecs) + res.type_name = type_name elif t == CTYPE_NAMEDTUPLE: if protocol_version >= (2, 0): str_len = hton.unpack_uint32(frb_read(spec, 4)) - name = cpythonx.PyUnicode_FromStringAndSize( + type_name = cpythonx.PyUnicode_FromStringAndSize( frb_read(spec, str_len), str_len) schema_defined = frb_read(spec, 1)[0] ancestor_count = hton.unpack_int16(frb_read(spec, 2)) @@ -351,6 +355,8 @@ cdef class CodecsRegistry: ancestor_pos = hton.unpack_int16( frb_read(spec, 2)) ancestor_codec = codecs_list[ancestor_pos] + else: + type_name = None els = hton.unpack_int16(frb_read(spec, 2)) codecs = cpython.PyTuple_New(els) names = cpython.PyTuple_New(els) @@ -368,11 +374,12 @@ cdef class CodecsRegistry: cpython.PyTuple_SetItem(codecs, i, sub_codec) res = NamedTupleCodec.new(tid, names, codecs) + res.type_name = type_name elif t == CTYPE_ENUM: if protocol_version >= (2, 0): str_len = hton.unpack_uint32(frb_read(spec, 4)) - name = cpythonx.PyUnicode_FromStringAndSize( + type_name = cpythonx.PyUnicode_FromStringAndSize( frb_read(spec, str_len), str_len) schema_defined = frb_read(spec, 1)[0] ancestor_count = hton.unpack_int16(frb_read(spec, 2)) @@ -380,6 +387,8 @@ cdef class CodecsRegistry: ancestor_pos = hton.unpack_int16( frb_read(spec, 2)) ancestor_codec = codecs_list[ancestor_pos] + else: + type_name = None els = hton.unpack_int16(frb_read(spec, 2)) names = cpython.PyTuple_New(els) for i in range(els): @@ -391,11 +400,12 @@ cdef class CodecsRegistry: cpython.PyTuple_SetItem(names, i, name) res = EnumCodec.new(tid, names) + res.type_name = type_name elif t == CTYPE_ARRAY: if protocol_version >= (2, 0): str_len = hton.unpack_uint32(frb_read(spec, 4)) - name = cpythonx.PyUnicode_FromStringAndSize( + type_name = cpythonx.PyUnicode_FromStringAndSize( frb_read(spec, str_len), str_len) schema_defined = frb_read(spec, 1)[0] ancestor_count = hton.unpack_int16(frb_read(spec, 2)) @@ -403,6 +413,8 @@ cdef class CodecsRegistry: ancestor_pos = hton.unpack_int16( frb_read(spec, 2)) ancestor_codec = codecs_list[ancestor_pos] + else: + type_name = None pos = hton.unpack_int16(frb_read(spec, 2)) els = hton.unpack_int16(frb_read(spec, 2)) if els != 1: @@ -412,11 +424,12 @@ cdef class CodecsRegistry: dim_len = hton.unpack_int32(frb_read(spec, 4)) sub_codec = codecs_list[pos] res = ArrayCodec.new(tid, sub_codec, dim_len) + res.type_name = type_name elif t == CTYPE_RANGE: if protocol_version >= (2, 0): str_len = hton.unpack_uint32(frb_read(spec, 4)) - name = cpythonx.PyUnicode_FromStringAndSize( + type_name = cpythonx.PyUnicode_FromStringAndSize( frb_read(spec, str_len), str_len) schema_defined = frb_read(spec, 1)[0] ancestor_count = hton.unpack_int16(frb_read(spec, 2)) @@ -424,9 +437,12 @@ cdef class CodecsRegistry: ancestor_pos = hton.unpack_int16( frb_read(spec, 2)) ancestor_codec = codecs_list[ancestor_pos] + else: + type_name = None pos = hton.unpack_int16(frb_read(spec, 2)) sub_codec = codecs_list[pos] res = RangeCodec.new(tid, sub_codec) + res.type_name = type_name elif t == CTYPE_OBJECT and protocol_version >= (2, 0): # Ignore diff --git a/edgedb/protocol/codecs/namedtuple.pyx b/edgedb/protocol/codecs/namedtuple.pyx index 930ee0ee..6514ef36 100644 --- a/edgedb/protocol/codecs/namedtuple.pyx +++ b/edgedb/protocol/codecs/namedtuple.pyx @@ -78,7 +78,7 @@ cdef class NamedTupleCodec(BaseNamedRecordCodec): def make_type(self, describe_context): return describe.NamedTupleType( desc_id=uuid.UUID(bytes=self.tid), - name=None, + name=self.type_name, element_types={ field: codec.make_type(describe_context) for field, codec in zip( diff --git a/edgedb/protocol/codecs/range.pyx b/edgedb/protocol/codecs/range.pyx index 9555d969..3d608fd0 100644 --- a/edgedb/protocol/codecs/range.pyx +++ b/edgedb/protocol/codecs/range.pyx @@ -143,6 +143,6 @@ cdef class RangeCodec(BaseCodec): def make_type(self, describe_context): return describe.RangeType( desc_id=uuid.UUID(bytes=self.tid), - name=None, + name=self.type_name, value_type=self.sub_codec.make_type(describe_context), ) diff --git a/edgedb/protocol/codecs/tuple.pyx b/edgedb/protocol/codecs/tuple.pyx index 68ed0352..d29415d6 100644 --- a/edgedb/protocol/codecs/tuple.pyx +++ b/edgedb/protocol/codecs/tuple.pyx @@ -81,7 +81,7 @@ cdef class TupleCodec(BaseRecordCodec): def make_type(self, describe_context): return describe.TupleType( desc_id=uuid.UUID(bytes=self.tid), - name=None, + name=self.type_name, element_types=tuple( codec.make_type(describe_context) for codec in self.fields_codecs