Add a way to export SQLAlchemy models from Gel (EdgeDB).

Currently just have a function `render_models` that takes a client in
order to print SQLAlchemy models. They are printed to stdout which can
be redirected as needed.

Add tests that setup a Gel database and then generate the SQLAlchemy
models from it. The individual tests use a SQLAlchemy session to access
the database using postgres protocol.

Author: vpetrovykh

.flake8

@@ -1,3 +1,3 @@
 [flake8]
 ignore = B008,B023,B306,E203,E402,E731,D100,D101,D102,D103,D104,D105,W503,W504,E252,F999,F541
-exclude = .git,__pycache__,build,dist,.eggs
+exclude = .git,__pycache__,build,dist,.eggs,generated

edgedb/_testbase.py

@@ -35,7 +35,7 @@
 import edgedb
 from edgedb import asyncio_client
 from edgedb import blocking_client
-
+from edgedb.gelalchemy.models import get_schema_json, render_models
 
 log = logging.getLogger(__name__)
 
@@ -598,6 +598,34 @@ def adapt_call(cls, result):
         return result
 
 
+class SQLATestCase(SyncQueryTestCase):
+    SQLAMODELS = None
+    DEFAULT_MODULE = 'default'
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+
+        class_set_up = os.environ.get('EDGEDB_TEST_CASES_SET_UP')
+        if not class_set_up:
+            # Now that the DB is setup, generate the SQLAlchemy models from it
+            spec = get_schema_json(cls.client)
+
+            with open(cls.SQLAMODELS, 'w') as f:
+                with contextlib.redirect_stdout(f):
+                    render_models(spec)
+
+    @classmethod
+    def get_dsn_for_sqla(cls):
+        cargs = cls.get_connect_args(database=cls.get_database_name())
+        dsn = (
+            f'postgresql://{cargs["user"]}:{cargs["password"]}'
+            f'@{cargs["host"]}:{cargs["port"]}/{cargs["database"]}'
+        )
+
+        return dsn
+
+
 _lock_cnt = 0
 
 

edgedb/gelalchemy/__init__.py

