fix: create fk model in benchmark script (#15804)

This commit is contained in:
Beto Dealmeida 2021-07-20 14:04:19 -07:00 committed by GitHub
parent c9dad05f2b
commit 5cc95bb378
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 14 additions and 2 deletions

View File

@ -94,11 +94,22 @@ def find_models(module: ModuleType) -> List[Type[Model]]:
engine = create_engine(sqlalchemy_uri)
Base = automap_base()
Base.prepare(engine, reflect=True)
for table in tables:
seen = set()
while tables:
table = tables.pop()
seen.add(table)
model = getattr(Base.classes, table)
model.__tablename__ = table
models.append(model)
# add other models referenced in foreign keys
inspector = inspect(model)
for column in inspector.columns.values():
for foreign_key in column.foreign_keys:
table = foreign_key.column.table.name
if table not in seen:
tables.add(table)
# sort topologically so we can create entities in order and
# maintain relationships (eg, create a database before creating
# a slice)
@ -108,7 +119,8 @@ def find_models(module: ModuleType) -> List[Type[Model]]:
dependent_tables: List[str] = []
for column in inspector.columns.values():
for foreign_key in column.foreign_keys:
dependent_tables.append(foreign_key.target_fullname.split(".")[0])
if foreign_key.column.table.name != model.__tablename__:
dependent_tables.append(foreign_key.column.table.name)
sorter.add(model.__tablename__, *dependent_tables)
order = list(sorter.static_order())
models.sort(key=lambda model: order.index(model.__tablename__))