diff --git a/edgedb/__init__.py b/edgedb/__init__.py index bf5659210..dc8601c36 100644 --- a/edgedb/__init__.py +++ b/edgedb/__init__.py @@ -26,6 +26,8 @@ ) from edgedb.datatypes.datatypes import Set, Object, Array, Link, LinkSet from edgedb.datatypes.range import Range +from edgedb.pgproto.pgproto import CodecContext +from edgedb.protocol.protocol import get_default_codec_context from .abstract import ( Executor, AsyncIOExecutor, ReadOnlyExecutor, AsyncIOReadOnlyExecutor, @@ -51,6 +53,7 @@ "AsyncIOReadOnlyExecutor", "Cardinality", "Client", + "CodecContext", "ConfigMemory", "DateDuration", "EdgeDBError", @@ -58,6 +61,7 @@ "ElementKind", "EnumValue", "Executor", + "get_default_codec_context", "IsolationLevel", "Link", "LinkSet", diff --git a/edgedb/abstract.py b/edgedb/abstract.py index 158a269a1..252f63245 100644 --- a/edgedb/abstract.py +++ b/edgedb/abstract.py @@ -25,6 +25,7 @@ from . import describe from . import enums from . import options +from .pgproto import pgproto from .protocol import protocol __all__ = ( @@ -64,12 +65,14 @@ class QueryContext(typing.NamedTuple): query_options: QueryOptions retry_options: typing.Optional[options.RetryOptions] state: typing.Optional[options.State] + codec_ctx: pgproto.CodecContext class ExecuteContext(typing.NamedTuple): query: QueryWithArgs cache: QueryCache state: typing.Optional[options.State] + codec_ctx: pgproto.CodecContext @dataclasses.dataclass @@ -129,9 +132,14 @@ def _get_query_cache(self) -> QueryCache: def _get_retry_options(self) -> typing.Optional[options.RetryOptions]: return None + @abc.abstractmethod def _get_state(self) -> options.State: ... + @abc.abstractmethod + def _get_codec_ctx(self) -> pgproto.CodecContext: + ... + class ReadOnlyExecutor(BaseReadOnlyExecutor): """Subclasses can execute *at least* read-only queries""" @@ -149,6 +157,7 @@ def query(self, query: str, *args, **kwargs) -> list: query_options=_query_opts, retry_options=self._get_retry_options(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) def query_single( @@ -160,6 +169,7 @@ def query_single( query_options=_query_single_opts, retry_options=self._get_retry_options(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) def query_required_single(self, query: str, *args, **kwargs) -> typing.Any: @@ -169,6 +179,7 @@ def query_required_single(self, query: str, *args, **kwargs) -> typing.Any: query_options=_query_required_single_opts, retry_options=self._get_retry_options(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) def query_json(self, query: str, *args, **kwargs) -> str: @@ -178,6 +189,7 @@ def query_json(self, query: str, *args, **kwargs) -> str: query_options=_query_json_opts, retry_options=self._get_retry_options(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) def query_single_json(self, query: str, *args, **kwargs) -> str: @@ -187,6 +199,7 @@ def query_single_json(self, query: str, *args, **kwargs) -> str: query_options=_query_single_json_opts, retry_options=self._get_retry_options(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) def query_required_single_json(self, query: str, *args, **kwargs) -> str: @@ -196,6 +209,7 @@ def query_required_single_json(self, query: str, *args, **kwargs) -> str: query_options=_query_required_single_json_opts, retry_options=self._get_retry_options(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) @abc.abstractmethod @@ -207,6 +221,7 @@ def execute(self, commands: str, *args, **kwargs) -> None: query=QueryWithArgs(commands, args, kwargs), cache=self._get_query_cache(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) @@ -232,6 +247,7 @@ async def query(self, query: str, *args, **kwargs) -> list: query_options=_query_opts, retry_options=self._get_retry_options(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) async def query_single(self, query: str, *args, **kwargs) -> typing.Any: @@ -241,6 +257,7 @@ async def query_single(self, query: str, *args, **kwargs) -> typing.Any: query_options=_query_single_opts, retry_options=self._get_retry_options(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) async def query_required_single( @@ -255,6 +272,7 @@ async def query_required_single( query_options=_query_required_single_opts, retry_options=self._get_retry_options(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) async def query_json(self, query: str, *args, **kwargs) -> str: @@ -264,6 +282,7 @@ async def query_json(self, query: str, *args, **kwargs) -> str: query_options=_query_json_opts, retry_options=self._get_retry_options(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) async def query_single_json(self, query: str, *args, **kwargs) -> str: @@ -273,6 +292,7 @@ async def query_single_json(self, query: str, *args, **kwargs) -> str: query_options=_query_single_json_opts, retry_options=self._get_retry_options(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) async def query_required_single_json( @@ -287,6 +307,7 @@ async def query_required_single_json( query_options=_query_required_single_json_opts, retry_options=self._get_retry_options(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) @abc.abstractmethod @@ -298,6 +319,7 @@ async def execute(self, commands: str, *args, **kwargs) -> None: query=QueryWithArgs(commands, args, kwargs), cache=self._get_query_cache(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) diff --git a/edgedb/base_client.py b/edgedb/base_client.py index 94d859721..d41a9ebb6 100644 --- a/edgedb/base_client.py +++ b/edgedb/base_client.py @@ -27,6 +27,7 @@ from . import enums from . import errors from . import options as _options +from .pgproto import pgproto from .protocol import protocol @@ -186,6 +187,7 @@ async def privileged_execute( qc=execute_context.cache.query_cache, output_format=protocol.OutputFormat.NONE, allow_capabilities=enums.Capability.ALL, + codec_ctx=execute_context.codec_ctx, ) def is_in_transaction(self) -> bool: @@ -214,6 +216,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, + codec_ctx=query_context.codec_ctx, ) if self._protocol.is_legacy: args["allow_capabilities"] = enums.Capability.LEGACY_EXECUTE @@ -284,6 +287,7 @@ async def _execute(self, execute_context: abstract.ExecuteContext) -> None: execute_context.state.as_dict() if execute_context.state else None ), + codec_ctx=execute_context.codec_ctx, ) async def describe( @@ -719,6 +723,9 @@ def _get_retry_options(self) -> typing.Optional[_options.RetryOptions]: def _get_state(self) -> _options.State: return self._options.state + def _get_codec_ctx(self) -> pgproto.CodecContext: + return self._options.codec_context + @property def max_concurrency(self) -> int: """Max number of connections in the pool.""" diff --git a/edgedb/codegen/cli.py b/edgedb/codegen/cli.py index 1a229bae4..3849d41b5 100644 --- a/edgedb/codegen/cli.py +++ b/edgedb/codegen/cli.py @@ -60,6 +60,12 @@ help="Add a mixin to generated dataclasses " "to skip Pydantic validation (default is to add the mixin).", ) + parser.add_argument( + "--handle-json", + action=argparse.BooleanOptionalAction, + default=False, + help="Choose to handle JSON in query arguments and results." + ) else: parser.add_argument( "--skip-pydantic-validation", @@ -74,6 +80,19 @@ help="Add a mixin to generated dataclasses " "to skip Pydantic validation (default is to add the mixin).", ) + parser.add_argument( + "--handle-json", + action="store_true", + default=False, + ) + parser.add_argument( + "--no-handle-json", + dest="handle_json", + action="store_false", + default=True, + help="Choose to handle JSON in query arguments and results. " + "(default: False)" + ) def main(): diff --git a/edgedb/codegen/generator.py b/edgedb/codegen/generator.py index 65720066c..2e4b9a1e1 100644 --- a/edgedb/codegen/generator.py +++ b/edgedb/codegen/generator.py @@ -123,6 +123,7 @@ def __init__(self, args: argparse.Namespace): self._default_module = "default" self._targets = args.target self._skip_pydantic_validation = args.skip_pydantic_validation + self._handle_json = args.handle_json self._async = False try: self._project_dir = pathlib.Path( @@ -339,6 +340,16 @@ def _generate( method = "query_single" rt = "return " + if self._handle_json: + print(f"{INDENT}client = client.with_codec_context(", file=buf) + print( + f"{INDENT}{INDENT}" + f"edgedb.get_default_codec_context(handle_json=True),", + file=buf, + ) + print(f"{INDENT}{INDENT}only_replace_default=True,", file=buf) + print(f"{INDENT})", file=buf) + if self._async: print(f"{INDENT}{rt}await client.{method}(", file=buf) else: @@ -369,7 +380,17 @@ def _generate_code( if type_.desc_id in self._cache: return self._cache[type_.desc_id] - if isinstance(type_, describe.BaseScalarType): + if ( + isinstance(type_, describe.BaseScalarType) + and type_.name == "std::json" + ): + if self._handle_json: + self._imports.add("typing") + rv = "typing.Any" + else: + rv = "str" + + elif isinstance(type_, describe.BaseScalarType): if type_.name in TYPE_IMPORTS: self._imports.add(TYPE_IMPORTS[type_.name]) rv = TYPE_MAPPING[type_.name] diff --git a/edgedb/options.py b/edgedb/options.py index 714168ab3..57dc2346b 100644 --- a/edgedb/options.py +++ b/edgedb/options.py @@ -5,6 +5,8 @@ from collections import namedtuple from . import errors +from .pgproto import pgproto +from .protocol import protocol _RetryRule = namedtuple("_RetryRule", ["attempts", "backoff"]) @@ -365,21 +367,41 @@ def without_globals(self, *global_names): ) return result + def with_codec_context( + self, codec_context: pgproto.CodecContext, only_replace_default=False + ): + if self._options.codec_context is codec_context: + return self + if only_replace_default: + default = protocol.get_default_codec_context() + if self._options.codec_context is not default: + return self + result = self._shallow_clone() + result._options = self._options.with_codec_context(codec_context) + return result + class _Options: """Internal class for storing connection options""" - __slots__ = ['_retry_options', '_transaction_options', '_state'] + __slots__ = [ + '_retry_options', + '_transaction_options', + '_state', + '_codec_context', + ] def __init__( self, retry_options: RetryOptions, transaction_options: TransactionOptions, state: State, + codec_context: pgproto.CodecContext, ): self._retry_options = retry_options self._transaction_options = transaction_options self._state = state + self._codec_context = codec_context @property def retry_options(self): @@ -393,11 +415,16 @@ def transaction_options(self): def state(self): return self._state + @property + def codec_context(self): + return self._codec_context + def with_retry_options(self, options: RetryOptions): return _Options( options, self._transaction_options, self._state, + self._codec_context, ) def with_transaction_options(self, options: TransactionOptions): @@ -405,6 +432,7 @@ def with_transaction_options(self, options: TransactionOptions): self._retry_options, options, self._state, + self._codec_context, ) def with_state(self, state: State): @@ -412,6 +440,15 @@ def with_state(self, state: State): self._retry_options, self._transaction_options, state, + self._codec_context, + ) + + def with_codec_context(self, codec_context: pgproto.CodecContext): + return _Options( + self._retry_options, + self._transaction_options, + self._state, + codec_context, ) @classmethod @@ -420,4 +457,5 @@ def defaults(cls): RetryOptions.defaults(), TransactionOptions.defaults(), State.defaults(), + protocol.get_default_codec_context(), ) diff --git a/edgedb/pgproto b/edgedb/pgproto index 1720f8af6..c4fcf43e4 160000 --- a/edgedb/pgproto +++ b/edgedb/pgproto @@ -1 +1 @@ -Subproject commit 1720f8af63725d79454884cfa787202a50eb5430 +Subproject commit c4fcf43e44a15d23fe98666600fa8c90b17711a1 diff --git a/edgedb/protocol/codecs/array.pxd b/edgedb/protocol/codecs/array.pxd index 44de315be..1c43864b0 100644 --- a/edgedb/protocol/codecs/array.pxd +++ b/edgedb/protocol/codecs/array.pxd @@ -23,7 +23,7 @@ cdef class BaseArrayCodec(BaseCodec): BaseCodec sub_codec int32_t cardinality - cdef _decode_array(self, FRBuffer *buf) + cdef _decode_array(self, FRBuffer *buf, pgproto.CodecContext ctx) @cython.final diff --git a/edgedb/protocol/codecs/array.pyx b/edgedb/protocol/codecs/array.pyx index 2f709531d..6de634e59 100644 --- a/edgedb/protocol/codecs/array.pyx +++ b/edgedb/protocol/codecs/array.pyx @@ -30,7 +30,7 @@ cdef class BaseArrayCodec(BaseCodec): self.sub_codec = None self.cardinality = -1 - cdef encode(self, WriteBuffer buf, object obj): + cdef encode(self, WriteBuffer buf, object obj, pgproto.CodecContext ctx): cdef: WriteBuffer elem_data int32_t ndims = 1 @@ -63,7 +63,7 @@ cdef class BaseArrayCodec(BaseCodec): elem_data.write_int32(-1) else: try: - self.sub_codec.encode(elem_data, item) + self.sub_codec.encode(elem_data, item, ctx) except TypeError as e: raise ValueError( 'invalid array element: {}'.format( @@ -79,10 +79,10 @@ cdef class BaseArrayCodec(BaseCodec): buf.write_buffer(elem_data) - cdef decode(self, FRBuffer *buf): - return self._decode_array(buf) + cdef decode(self, FRBuffer *buf, pgproto.CodecContext ctx): + return self._decode_array(buf, ctx) - cdef inline _decode_array(self, FRBuffer *buf): + cdef inline _decode_array(self, FRBuffer *buf, pgproto.CodecContext ctx): cdef: Py_ssize_t elem_count int32_t ndims = hton.unpack_int32(frb_read(buf, 4)) @@ -118,7 +118,7 @@ cdef class BaseArrayCodec(BaseCodec): elem = None else: frb_slice_from(&elem_buf, buf, elem_len) - elem = self.sub_codec.decode(&elem_buf) + elem = self.sub_codec.decode(&elem_buf, ctx) if frb_get_len(&elem_buf): raise RuntimeError( f'unexpected trailing data in buffer after ' diff --git a/edgedb/protocol/codecs/base.pxd b/edgedb/protocol/codecs/base.pxd index 962753de7..ce3ea1101 100644 --- a/edgedb/protocol/codecs/base.pxd +++ b/edgedb/protocol/codecs/base.pxd @@ -27,8 +27,8 @@ cdef class BaseCodec: cdef inline bytes get_tid(self): return self.tid - cdef encode(self, WriteBuffer buf, object obj) - cdef decode(self, FRBuffer *buf) + cdef encode(self, WriteBuffer buf, object obj, pgproto.CodecContext ctx) + cdef decode(self, FRBuffer *buf, pgproto.CodecContext ctx) cdef dump(self, int level=?) @@ -77,3 +77,12 @@ cdef class EdegDBCodecContext(pgproto.CodecContext): cdef: object _codec + + +@cython.final +cdef class EdegDBJSONCodecContext(pgproto.CodecContext): + + cdef: + object _codec + object _json_decoder + object _json_encoder diff --git a/edgedb/protocol/codecs/base.pyx b/edgedb/protocol/codecs/base.pyx index 3bd52bb0b..6b7f9a037 100644 --- a/edgedb/protocol/codecs/base.pyx +++ b/edgedb/protocol/codecs/base.pyx @@ -18,6 +18,7 @@ import codecs +import json cdef uint64_t RECORD_ENCODER_CHECKED = 1 << 0 @@ -40,10 +41,10 @@ cdef class BaseCodec: self.tid = None self.name = None - cdef encode(self, WriteBuffer buf, object obj): + cdef encode(self, WriteBuffer buf, object obj, pgproto.CodecContext ctx): raise NotImplementedError - cdef decode(self, FRBuffer *buf): + cdef decode(self, FRBuffer *buf, pgproto.CodecContext ctx): raise NotImplementedError cdef dump(self, int level = 0): @@ -60,11 +61,11 @@ cdef class CodecPythonOverride(BaseCodec): self.encoder = None self.decoder = None - cdef encode(self, WriteBuffer buf, object obj): - self.codec.encode(buf, self.encoder(obj)) + cdef encode(self, WriteBuffer buf, object obj, pgproto.CodecContext ctx): + self.codec.encode(buf, self.encoder(obj), ctx) - cdef decode(self, FRBuffer *buf): - return self.decoder(self.codec.decode(buf)) + cdef decode(self, FRBuffer *buf, pgproto.CodecContext ctx): + return self.decoder(self.codec.decode(buf, ctx)) cdef dump(self, int level = 0): return f'{level * " "}{self.name}' @@ -97,7 +98,7 @@ cdef class EmptyTupleCodec(BaseCodec): self.name = 'no-input' self.empty_tup = None - cdef encode(self, WriteBuffer buf, object obj): + cdef encode(self, WriteBuffer buf, object obj, pgproto.CodecContext ctx): if type(obj) is not tuple: raise RuntimeError( f'cannot encode empty Tuple: expected a tuple, ' @@ -108,7 +109,7 @@ cdef class EmptyTupleCodec(BaseCodec): f'got {len(obj)}') buf.write_bytes(EMPTY_RECORD_DATA) - cdef decode(self, FRBuffer *buf): + cdef decode(self, FRBuffer *buf, pgproto.CodecContext ctx): elem_count = hton.unpack_int32(frb_read(buf, 4)) if elem_count != 0: raise RuntimeError( @@ -159,7 +160,7 @@ cdef class BaseRecordCodec(BaseCodec): raise TypeError( 'argument tuples do not support objects') - cdef encode(self, WriteBuffer buf, object obj): + cdef encode(self, WriteBuffer buf, object obj, pgproto.CodecContext ctx): cdef: WriteBuffer elem_data Py_ssize_t objlen @@ -195,7 +196,7 @@ cdef class BaseRecordCodec(BaseCodec): else: sub_codec = (self.fields_codecs[i]) try: - sub_codec.encode(elem_data, item) + sub_codec.encode(elem_data, item, ctx) except (TypeError, ValueError) as e: value_repr = repr(item) if len(value_repr) > 40: @@ -238,5 +239,47 @@ cdef class EdegDBCodecContext(pgproto.CodecContext): cdef is_encoding_utf8(self): return True + cdef is_decoding_json(self): + return False + + cdef is_encoding_json(self): + return False + + +@cython.final +cdef class EdegDBJSONCodecContext(pgproto.CodecContext): + + def __cinit__(self): + self._codec = codecs.lookup('utf-8') + self._json_decoder = json.JSONDecoder() + self._json_encoder = json.JSONEncoder() + + cpdef get_text_codec(self): + return self._codec + + cdef is_encoding_utf8(self): + return True + + cpdef get_json_decoder(self): + return self._json_decoder + + cdef is_decoding_json(self): + return True + + cpdef get_json_encoder(self): + return self._json_encoder + + cdef is_encoding_json(self): + return True + cdef EdegDBCodecContext DEFAULT_CODEC_CONTEXT = EdegDBCodecContext() +cdef EdegDBJSONCodecContext DEFAULT_JSON_CODEC_CONTEXT = \ + EdegDBJSONCodecContext() + + +def get_default_codec_context(handle_json=False): + if handle_json: + return DEFAULT_JSON_CODEC_CONTEXT + else: + return DEFAULT_CODEC_CONTEXT diff --git a/edgedb/protocol/codecs/enum.pyx b/edgedb/protocol/codecs/enum.pyx index 78a30f6b0..cbd77841e 100644 --- a/edgedb/protocol/codecs/enum.pyx +++ b/edgedb/protocol/codecs/enum.pyx @@ -22,15 +22,15 @@ import enum @cython.final cdef class EnumCodec(BaseCodec): - cdef encode(self, WriteBuffer buf, object obj): + cdef encode(self, WriteBuffer buf, object obj, pgproto.CodecContext ctx): if not isinstance(obj, (self.cls, str)): raise TypeError( f'a str or edgedb.EnumValue(__tid__={self.cls.__tid__}) is ' f'expected as a valid enum argument, got {type(obj).__name__}') - pgproto.text_encode(DEFAULT_CODEC_CONTEXT, buf, str(obj)) + pgproto.text_encode(ctx, buf, str(obj)) - cdef decode(self, FRBuffer *buf): - label = pgproto.text_decode(DEFAULT_CODEC_CONTEXT, buf) + cdef decode(self, FRBuffer *buf, pgproto.CodecContext ctx): + label = pgproto.text_decode(ctx, buf) return self.cls(label) @staticmethod diff --git a/edgedb/protocol/codecs/namedtuple.pyx b/edgedb/protocol/codecs/namedtuple.pyx index 930ee0ee5..d61d77af3 100644 --- a/edgedb/protocol/codecs/namedtuple.pyx +++ b/edgedb/protocol/codecs/namedtuple.pyx @@ -20,7 +20,7 @@ @cython.final cdef class NamedTupleCodec(BaseNamedRecordCodec): - cdef decode(self, FRBuffer *buf): + cdef decode(self, FRBuffer *buf, pgproto.CodecContext ctx): cdef: object result Py_ssize_t elem_count @@ -48,7 +48,7 @@ cdef class NamedTupleCodec(BaseNamedRecordCodec): else: elem_codec = fields_codecs[i] elem = elem_codec.decode( - frb_slice_from(&elem_buf, buf, elem_len)) + frb_slice_from(&elem_buf, buf, elem_len), ctx) if frb_get_len(&elem_buf): raise RuntimeError( f'unexpected trailing data in buffer after named ' diff --git a/edgedb/protocol/codecs/object.pxd b/edgedb/protocol/codecs/object.pxd index 18b782890..210d3a6c7 100644 --- a/edgedb/protocol/codecs/object.pxd +++ b/edgedb/protocol/codecs/object.pxd @@ -23,7 +23,7 @@ cdef class ObjectCodec(BaseNamedRecordCodec): bint is_sparse object cached_dataclass_fields - cdef encode_args(self, WriteBuffer buf, dict obj) + cdef encode_args(self, WriteBuffer buf, dict obj, pgproto.CodecContext ctx) @staticmethod cdef BaseCodec new(bytes tid, tuple names, tuple flags, diff --git a/edgedb/protocol/codecs/object.pyx b/edgedb/protocol/codecs/object.pyx index 8baf63b05..554c4e10c 100644 --- a/edgedb/protocol/codecs/object.pyx +++ b/edgedb/protocol/codecs/object.pyx @@ -31,7 +31,7 @@ cdef dict CARDS_MAP = { @cython.final cdef class ObjectCodec(BaseNamedRecordCodec): - cdef encode(self, WriteBuffer buf, object obj): + cdef encode(self, WriteBuffer buf, object obj, pgproto.CodecContext ctx): cdef: WriteBuffer elem_data Py_ssize_t objlen = 0 @@ -53,7 +53,7 @@ cdef class ObjectCodec(BaseNamedRecordCodec): if arg is not None: sub_codec = (self.fields_codecs[i]) try: - sub_codec.encode(elem_data, arg) + sub_codec.encode(elem_data, arg, ctx) except (TypeError, ValueError) as e: value_repr = repr(arg) if len(value_repr) > 40: @@ -68,7 +68,9 @@ cdef class ObjectCodec(BaseNamedRecordCodec): buf.write_int32(objlen) buf.write_buffer(elem_data) - cdef encode_args(self, WriteBuffer buf, dict obj): + cdef encode_args( + self, WriteBuffer buf, dict obj, pgproto.CodecContext ctx + ): cdef: WriteBuffer elem_data Py_ssize_t objlen @@ -106,7 +108,7 @@ cdef class ObjectCodec(BaseNamedRecordCodec): else: sub_codec = (self.fields_codecs[i]) try: - sub_codec.encode(elem_data, arg) + sub_codec.encode(elem_data, arg, ctx) except (TypeError, ValueError) as e: value_repr = repr(arg) if len(value_repr) > 40: @@ -149,7 +151,7 @@ cdef class ObjectCodec(BaseNamedRecordCodec): return errors.QueryArgumentError(error_message) - cdef decode(self, FRBuffer *buf): + cdef decode(self, FRBuffer *buf, pgproto.CodecContext ctx): cdef: object result Py_ssize_t elem_count @@ -181,7 +183,7 @@ cdef class ObjectCodec(BaseNamedRecordCodec): else: elem_codec = fields_codecs[i] elem = elem_codec.decode( - frb_slice_from(&elem_buf, buf, elem_len)) + frb_slice_from(&elem_buf, buf, elem_len), ctx) if frb_get_len(&elem_buf): raise RuntimeError( f'unexpected trailing data in buffer after ' diff --git a/edgedb/protocol/codecs/range.pyx b/edgedb/protocol/codecs/range.pyx index 9555d969e..c8b7e31e7 100644 --- a/edgedb/protocol/codecs/range.pyx +++ b/edgedb/protocol/codecs/range.pyx @@ -46,7 +46,7 @@ cdef class RangeCodec(BaseCodec): return codec - cdef encode(self, WriteBuffer buf, object obj): + cdef encode(self, WriteBuffer buf, object obj, pgproto.CodecContext ctx): cdef: uint8_t flags = 0 WriteBuffer sub_data @@ -78,14 +78,14 @@ cdef class RangeCodec(BaseCodec): sub_data = WriteBuffer.new() if lower is not None: try: - self.sub_codec.encode(sub_data, lower) + self.sub_codec.encode(sub_data, lower, ctx) except TypeError as e: raise ValueError( 'invalid range lower bound: {}'.format( e.args[0])) from None if upper is not None: try: - self.sub_codec.encode(sub_data, upper) + self.sub_codec.encode(sub_data, upper, ctx) except TypeError as e: raise ValueError( 'invalid range upper bound: {}'.format( @@ -95,7 +95,7 @@ cdef class RangeCodec(BaseCodec): buf.write_byte(flags) buf.write_buffer(sub_data) - cdef decode(self, FRBuffer *buf): + cdef decode(self, FRBuffer *buf, pgproto.CodecContext ctx): cdef: uint8_t flags = frb_read(buf, 1)[0] bint empty = (flags & RANGE_EMPTY) != 0 @@ -113,7 +113,7 @@ cdef class RangeCodec(BaseCodec): sub_len = hton.unpack_int32(frb_read(buf, 4)) if sub_len != -1: frb_slice_from(&sub_buf, buf, sub_len) - lower = sub_codec.decode(&sub_buf) + lower = sub_codec.decode(&sub_buf, ctx) if frb_get_len(&sub_buf): raise RuntimeError( f'unexpected trailing data in buffer after ' @@ -123,7 +123,7 @@ cdef class RangeCodec(BaseCodec): sub_len = hton.unpack_int32(frb_read(buf, 4)) if sub_len != -1: frb_slice_from(&sub_buf, buf, sub_len) - upper = sub_codec.decode(&sub_buf) + upper = sub_codec.decode(&sub_buf, ctx) if frb_get_len(&sub_buf): raise RuntimeError( f'unexpected trailing data in buffer after ' diff --git a/edgedb/protocol/codecs/scalar.pyx b/edgedb/protocol/codecs/scalar.pyx index 9dcebbbfc..3ba7cc2b0 100644 --- a/edgedb/protocol/codecs/scalar.pyx +++ b/edgedb/protocol/codecs/scalar.pyx @@ -24,11 +24,11 @@ cdef class ScalarCodec(BaseCodec): self.c_encoder = NULL self.c_decoder = NULL - cdef encode(self, WriteBuffer buf, object obj): - self.c_encoder(DEFAULT_CODEC_CONTEXT, buf, obj) + cdef encode(self, WriteBuffer buf, object obj, pgproto.CodecContext ctx): + self.c_encoder(ctx, buf, obj) - cdef decode(self, FRBuffer *buf): - return self.c_decoder(DEFAULT_CODEC_CONTEXT, buf) + cdef decode(self, FRBuffer *buf, pgproto.CodecContext ctx): + return self.c_decoder(ctx, buf) cdef derive(self, bytes tid): cdef ScalarCodec rv diff --git a/edgedb/protocol/codecs/set.pxd b/edgedb/protocol/codecs/set.pxd index 9a257d8a4..4da86afb6 100644 --- a/edgedb/protocol/codecs/set.pxd +++ b/edgedb/protocol/codecs/set.pxd @@ -20,7 +20,9 @@ @cython.final cdef class SetCodec(BaseArrayCodec): - cdef inline _decode_array_set(self, FRBuffer *buf) + cdef inline _decode_array_set( + self, FRBuffer *buf, pgproto.CodecContext ctx + ) @staticmethod cdef BaseCodec new(bytes tid, BaseCodec sub_codec) diff --git a/edgedb/protocol/codecs/set.pyx b/edgedb/protocol/codecs/set.pyx index 03b0e326a..a30975acc 100644 --- a/edgedb/protocol/codecs/set.pyx +++ b/edgedb/protocol/codecs/set.pyx @@ -33,16 +33,18 @@ cdef class SetCodec(BaseArrayCodec): return codec - cdef decode(self, FRBuffer *buf): + cdef decode(self, FRBuffer *buf, pgproto.CodecContext ctx): if type(self.sub_codec) is ArrayCodec: # This is a set of arrays encoded as an array # of single-element records. - return self._decode_array_set(buf) + return self._decode_array_set(buf, ctx) else: # Set of non-arrays. - return self._decode_array(buf) + return self._decode_array(buf, ctx) - cdef inline _decode_array_set(self, FRBuffer *buf): + cdef inline _decode_array_set( + self, FRBuffer *buf, pgproto.CodecContext ctx + ): cdef: object result object elem @@ -85,7 +87,7 @@ cdef class SetCodec(BaseArrayCodec): 'unexpected NULL value in array set element ') frb_slice_from(&elem_buf, buf, elem_len) - elem = sub_codec.decode(&elem_buf) + elem = sub_codec.decode(&elem_buf, ctx) if frb_get_len(&elem_buf): raise RuntimeError( f'unexpected trailing data in buffer after ' diff --git a/edgedb/protocol/codecs/tuple.pyx b/edgedb/protocol/codecs/tuple.pyx index 68ed0352d..0b260e2a2 100644 --- a/edgedb/protocol/codecs/tuple.pyx +++ b/edgedb/protocol/codecs/tuple.pyx @@ -20,7 +20,7 @@ @cython.final cdef class TupleCodec(BaseRecordCodec): - cdef decode(self, FRBuffer *buf): + cdef decode(self, FRBuffer *buf, pgproto.CodecContext ctx): cdef: object result Py_ssize_t elem_count @@ -48,7 +48,7 @@ cdef class TupleCodec(BaseRecordCodec): else: elem_codec = fields_codecs[i] elem = elem_codec.decode( - frb_slice_from(&elem_buf, buf, elem_len)) + frb_slice_from(&elem_buf, buf, elem_len), ctx) if frb_get_len(&elem_buf): raise RuntimeError( f'unexpected trailing data in buffer after ' diff --git a/edgedb/protocol/protocol.pxd b/edgedb/protocol/protocol.pxd index 03f049647..e301be7af 100644 --- a/edgedb/protocol/protocol.pxd +++ b/edgedb/protocol/protocol.pxd @@ -37,7 +37,9 @@ include "./lru.pxd" include "./codecs/codecs.pxd" -ctypedef object (*decode_row_method)(BaseCodec, FRBuffer *buf) +ctypedef object (*decode_row_method)( + BaseCodec, FRBuffer *buf, pgproto.CodecContext ctx +) cpdef enum OutputFormat: @@ -107,10 +109,19 @@ cdef class SansIOProtocol: BaseCodec state_codec object state_cache - cdef encode_args(self, BaseCodec in_dc, WriteBuffer buf, args, kwargs) + cdef encode_args( + self, + BaseCodec in_dc, + WriteBuffer buf, + args, + kwargs, + pgproto.CodecContext codec_ctx + ) cdef encode_state(self, state) - cdef parse_data_messages(self, BaseCodec out_dc, result) + cdef parse_data_messages( + self, BaseCodec out_dc, result, pgproto.CodecContext codec_ctx + ) cdef parse_sync_message(self) cdef parse_command_complete_message(self) cdef parse_describe_type_message(self, CodecsRegistry reg) diff --git a/edgedb/protocol/protocol.pyx b/edgedb/protocol/protocol.pyx index 6bc5cd75f..c27c39dc9 100644 --- a/edgedb/protocol/protocol.pyx +++ b/edgedb/protocol/protocol.pyx @@ -347,6 +347,7 @@ cdef class SansIOProtocol: in_dc: BaseCodec, out_dc: BaseCodec, state: typing.Optional[dict] = None, + codec_ctx: pgproto.CodecContext, ): cdef: WriteBuffer packet @@ -375,7 +376,7 @@ cdef class SansIOProtocol: buf.write_bytes(in_dc.get_tid()) buf.write_bytes(out_dc.get_tid()) - self.encode_args(in_dc, buf, args, kwargs) + self.encode_args(in_dc, buf, args, kwargs, codec_ctx) buf.end_message() @@ -413,7 +414,7 @@ cdef class SansIOProtocol: elif mtype == DATA_MSG: if exc is None: try: - self.parse_data_messages(out_dc, result) + self.parse_data_messages(out_dc, result, codec_ctx) except Exception as ex: # An error during data decoding. We need to # handle this as gracefully as possible: @@ -439,7 +440,9 @@ cdef class SansIOProtocol: if not isinstance(in_dc, NullCodec): buf = WriteBuffer.new() try: - self.encode_args(in_dc, buf, args, kwargs) + self.encode_args( + in_dc, buf, args, kwargs, codec_ctx + ) except errors.QueryArgumentError as ex: exc = ex finally: @@ -472,7 +475,7 @@ cdef class SansIOProtocol: else: assert self.state_codec is not None buf = WriteBuffer.new() - self.state_codec.encode(buf, state) + self.state_codec.encode(buf, state, DEFAULT_CODEC_CONTEXT) state_data = bytes(buf) self.state_cache = (state, state_data) return self.state_type_id, state_data @@ -495,6 +498,7 @@ cdef class SansIOProtocol: inline_typeids: bool = False, allow_capabilities: enums.Capability = enums.Capability.ALL, state: typing.Optional[dict] = None, + codec_ctx: pgproto.CodecContext, ): cdef: BaseCodec in_dc @@ -570,6 +574,7 @@ cdef class SansIOProtocol: in_dc=in_dc, out_dc=out_dc, state=state, + codec_ctx=codec_ctx, ) async def query( @@ -588,6 +593,7 @@ cdef class SansIOProtocol: inline_typeids: bool = False, allow_capabilities: enums.Capability = enums.Capability.ALL, state: typing.Optional[dict] = None, + codec_ctx: pgproto.CodecContext, ): ret = await self.execute( query=query, @@ -603,6 +609,7 @@ cdef class SansIOProtocol: inline_typeids=inline_typeids, allow_capabilities=allow_capabilities, state=state, + codec_ctx=codec_ctx, ) if expect_one: @@ -1018,7 +1025,7 @@ cdef class SansIOProtocol: frb_init(rbuf, buf, buf_len) - return decoder(codec, rbuf) + return decoder(codec, rbuf, DEFAULT_CODEC_CONTEXT) cdef parse_server_settings(self, str name, bytes val): if name == 'suggested_pool_concurrency': @@ -1071,7 +1078,14 @@ cdef class SansIOProtocol: raise errors.ProtocolError( f'unexpected message type {chr(mtype)!r}') - cdef encode_args(self, BaseCodec in_dc, WriteBuffer buf, args, kwargs): + cdef encode_args( + self, + BaseCodec in_dc, + WriteBuffer buf, + args, + kwargs, + pgproto.CodecContext codec_ctx + ): if args and kwargs: raise errors.QueryArgumentError( 'either positional or named arguments are supported; ' @@ -1097,7 +1111,7 @@ cdef class SansIOProtocol: if args: kwargs = {str(i): v for i, v in enumerate(args)} - (in_dc).encode_args(buf, kwargs) + (in_dc).encode_args(buf, kwargs, codec_ctx) cdef parse_describe_type_message(self, CodecsRegistry reg): assert self.buffer.get_message_type() == COMMAND_DATA_DESC_MSG @@ -1158,7 +1172,9 @@ cdef class SansIOProtocol: return in_dc, out_dc - cdef parse_data_messages(self, BaseCodec out_dc, result): + cdef parse_data_messages( + self, BaseCodec out_dc, result, pgproto.CodecContext codec_ctx + ): cdef: ReadBuffer buf = self.buffer @@ -1209,7 +1225,7 @@ cdef class SansIOProtocol: # so we want to skip first 6 bytes: frb_init(rbuf, cbuf + 6, cbuf_len - 6) - row = decoder(out_dc, rbuf) + row = decoder(out_dc, rbuf, codec_ctx) result.append(row) if frb_get_len(rbuf): diff --git a/edgedb/protocol/protocol_v0.pyx b/edgedb/protocol/protocol_v0.pyx index 399f803c3..6c67c8ee9 100644 --- a/edgedb/protocol/protocol_v0.pyx +++ b/edgedb/protocol/protocol_v0.pyx @@ -154,7 +154,12 @@ cdef class SansIOProtocolBackwardsCompatible(SansIOProtocol): return cardinality, in_dc, out_dc, attrs async def _legacy_execute( - self, BaseCodec in_dc, BaseCodec out_dc, args, kwargs + self, + BaseCodec in_dc, + BaseCodec out_dc, + args, + kwargs, + pgproto.CodecContext codec_ctx, ): cdef: WriteBuffer packet @@ -169,7 +174,7 @@ cdef class SansIOProtocolBackwardsCompatible(SansIOProtocol): buf = WriteBuffer.new_message(LEGACY_EXECUTE_MSG) buf.write_int16(0) # no headers buf.write_len_prefixed_bytes(b'') # stmt_name - self.encode_args(in_dc, buf, args, kwargs) + self.encode_args(in_dc, buf, args, kwargs, codec_ctx) packet.write_buffer(buf.end_message()) packet.write_bytes(SYNC_MESSAGE) @@ -187,7 +192,7 @@ cdef class SansIOProtocolBackwardsCompatible(SansIOProtocol): if mtype == DATA_MSG: if exc is None: try: - self.parse_data_messages(out_dc, result) + self.parse_data_messages(out_dc, result, codec_ctx) except Exception as ex: # An error during data decoding. We need to # handle this as gracefully as possible: @@ -242,6 +247,7 @@ cdef class SansIOProtocolBackwardsCompatible(SansIOProtocol): allow_capabilities: typing.Optional[int] = None, in_dc: BaseCodec, out_dc: BaseCodec, + pgproto.CodecContext codec_ctx, ): cdef: WriteBuffer packet @@ -261,7 +267,7 @@ cdef class SansIOProtocolBackwardsCompatible(SansIOProtocol): buf.write_len_prefixed_utf8(query) buf.write_bytes(in_dc.get_tid()) buf.write_bytes(out_dc.get_tid()) - self.encode_args(in_dc, buf, args, kwargs) + self.encode_args(in_dc, buf, args, kwargs, codec_ctx) buf.end_message() packet = WriteBuffer.new() @@ -302,7 +308,7 @@ cdef class SansIOProtocolBackwardsCompatible(SansIOProtocol): assert not re_exec if exc is None: try: - self.parse_data_messages(out_dc, result) + self.parse_data_messages(out_dc, result, codec_ctx) except Exception as ex: # An error during data decoding. We need to # handle this as gracefully as possible: @@ -347,7 +353,9 @@ cdef class SansIOProtocolBackwardsCompatible(SansIOProtocol): raise errors.InterfaceError( f'query cannot be executed with {methname}() as it ' f'does not return any data') - return await self._legacy_execute(in_dc, out_dc, args, kwargs) + return await self._legacy_execute( + in_dc, out_dc, args, kwargs, codec_ctx + ) else: return result @@ -366,6 +374,7 @@ cdef class SansIOProtocolBackwardsCompatible(SansIOProtocol): inline_typenames: bool = False, inline_typeids: bool = False, allow_capabilities: enums.Capability = enums.Capability.ALL, + pgproto.CodecContext codec_ctx, ): cdef: BaseCodec in_dc @@ -418,7 +427,9 @@ cdef class SansIOProtocolBackwardsCompatible(SansIOProtocol): capabilities, ) - ret = await self._legacy_execute(in_dc, out_dc, args, kwargs) + ret = await self._legacy_execute( + in_dc, out_dc, args, kwargs, codec_ctx + ) else: has_na_cardinality = codecs[0] diff --git a/edgedb/transaction.py b/edgedb/transaction.py index 511b8f42e..e3e13aacb 100644 --- a/edgedb/transaction.py +++ b/edgedb/transaction.py @@ -22,6 +22,7 @@ from . import abstract from . import errors from . import options +from .pgproto import pgproto class TransactionState(enum.Enum): @@ -185,6 +186,9 @@ def _get_query_cache(self) -> abstract.QueryCache: def _get_state(self) -> options.State: return self._client._get_state() + def _get_codec_ctx(self) -> pgproto.CodecContext: + return self._client._get_codec_ctx() + async def _query(self, query_context: abstract.QueryContext): await self._ensure_transaction() return await self._connection.raw_query(query_context) @@ -198,6 +202,7 @@ async def _privileged_execute(self, query: str) -> None: query=abstract.QueryWithArgs(query, (), {}), cache=self._get_query_cache(), state=self._get_state(), + codec_ctx=self._get_codec_ctx(), )) diff --git a/tests/codegen/test-project1/generated_async_edgeql.py.assert b/tests/codegen/test-project1/generated_async_edgeql.py.assert index 9a25de466..d98a36c02 100644 --- a/tests/codegen/test-project1/generated_async_edgeql.py.assert +++ b/tests/codegen/test-project1/generated_async_edgeql.py.assert @@ -1,5 +1,5 @@ # AUTOGENERATED FROM 'select_scalar.edgeql' WITH: -# $ edgedb-py --target async --file --no-skip-pydantic-validation +# $ edgedb-py --target async --file --no-skip-pydantic-validation --handle-json from __future__ import annotations @@ -9,6 +9,10 @@ import edgedb async def select_scalar( client: edgedb.AsyncIOClient, ) -> int: + client = client.with_codec_context( + edgedb.get_default_codec_context(handle_json=True), + only_replace_default=True, + ) return await client.query_single( """\ select 1;\ diff --git a/tests/codegen/test-project2/generated_async_edgeql.py.assert b/tests/codegen/test-project2/generated_async_edgeql.py.assert index b337aa8bc..7e6f64cb8 100644 --- a/tests/codegen/test-project2/generated_async_edgeql.py.assert +++ b/tests/codegen/test-project2/generated_async_edgeql.py.assert @@ -8,7 +8,7 @@ # 'scalar/select_scalar.edgeql' # 'scalar/select_scalars.edgeql' # WITH: -# $ edgedb-py --target async --file --no-skip-pydantic-validation +# $ edgedb-py --target async --file --no-skip-pydantic-validation --handle-json from __future__ import annotations @@ -137,6 +137,10 @@ class SelectObjectResultParamsItem: async def link_prop( client: edgedb.AsyncIOClient, ) -> list[LinkPropResult]: + client = client.with_codec_context( + edgedb.get_default_codec_context(handle_json=True), + only_replace_default=True, + ) return await client.query( """\ create type Person { @@ -215,6 +219,10 @@ async def my_query( aw: edgedb.Range[datetime.date], ax: edgedb.Range[datetime.date], ) -> MyQueryResult: + client = client.with_codec_context( + edgedb.get_default_codec_context(handle_json=True), + only_replace_default=True, + ) return await client.query_single( """\ create scalar type MyScalar extending int64; @@ -335,6 +343,10 @@ async def query_one( *, arg_name_with_underscores: int, ) -> int: + client = client.with_codec_context( + edgedb.get_default_codec_context(handle_json=True), + only_replace_default=True, + ) return await client.query_single( """\ select $arg_name_with_underscores\ @@ -349,6 +361,10 @@ async def select_args( arg_str: str, arg_datetime: datetime.datetime, ) -> SelectArgsResult: + client = client.with_codec_context( + edgedb.get_default_codec_context(handle_json=True), + only_replace_default=True, + ) return await client.query_single( """\ select { @@ -364,6 +380,10 @@ async def select_args( async def select_object( client: edgedb.AsyncIOClient, ) -> SelectObjectResult | None: + client = client.with_codec_context( + edgedb.get_default_codec_context(handle_json=True), + only_replace_default=True, + ) return await client.query_single( """\ select schema::Function { @@ -382,6 +402,10 @@ async def select_object( async def select_objects( client: edgedb.AsyncIOClient, ) -> list[SelectObjectResult]: + client = client.with_codec_context( + edgedb.get_default_codec_context(handle_json=True), + only_replace_default=True, + ) return await client.query( """\ select schema::Function { @@ -399,6 +423,10 @@ async def select_objects( async def select_scalar( client: edgedb.AsyncIOClient, ) -> int: + client = client.with_codec_context( + edgedb.get_default_codec_context(handle_json=True), + only_replace_default=True, + ) return await client.query_single( """\ select 1;\ @@ -409,6 +437,10 @@ async def select_scalar( async def select_scalars( client: edgedb.AsyncIOClient, ) -> list[edgedb.ConfigMemory]: + client = client.with_codec_context( + edgedb.get_default_codec_context(handle_json=True), + only_replace_default=True, + ) return await client.query( """\ select {1, 2, 3};\ diff --git a/tests/test_async_query.py b/tests/test_async_query.py index af3342166..2968f3977 100644 --- a/tests/test_async_query.py +++ b/tests/test_async_query.py @@ -925,6 +925,7 @@ async def test_json_elements(self): ), retry_options=None, state=None, + codec_ctx=edgedb.get_default_codec_context(), ) ) self.assertEqual( diff --git a/tests/test_codegen.py b/tests/test_codegen.py index dccf3a98b..515dd982b 100644 --- a/tests/test_codegen.py +++ b/tests/test_codegen.py @@ -82,6 +82,7 @@ async def run(*args, extra_env=None): "async", "--file", "--no-skip-pydantic-validation", + "--handle-json", extra_env={"EDGEDB_PYTHON_CODEGEN_PY_VER": "3.10.3"}, ) diff --git a/tests/test_sync_query.py b/tests/test_sync_query.py index 0d7936167..6b739fd56 100644 --- a/tests/test_sync_query.py +++ b/tests/test_sync_query.py @@ -794,6 +794,7 @@ def test_json_elements(self): ), retry_options=None, state=None, + codec_ctx=edgedb.get_default_codec_context(), ) ) )