Print table schema in DB initialization test

.github/workflows/master.yml

@@ -168,10 +168,11 @@ jobs:
         # Run tests
         cd tests/db
         docker-compose build
-        docker-compose run mlflow-postgres python log.py
-        docker-compose run mlflow-mysql python log.py
+        docker-compose run mlflow-sqlite python run_checks.py --schema-output schemas/sqlite.sql
+        docker-compose run mlflow-postgres python run_checks.py --schema-output schemas/postgres.sql
+        docker-compose run mlflow-mysql python run_checks.py --schema-output schemas/mysql.sql
         docker-compose run mlflow-mssql ./init-mssql-db.sh
-        docker-compose run mlflow-mssql python log.py
+        docker-compose run mlflow-mssql python run_checks.py --schema-output schemas/mssql.sql
 
         # Clean up
         docker-compose down --rmi all --volumes

tests/db/.dockerignore

@@ -1,6 +1,5 @@
 **
 
 !dist/*.whl
-!log.py
 !init-mssql-db.sh
 !init-mssql-db.sql

tests/db/.gitignore

@@ -0,0 +1,2 @@
+schemas
+mlflowdb

tests/db/Dockerfile

@@ -7,5 +7,3 @@ COPY dist ./dist
 RUN pip install dist/*.whl
 RUN pip install psycopg2 pymysql mysqlclient
 RUN pip list
-
-COPY log.py .

tests/db/Dockerfile.mssql

@@ -19,9 +19,3 @@ RUN apt-get update && ACCEPT_EULA=Y apt-get install -y mssql-tools unixodbc-dev
 RUN pip install dist/*.whl
 RUN pip install pyodbc
 RUN pip list
-
-COPY log.py .
-COPY init-mssql-db.sh .
-COPY init-mssql-db.sql .
-
-RUN chmod +x init-mssql-db.sh

tests/db/docker-compose.yml

@@ -13,6 +13,8 @@ services:
       - postgres
     build:
       context: .
+    volumes:
+      - .:/tmp/mlflow
     environment:
       MLFLOW_TRACKING_URI: postgresql://mlflowuser:mlflowpassword@postgres:5432/mlflowdb
 
@@ -30,6 +32,8 @@ services:
       - mysql
     build:
       context: .
+    volumes:
+      - .:/tmp/mlflow
     environment:
       MLFLOW_TRACKING_URI: mysql://mlflowuser:mlflowpassword@mysql:3306/mlflowdb
 
@@ -46,5 +50,17 @@ services:
     build:
       context: .
       dockerfile: Dockerfile.mssql
+    volumes:
+      - .:/tmp/mlflow
     environment:
       MLFLOW_TRACKING_URI: mssql+pyodbc://mlflowuser:Mlfl*wpassword1@mssql/mlflowdb?driver=ODBC+Driver+17+for+SQL+Server
+
+  mlflow-sqlite:
+    depends_on:
+      - postgres
+    build:
+      context: .
+    volumes:
+      - .:/tmp/mlflow
+    environment:
+      MLFLOW_TRACKING_URI: sqlite:////tmp/mlflow/mlflowdb

tests/db/run_checks.py

@@ -0,0 +1,67 @@
+import os
+import argparse
+
+import sqlalchemy
+from sqlalchemy.schema import MetaData, CreateTable
+
+import mlflow
+from mlflow.tracking._tracking_service.utils import _TRACKING_URI_ENV_VAR
+
+
+class MockModel(mlflow.pyfunc.PythonModel):
+    def load_context(self, context):
+        pass
+
+    def predict(self, context, model_input):
+        pass
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--schema-output", required=True, help="Output path of DB schema")
+    return parser.parse_args()
+
+
+def run_logging_operations():
+    with mlflow.start_run() as run:
+        print("Tracking URI:", mlflow.get_tracking_uri())
+        mlflow.log_param("p", "param")
+        mlflow.log_metric("m", 1.0)
+        mlflow.set_tag("t", "tag")
+        mlflow.pyfunc.log_model(
+            artifact_path="model",
+            python_model=MockModel(),
+            registered_model_name="mock",
+        )
+        print(mlflow.get_run(run.info.run_id))
+
+
+def get_db_schema():
+    engine = sqlalchemy.create_engine(mlflow.get_tracking_uri())
+    created_tables_metadata = MetaData(bind=engine)
+    created_tables_metadata.reflect()
+    # 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 += list(map(str.rstrip, str(CreateTable(ti)).splitlines()))
+    return "\n".join(lines)
+
+
+def main():
+    assert _TRACKING_URI_ENV_VAR in os.environ
+
+    run_logging_operations()
+    schema = get_db_schema()
+    title = "Schema"
+    print("=" * 10, title, "=" * 10)
+    print(schema)
+    print("=" * (20 + 2 + len(title)))
+    args = parse_args()
+    os.makedirs(os.path.dirname(args.schema_output), exist_ok=True)
+    with open(args.schema_output, "w") as f:
+        f.write(schema)
+
+
+if __name__ == "__main__":
+    main()