Add support for protocol 3.0 (#534)

This is basically just adding support for specifying the input language in `Parse` and `Execute`. Expose that as the new `query_sql()` and `execute_sql()` client methods.
geldata · Nov 13, 2024 · 2e6a877 · 2e6a877
1 parent d0b3961
commit 2e6a877
Show file tree

Hide file tree

Showing 6 changed files with 78 additions and 3 deletions.
diff --git a/edgedb/abstract.py b/edgedb/abstract.py
@@ -45,6 +45,7 @@ class QueryWithArgs(typing.NamedTuple):
     query: str
     args: typing.Tuple
     kwargs: typing.Dict[str, typing.Any]
+    input_language: protocol.InputLanguage = protocol.InputLanguage.EDGEQL
 
 
 class QueryCache(typing.NamedTuple):
@@ -75,6 +76,7 @@ def lower(
             kwargs=self.query.kwargs,
             reg=self.cache.codecs_registry,
             qc=self.cache.query_cache,
+            input_language=self.query.input_language,
             output_format=self.query_options.output_format,
             expect_one=self.query_options.expect_one,
             required_one=self.query_options.required_one,
@@ -98,6 +100,7 @@ def lower(
             kwargs=self.query.kwargs,
             reg=self.cache.codecs_registry,
             qc=self.cache.query_cache,
+            input_language=self.query.input_language,
             output_format=protocol.OutputFormat.NONE,
             allow_capabilities=allow_capabilities,
             state=self.state.as_dict() if self.state else None,
@@ -109,6 +112,7 @@ class DescribeContext:
     query: str
     state: typing.Optional[options.State]
     inject_type_names: bool
+    input_language: protocol.InputLanguage
     output_format: protocol.OutputFormat
     expect_one: bool
 
@@ -121,6 +125,7 @@ def lower(
             kwargs=None,
             reg=protocol.CodecsRegistry(),
             qc=protocol.LRUMapping(maxsize=1),
+            input_language=self.input_language,
             output_format=self.output_format,
             expect_one=self.expect_one,
             inline_typenames=self.inject_type_names,
@@ -259,6 +264,21 @@ def query_required_single_json(self, query: str, *args, **kwargs) -> str:
             warning_handler=self._get_warning_handler(),
         ))
 
+    def query_sql(self, query: str, *args, **kwargs) -> typing.Any:
+        return self._query(QueryContext(
+            query=QueryWithArgs(
+                query,
+                args,
+                kwargs,
+                input_language=protocol.InputLanguage.SQL,
+            ),
+            cache=self._get_query_cache(),
+            query_options=_query_opts,
+            retry_options=self._get_retry_options(),
+            state=self._get_state(),
+            warning_handler=self._get_warning_handler(),
+        ))
+
     @abc.abstractmethod
     def _execute(self, execute_context: ExecuteContext):
         ...
@@ -271,6 +291,19 @@ def execute(self, commands: str, *args, **kwargs) -> None:
             warning_handler=self._get_warning_handler(),
         ))
 
+    def execute_sql(self, commands: str, *args, **kwargs) -> None:
+        self._execute(ExecuteContext(
+            query=QueryWithArgs(
+                commands,
+                args,
+                kwargs,
+                input_language=protocol.InputLanguage.SQL,
+            ),
+            cache=self._get_query_cache(),
+            state=self._get_state(),
+            warning_handler=self._get_warning_handler(),
+        ))
+
 
 class Executor(ReadOnlyExecutor):
     """Subclasses can execute both read-only and modification queries"""
@@ -357,6 +390,21 @@ async def query_required_single_json(
             warning_handler=self._get_warning_handler(),
         ))
 
+    async def query_sql(self, query: str, *args, **kwargs) -> typing.Any:
+        return await self._query(QueryContext(
+            query=QueryWithArgs(
+                query,
+                args,
+                kwargs,
+                input_language=protocol.InputLanguage.SQL,
+            ),
+            cache=self._get_query_cache(),
+            query_options=_query_opts,
+            retry_options=self._get_retry_options(),
+            state=self._get_state(),
+            warning_handler=self._get_warning_handler(),
+        ))
+
     @abc.abstractmethod
     async def _execute(self, execute_context: ExecuteContext) -> None:
         ...
@@ -369,6 +417,19 @@ async def execute(self, commands: str, *args, **kwargs) -> None:
             warning_handler=self._get_warning_handler(),
         ))
 
+    async def execute_sql(self, commands: str, *args, **kwargs) -> None:
+        await self._execute(ExecuteContext(
+            query=QueryWithArgs(
+                commands,
+                args,
+                kwargs,
+                input_language=protocol.InputLanguage.SQL,
+            ),
+            cache=self._get_query_cache(),
+            state=self._get_state(),
+            warning_handler=self._get_warning_handler(),
+        ))
+
 
 class AsyncIOExecutor(AsyncIOReadOnlyExecutor):
     """Subclasses can execute both read-only and modification queries"""

