mirror of https://github.com/apache/superset.git
fix: create fk model in benchmark script (#15804)
This commit is contained in:
parent
c9dad05f2b
commit
5cc95bb378
|
@ -94,11 +94,22 @@ def find_models(module: ModuleType) -> List[Type[Model]]:
|
||||||
engine = create_engine(sqlalchemy_uri)
|
engine = create_engine(sqlalchemy_uri)
|
||||||
Base = automap_base()
|
Base = automap_base()
|
||||||
Base.prepare(engine, reflect=True)
|
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 = getattr(Base.classes, table)
|
||||||
model.__tablename__ = table
|
model.__tablename__ = table
|
||||||
models.append(model)
|
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
|
# sort topologically so we can create entities in order and
|
||||||
# maintain relationships (eg, create a database before creating
|
# maintain relationships (eg, create a database before creating
|
||||||
# a slice)
|
# a slice)
|
||||||
|
@ -108,7 +119,8 @@ def find_models(module: ModuleType) -> List[Type[Model]]:
|
||||||
dependent_tables: List[str] = []
|
dependent_tables: List[str] = []
|
||||||
for column in inspector.columns.values():
|
for column in inspector.columns.values():
|
||||||
for foreign_key in column.foreign_keys:
|
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)
|
sorter.add(model.__tablename__, *dependent_tables)
|
||||||
order = list(sorter.static_order())
|
order = list(sorter.static_order())
|
||||||
models.sort(key=lambda model: order.index(model.__tablename__))
|
models.sort(key=lambda model: order.index(model.__tablename__))
|
||||||
|
|
Loading…
Reference in New Issue