wehub-resource-sync
40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
"""Script that generates a dump of the MLflow tracking database schema"""
|
|
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
|
|
import sqlalchemy
|
|
from sqlalchemy.schema import CreateTable, MetaData
|
|
|
|
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
|
|
|
|
|
def dump_db_schema(db_url, dst_file):
|
|
engine = sqlalchemy.create_engine(db_url)
|
|
created_tables_metadata = MetaData()
|
|
created_tables_metadata.reflect(bind=engine)
|
|
# Write out table schema as described in
|
|
# https://docs.sqlalchemy.org/en/13/faq/
|
|
# metadata_schema.html#how-can-i-get-the-create-table-drop-table-output-as-a-string
|
|
lines = []
|
|
for ti in created_tables_metadata.sorted_tables:
|
|
lines.extend(line.rstrip() + "\n" for line in str(CreateTable(ti)).splitlines())
|
|
schema = "".join(lines)
|
|
with open(dst_file, "w") as handle:
|
|
handle.write(schema)
|
|
|
|
|
|
def dump_sqlalchemy_store_schema(dst_file):
|
|
with tempfile.TemporaryDirectory() as db_tmpdir:
|
|
path = os.path.join(db_tmpdir, "db_file")
|
|
db_url = f"sqlite:///{path}"
|
|
SqlAlchemyStore(db_url, db_tmpdir)
|
|
dump_db_schema(db_url, dst_file)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) != 2:
|
|
sys.exit(1)
|
|
dump_sqlalchemy_store_schema(sys.argv[1])
|