diff --git a/edgedb/asyncio_client.py b/edgedb/asyncio_client.py
@@ -30,7 +30,7 @@
 from . import errors
 from . import transaction
 from .protocol import asyncio_proto
-from .protocol.protocol import OutputFormat
+from .protocol.protocol import InputLanguage, OutputFormat
 
 
 __all__ = (
@@ -392,13 +392,15 @@ async def _describe_query(
         query: str,
         *,
         inject_type_names: bool = False,
+        input_language: InputLanguage = InputLanguage.EDGEQL,
         output_format: OutputFormat = OutputFormat.BINARY,
         expect_one: bool = False,
     ) -> abstract.DescribeResult:
         return await self._describe(abstract.DescribeContext(
             query=query,
             state=self._get_state(),
             inject_type_names=inject_type_names,
+            input_language=input_language,
             output_format=output_format,
             expect_one=expect_one,
         ))

diff --git a/edgedb/blocking_client.py b/edgedb/blocking_client.py
@@ -32,7 +32,7 @@
 from . import errors
 from . import transaction
 from .protocol import blocking_proto
-from .protocol.protocol import OutputFormat
+from .protocol.protocol import InputLanguage, OutputFormat
 
 
 DEFAULT_PING_BEFORE_IDLE_TIMEOUT = datetime.timedelta(seconds=5)
@@ -434,13 +434,15 @@ def _describe_query(
         query: str,
         *,
         inject_type_names: bool = False,
+        input_language: InputLanguage = InputLanguage.EDGEQL,
         output_format: OutputFormat = OutputFormat.BINARY,
         expect_one: bool = False,
     ) -> abstract.DescribeResult:
         return self._iter_coroutine(self._describe(abstract.DescribeContext(
             query=query,
             state=self._get_state(),
             inject_type_names=inject_type_names,
+            input_language=input_language,
             output_format=output_format,
             expect_one=expect_one,
         )))

diff --git a/edgedb/protocol/consts.pxi b/edgedb/protocol/consts.pxi
@@ -61,7 +61,7 @@ DEF TRANS_STATUS_IDLE = b'I'
 DEF TRANS_STATUS_INTRANS = b'T'
 DEF TRANS_STATUS_ERROR = b'E'
 
-DEF PROTO_VER_MAJOR = 2
+DEF PROTO_VER_MAJOR = 3
 DEF PROTO_VER_MINOR = 0
 
 DEF MIN_PROTO_VER_MAJOR = 0

diff --git a/edgedb/protocol/protocol.pxd b/edgedb/protocol/protocol.pxd
@@ -40,6 +40,11 @@ include "./codecs/codecs.pxd"
 ctypedef object (*decode_row_method)(BaseCodec, FRBuffer *buf)
 
 
+cpdef enum InputLanguage:
+    EDGEQL = 0x45  # b'E'
+    SQL = 0x53  # b'S'
+
+
 cpdef enum OutputFormat:
     BINARY = 98  # b'b'
     JSON = 106  # b'j'
@@ -75,6 +80,7 @@ cdef class ExecuteContext:
         object kwargs
         CodecsRegistry reg
         LRUMapping qc
+        InputLanguage input_language
         OutputFormat output_format
         bint expect_one
         bint required_one

diff --git a/edgedb/protocol/protocol.pyx b/edgedb/protocol/protocol.pyx
@@ -100,6 +100,7 @@ cdef class ExecuteContext:
         kwargs,
         reg: CodecsRegistry,
         qc: LRUMapping,
+        input_language: InputLanguage,
         output_format: OutputFormat,
         expect_one: bool = False,
         required_one: bool = False,
@@ -114,6 +115,7 @@ cdef class ExecuteContext:
         self.kwargs = kwargs
         self.reg = reg
         self.qc = qc
+        self.input_language = input_language
         self.output_format = output_format
         self.expect_one = bool(expect_one)
         self.required_one = bool(required_one)
@@ -271,6 +273,8 @@ cdef class SansIOProtocol:
         buf.write_int64(<int64_t>ctx.allow_capabilities)
         buf.write_int64(<int64_t><uint64_t>compilation_flags)
         buf.write_int64(<int64_t>ctx.implicit_limit)
+        if self.protocol_version >= (3, 0):
+            buf.write_byte(ctx.input_language)
         buf.write_byte(ctx.output_format)
         buf.write_byte(CARDINALITY_ONE if ctx.expect_one else CARDINALITY_MANY)
         buf.write_len_prefixed_utf8(ctx.query)