@@ -0,0 +1,271 @@
+import edgedb
+import json
+
+
+DSN = 'edgedb://edgedb@localhost:5656/edgedb?tls_security=insecure'
+client = edgedb.create_client(DSN)
+INDENT = ' ' * 4
+
+GEL_SCALAR_MAP = {
+    'std::bool': ('bool', 'Boolean'),
+    'std::str': ('str', 'String'),
+    'std::int16': ('int', 'Integer'),
+    'std::int32': ('int', 'Integer'),
+    'std::int64': ('int', 'Integer'),
+    'std::float32': ('float', 'Float'),
+    'std::float64': ('float', 'Float'),
+    'std::uuid': ('uuid.UUID', 'Uuid'),
+}
+
+INTRO_QUERY = '''
+with module schema
+select ObjectType {
+    name,
+    links: {
+        name,
+        readonly,
+        required,
+        cardinality,
+        exclusive := exists (
+            select .constraints
+            filter .name = 'std::exclusive'
+        ),
+        target: {name},
+
+        properties: {
+            name,
+            readonly,
+            required,
+            cardinality,
+            exclusive := exists (
+                select .constraints
+                filter .name = 'std::exclusive'
+            ),
+            target: {name},
+        } filter .name not in {'source', 'target'},
+    },
+    properties: {
+        name,
+        readonly,
+        required,
+        cardinality,
+        exclusive := exists (
+            select .constraints
+            filter .name = 'std::exclusive'
+        ),
+        target: {name},
+    },
+    backlinks := <array<str>>[],
+}
+filter
+    not .builtin
+    and
+    not .internal
+    and
+    not re_test('^(std|cfg|sys|schema)::', .name);
+'''
+
+
+def get_object_type_name(name):
+    # for now just assume stuff is in "defualt" module
+    return name.replace('default::', '')
+
+
+def get_schema_json(client):
+    types = json.loads(client.query_json(INTRO_QUERY))
+
+    return _process_links(types)
+
+
+async def async_get_schema_json(client):
+    types = json.loads(await client.query_json(INTRO_QUERY))
+
+    return _process_links(types)
+
+
+def _process_links(types):
+    # Figure out all the backlinks and link tables
+    link_tables = []
+    type_map = {}
+    for spec in types:
+        type_map[get_object_type_name(spec['name'])] = spec
+
+    for spec in types:
+        for link in spec['links']:
+            if link['name'] != '__type__':
+                target = get_object_type_name(link['target']['name'])
+                cardinality = link['cardinality']
+                exclusive = link['exclusive']
+
+                objtype = type_map[target]
+                objtype['backlinks'].append({
+                    'name': f'backlink_via_{link["name"]}',
+                    'cardinality': 'One' if exclusive else 'Many',
+                    'exclusive': cardinality == 'One',
+                    'target': {'name': spec['name']},
+                })
+
+                # Add a link table for One-to-Many and Many-to-Many
+                if cardinality == 'Many':
+                    source = get_object_type_name(spec["name"])
+                    link_tables.append({
+                        'name': f'{source}_{link["name"]}_table',
+                        'table': f'{source}.{link["name"]}',
+                        'source': f'{source}.id',
+                        'target': f'{target}.id',
+                    })
+
+    return {
+        'object_types': types,
+        'link_tables': link_tables,
+    }
+
+
+def render_link_table(spec):
+    print(f'''\
+{spec["name"]} = Table(
+    {spec["table"]!r},
+    Base.metadata,
+    Column("source", ForeignKey({spec["source"]!r})),
+    Column("target", ForeignKey({spec["target"]!r})),
+)\
+    ''')
+
+
+def render_type(spec):
+    # assume nice names for now
+    name = get_object_type_name(spec['name'])
+
+    print(f'class {name}(Base):')
+    print(f'{INDENT}__tablename__ = {name!r}')
+
+    # Add two fields that all objects have
+    print(
+        f'''
+    id: Mapped[uuid.UUID] = mapped_column(
+        Uuid(), primary_key=True, server_default='uuid_generate_v4()')
+    # This is maintained entirely by Gel, the server_default simply indicates
+    # to SQLAlchemy that this value may be omitted.
+    gel_type_id: Mapped[uuid.UUID] = mapped_column(
+        '__type__', Uuid(), unique=True, server_default='PLACEHOLDER')'''
+    )
+
+    if spec['properties']:
+        print(f'\n    # Properties:')
+        for prop in spec['properties']:
+            if prop['name'] != 'id':
+                render_prop(prop)
+
+    if spec['links']:
+        print(f'\n    # Links:')
+        for link in spec['links']:
+            if link['name'] != '__type__':
+                render_link(link, name)
+
+    if spec['backlinks']:
+        print(f'\n    # Back-links:')
+        for link in spec['backlinks']:
+            render_backlink(link)
+
+
+def render_prop(spec):
+    name = spec['name']
+    nullable = not spec['required']
+
+    pytype, sqlatype = GEL_SCALAR_MAP[spec['target']['name']]
+
+    print(
+        f'\
+    {name}: Mapped[{pytype}] = '
+        f'mapped_column({sqlatype}(), nullable={nullable})'
+    )
+
+
+def render_link(spec, parent):
+    name = spec['name']
+    nullable = not spec['required']
+    target = get_object_type_name(spec['target']['name'])
+    cardinality = spec['cardinality']
+    bklink = f'backlink_via_{name}'
+
+    if cardinality == 'One':
+        print(
+            f'{INDENT}{name}_id: Mapped[uuid.UUID] = '
+            f'mapped_column(Uuid(), ForeignKey("{target}.id"), '
+            f'nullable={nullable})'
+        )
+        print(
+            f'{INDENT}{name}: Mapped[{target!r}] = '
+            f'relationship(back_populates={bklink!r})'
+        )
+    elif cardinality == 'Many':
+        secondary = f'{parent}_{name}_table'
+        print(f'{INDENT}{name}: Mapped[List[{target!r}]] = relationship(')
+        print(
+            f'{INDENT * 2}{target!r}, secondary={secondary}, '
+            f'back_populates={bklink!r},'
+        )
+        print(f'{INDENT})')
+
+
+def render_backlink(spec):
+    name = spec['name']
+    target = get_object_type_name(spec['target']['name'])
+    cardinality = spec['cardinality']
+    exclusive = spec['exclusive']
+    bklink = name.replace('backlink_via_', '', 1)
+
+    if exclusive:
+        # This is a backlink from a single link. There is no link table
+        # involved.
+        if cardinality == 'One':
+            print(
+                f'{INDENT}{name}: Mapped[{target!r}] = '
+                f'relationship(back_populates={bklink!r})'
+            )
+        elif cardinality == 'Many':
+            print(
+                f'{INDENT}{name}: Mapped[List[{target!r}]] = '
+                f'relationship(back_populates={bklink!r})'
+            )
+
+    else:
+        # This backlink involves a link table, so we still treat it as a
+        # Many-to-Many.
+        secondary = f'{target}_{bklink}_table'
+        print(f'{INDENT}{name}: Mapped[List[{target!r}]] = relationship(')
+        print(
+            f'{INDENT * 2}{target!r}, secondary={secondary}, '
+            f'back_populates={bklink!r},'
+        )
+        print(f'{INDENT})')
+
+
+def render_models(spec):
+    print(f'''\
+#
+# Automatically generated from Gel schema.
+#
+
+import uuid
+
+from typing import List
+from typing import Optional
+
+from sqlalchemy import MetaData, Table, Column, ForeignKey
+from sqlalchemy import String, Uuid, Integer, Float, Boolean
+from sqlalchemy.orm import (DeclarativeBase, Mapped, mapped_column,
+                            relationship)
+
+
+class Base(DeclarativeBase):
+    pass
+    ''')
+
+    for rec in spec['link_tables']:
+        render_link_table(rec)
+        print('\n')
+
+    for rec in spec['object_types']:
+        render_type(rec)
+        print('\n')

