Skip to content

Commit

Permalink
Add a way to export SQLAlchemy models from Gel (EdgeDB).
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
vpetrovykh committed Nov 14, 2024
1 parent 96161a2 commit be2af63
Show file tree
Hide file tree
Showing 8 changed files with 735 additions and 2 deletions.
2 changes: 1 addition & 1 deletion .flake8
Original file line number Diff line number Diff line change
@@ -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
30 changes: 29 additions & 1 deletion edgedb/_testbase.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand Down Expand Up @@ -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


Expand Down
Empty file added edgedb/gelalchemy/__init__.py
Empty file.
271 changes: 271 additions & 0 deletions edgedb/gelalchemy/models.py
Original file line number Diff line number Diff line change
@@ -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')
44 changes: 44 additions & 0 deletions tests/dbsetup/base.edgeql
Original file line number Diff line number Diff line change
@@ -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*',
};
23 changes: 23 additions & 0 deletions tests/dbsetup/base.esdl
Original file line number Diff line number Diff line change
@@ -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;
}
Empty file added tests/generated/__init__.py
Empty file.
Loading

0 comments on commit be2af63

Please sign in to comment.