Merge pull request #470 from opensafely-core/migrate-cmd-fix

fix: fix migrate command and improve migration logging

jobrunner/cli/migrate.py

@@ -2,13 +2,14 @@
 Run any pending database migrations
 """
 import argparse
+import sys
 from pathlib import Path
 
 from jobrunner import config
 from jobrunner.lib import database, log_utils
 
 
-def run():
+def run(argv):
     log_utils.configure_logging()
     parser = argparse.ArgumentParser(description=__doc__.partition("\n\n")[0])
     parser.add_argument(
@@ -17,9 +18,9 @@ def run():
         default=config.DATABASE_FILE,
         help="db file to migrate (defaults to configured db)",
     )
-    args = parser.parse_args()
-    database.ensure_db(args.dbpath)
+    args = parser.parse_args(argv)
+    database.ensure_db(args.dbpath, verbose=True)
 
 
 if __name__ == "__main__":
-    run()
+    run(sys.argv[1:])

jobrunner/lib/database.py

@@ -10,6 +10,7 @@
 """
 import dataclasses
 import json
+import logging
 import sqlite3
 import threading
 from enum import Enum
@@ -18,6 +19,8 @@
 from jobrunner import config
 
 
+log = logging.getLogger(__name__)
+
 CONNECTION_CACHE = threading.local()
 TABLES = {}
 MIGRATIONS = {}
@@ -214,7 +217,7 @@ def ensure_valid_db(filename=None, migrations=MIGRATIONS):
         )
 
 
-def ensure_db(filename=None, migrations=MIGRATIONS):
+def ensure_db(filename=None, migrations=MIGRATIONS, verbose=False):
     """Ensure db is created and up to date with migrations
 
     Will create new tables, or migrate the exisiting ones as needed.
@@ -230,14 +233,14 @@ def ensure_db(filename=None, migrations=MIGRATIONS):
     conn = get_connection(filename)
 
     if db_exists:
-        migrate_db(conn, migrations)
+        migrate_db(conn, migrations, verbose=verbose)
     else:  # new db
         for table in TABLES.values():
             create_table(conn, table)
         # set migration level to highest migration version
         conn.execute(f"PRAGMA user_version={max(migrations, default=0)}")
-        # print("created new db at {filename}")
-
+        if verbose:
+            log.info(f"created new db at {filename}")
     return conn
 
 
@@ -254,7 +257,10 @@ def create_table(conn, cls):
 """
 
 
-def migrate_db(conn, migrations=None):
+def migrate_db(conn, migrations=None, verbose=False):
+
+    # we store migrations in models, so make sure this has been imported to collect them
+    import jobrunner.models  # noqa: F401
 
     current_version = conn.execute("PRAGMA user_version").fetchone()[0]
     applied = []
@@ -264,9 +270,11 @@ def migrate_db(conn, migrations=None):
             transaction_sql = MIGRATION_SQL.format(sql=sql, version=version)
             conn.executescript(transaction_sql)
             applied.append(version)
-            print(f"Applied {migration} as migration {version}:\n{sql}")
+            if verbose:
+                log.info(f"Applied migration {version}:\n{sql}")
         else:
-            print(f"Skipping {migration} as already applied")
+            if verbose:
+                log.info(f"Skipping migration {version} as already applied")
 
     return applied
 

tests/lib/test_database.py

@@ -1,3 +1,4 @@
+import logging
 import sqlite3
 
 import pytest
@@ -87,6 +88,22 @@ def test_ensure_db_new_db(tmp_path):
     assert conn.execute("PRAGMA user_version").fetchone()[0] == 1
 
 
+def test_ensure_db_verbose(tmp_path, caplog):
+    caplog.set_level(logging.INFO)
+    db = tmp_path / "db.sqlite"
+
+    ensure_db(db, {1: "should not run"}, verbose=True)
+    assert "created new db" in caplog.records[-1].msg
+
+    ensure_db(
+        db,
+        {1: "should not run", 2: "ALTER TABLE job ADD COLUM test TEXT"},
+        verbose=True,
+    )
+    assert "Skipping migration 1" in caplog.records[-2].msg
+    assert "Applied migration 2" in caplog.records[-1].msg
+
+
 def test_ensure_db_new_db_memory():
     db = "file:test?mode=memory&cached=shared"
     conn = ensure_db(db, {1: "should not run"})