edgedb/gelalchemy/models.py

@@ -0,0 +1,44 @@
+insert User {name := 'Alice'};
+insert User {name := 'Billie'};
+insert User {name := 'Cameron'};
+insert User {name := 'Dana'};
+insert User {name := 'Elsa'};
+insert User {name := 'Zoe'};
+
+insert UserGroup {
+    name := 'red',
+    users := (select User filter .name != 'Zoe'),
+};
+insert UserGroup {
+    name := 'green',
+    users := (select User filter .name in {'Alice', 'Billie'}),
+};
+insert UserGroup {
+    name := 'blue',
+};
+
+insert GameSession {
+    num := 123,
+    players := (select User filter .name in {'Alice', 'Billie'}),
+};
+insert GameSession {
+    num := 456,
+    players := (select User filter .name in {'Dana'}),
+};
+
+insert Post {
+    author := assert_single((select User filter .name = 'Alice')),
+    body := 'Hello',
+};
+insert Post {
+    author := assert_single((select User filter .name = 'Alice')),
+    body := "I'm Alice",
+};
+insert Post {
+    author := assert_single((select User filter .name = 'Cameron')),
+    body := "I'm Cameron",
+};
+insert Post {
+    author := assert_single((select User filter .name = 'Elsa')),
+    body := '*magic stuff*',
+};

tests/dbsetup/base.edgeql

@@ -0,0 +1,23 @@
+abstract type Named {
+    required name: str;
+}
+
+type UserGroup extending Named {
+    # many-to-many
+    multi link users: User;
+}
+
+type GameSession {
+    required num: int64;
+    # one-to-many
+    multi link players: User {
+        constraint exclusive;
+    };
+}
+
+type User extending Named;
+
+type Post {
+    required body: str;
+    required link author: User;
+}