Codegen with the describe_query() API (#363)

* Change the lower bound of type annotations tag to 0x80, refs geldata/gel#1958
* Add `client._describe()` API
* Add missing copyright header
* Implement `_describe()` API on codecs
* Implement codegen using `_describe()` API
* Support `edgedb.Range` type (Pydantic support for edgedb.Range type is not included out of the box)
* Support Python 3.9 and 3.10 typing simplification
* Add header comments to generated files
* Add e2e test for codegen (test case copied from edgedb-go)
* Support typed link properties

Co-authored-by: fmoor <fmoor@gmx.com>

Author: fantix

edgedb/__init__.py

@@ -28,7 +28,7 @@
 from edgedb.datatypes.range import Range
 
 from .abstract import (
-    Executor, AsyncIOExecutor, ReadOnlyExecutor, AsyncIOReadOnlyExecutor
+    Executor, AsyncIOExecutor, ReadOnlyExecutor, AsyncIOReadOnlyExecutor,
 )
 
 from .asyncio_client import (
@@ -37,6 +37,7 @@
 )
 
 from .blocking_client import create_client, Client
+from .enums import Cardinality, ElementKind
 from .options import RetryCondition, IsolationLevel, default_backoff
 from .options import RetryOptions, TransactionOptions
 from .options import State

edgedb/_testbase.py

@@ -584,3 +584,8 @@ def gen_lock_key():
 if os.environ.get('USE_UVLOOP'):
     import uvloop
     uvloop.install()
+elif sys.platform == 'win32' and sys.version_info[:2] == (3, 7):
+    # The default policy on win32 of Python 3.7 is SelectorEventLoop, which
+    # does not implement some subprocess functions required by the tests,
+    # so we have to manually set it to the proactor one (default in 3.8).
+    asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())

edgedb/abstract.py

@@ -1,6 +1,29 @@
+#
+# This source file is part of the EdgeDB open source project.
+#
+# Copyright 2020-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.
+#
+
+
+from __future__ import annotations
 import abc
+import dataclasses
 import typing
 
+from . import describe
+from . import enums
 from . import options
 from .protocol import protocol
 
@@ -13,6 +36,8 @@
     "AsyncIOExecutor",
     "ReadOnlyExecutor",
     "AsyncIOReadOnlyExecutor",
+    "DescribeContext",
+    "DescribeResult",
 )
 
 
@@ -47,6 +72,21 @@ class ExecuteContext(typing.NamedTuple):
     state: typing.Optional[options.State]
 
 
+@dataclasses.dataclass
+class DescribeContext:
+    query: str
+    state: typing.Optional[options.State]
+    inject_type_names: bool
+
+
+@dataclasses.dataclass
+class DescribeResult:
+    input_type: typing.Optional[describe.AnyType]
+    output_type: typing.Optional[describe.AnyType]
+    output_cardinality: enums.Cardinality
+    capabilities: enums.Capability
+
+
 _query_opts = QueryOptions(
     output_format=protocol.OutputFormat.BINARY,
     expect_one=False,

edgedb/asyncio_client.py

@@ -347,6 +347,15 @@ async def __aenter__(self):
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         await self.aclose()
 
+    async def _describe_query(
+        self, query: str, *, inject_type_names: bool = False
+    ) -> abstract.DescribeResult:
+        return await self._describe(abstract.DescribeContext(
+            query=query,
+            state=self._get_state(),
+            inject_type_names=inject_type_names,
+        ))
+
 
 def create_async_client(
     dsn=None,

edgedb/base_client.py

@@ -287,6 +287,25 @@ async def _execute(self, execute_context: abstract.ExecuteContext) -> None:
                 ),
             )
 
+    async def describe(
+        self, describe_context: abstract.DescribeContext
+    ) -> abstract.DescribeResult:
+        cardinality, in_dc, out_dc, capabilities = await self._protocol._parse(
+            describe_context.query,
+            reg=protocol.CodecsRegistry(),
+            inline_typenames=describe_context.inject_type_names,
+            state=(
+                describe_context.state.as_dict()
+                if describe_context.state else None
+            ),
+        )
+        return abstract.DescribeResult(
+            input_type=in_dc.make_type(describe_context),
+            output_type=out_dc.make_type(describe_context),
+            output_cardinality=enums.Cardinality(cardinality[0]),
+            capabilities=capabilities,
+        )
+
     def terminate(self):
         if not self.is_closed():
             try:
@@ -739,6 +758,15 @@ async def _execute(self, execute_context: abstract.ExecuteContext) -> None:
         finally:
             await self._impl.release(con)
 
+    async def _describe(
+        self, describe_context: abstract.DescribeContext
+    ) -> abstract.DescribeResult:
+        con = await self._impl.acquire()
+        try:
+            return await con.describe(describe_context)
+        finally:
+            await self._impl.release(con)
+
     def terminate(self):
         """Terminate all connections in the pool."""
         self._impl.terminate()

edgedb/blocking_client.py

@@ -327,6 +327,15 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
+    def _describe_query(
+        self, query: str, *, inject_type_names: bool = False
+    ) -> abstract.DescribeResult:
+        return self._iter_coroutine(self._describe(abstract.DescribeContext(
+            query=query,
+            state=self._get_state(),
+            inject_type_names=inject_type_names,
+        )))
+
 
 def create_client(
     dsn=None,

edgedb/codegen/__init__.py

@@ -0,0 +1,17 @@
+#
+# 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.
+#

edgedb/codegen/__main__.py

@@ -0,0 +1,23 @@
+#
+# 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.
+#
+
+
+from .cli import main
+
+if __name__ == "__main__":
+    main()

edgedb/codegen/cli.py

@@ -0,0 +1,53 @@
+#
+# 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.
+#
+
+
+import argparse
+
+from . import generator
+
+
+parser = argparse.ArgumentParser(
+    description="Generate Python code for .edgeql files."
+)
+parser.add_argument("--dsn")
+parser.add_argument("--credentials_file", metavar="PATH")
+parser.add_argument("-I", "--instance", metavar="NAME")
+parser.add_argument("-H", "--host")
+parser.add_argument("-P", "--port")
+parser.add_argument("-d", "--database", metavar="NAME")
+parser.add_argument("-u", "--user")
+parser.add_argument("--password")
+parser.add_argument("--password-from-stdin", action="store_true")
+parser.add_argument("--tls-ca-file", metavar="PATH")
+parser.add_argument(
+    "--tls-security",
+    choices=["default", "strict", "no_host_verification", "insecure"],
+)
+parser.add_argument("--file", action="store_true")
+parser.add_argument(
+    "--target",
+    choices=["blocking", "async", "pydantic"],
+    nargs="*",
+    default=["async", "pydantic"],
+)
+
+
+def main():
+    args = parser.parse_args()
+    generator.Generator(args).run()