Don't auto-upgrade db for new Experiments [Resolves #695] (#698)

* Don't auto-upgrade db for new Experiments [Resolves #695]

To avoid the problem of time-consuming database upgrades happening when
we don't want them, the Experiment now:

1. Checks to see if the results_schema_versions table exists at all. if
it doesn't exist, upgrade. This is because means the results schema
should be clean in this case, and new users won't have to always run a
new thing when they first try Triage.
2. If it does exist, and the version number doesn't match what the
code's current HEAD revision is, throw an error. The error message is
customized to whether the database revision is a known revision to the
code (easy case, just upgrade if you have time) or not (you probably
upgraded on a different branch and need to go check out that branch to
downgrade).

src/tests/results_tests/test_upgrade_if_clean.py

@@ -0,0 +1,17 @@
+from triage.component import results_schema
+from alembic import command, script
+import pytest
+
+
+def test_upgrade_if_clean_upgrades_if_clean(db_engine):
+    results_schema.upgrade_if_clean(db_engine.url)
+    db_version = db_engine.execute("select version_num from results_schema_versions").scalar()
+    alembic_cfg = results_schema.alembic_config(db_engine.url)
+    assert db_version == script.ScriptDirectory.from_config(alembic_cfg).get_current_head()
+
+
+def test_upgrade_if_clean_does_not_upgrade_if_not_clean(db_engine):
+    command.upgrade(results_schema.alembic_config(dburl=db_engine.url), "head")
+    command.downgrade(results_schema.alembic_config(dburl=db_engine.url), "-1")
+    with pytest.raises(ValueError):
+        results_schema.upgrade_if_clean(db_engine.url)

src/triage/component/results_schema/__init__.py

@@ -1,9 +1,11 @@
 import os.path
 
+import logging
 from alembic.config import Config
+from alembic import script
 from alembic import command
-import yaml
-from sqlalchemy.engine.url import URL
+from triage import create_engine
+from triage.database_reflection import table_exists
 
 
 from .schema import (
@@ -88,6 +90,51 @@ def stamp_db(revision, dburl):
     command.stamp(alembic_config(dburl=dburl), revision)
 
 
+def upgrade_if_clean(dburl):
+    """Upgrade the database only if the results schema hasn't been created yet.
+
+    Raises: ValueError if the database results schema version does not equal the code's version
+    """
+    alembic_cfg = alembic_config(dburl)
+    engine = create_engine(dburl)
+    script_ = script.ScriptDirectory.from_config(alembic_cfg)
+    if not table_exists('results_schema_versions', engine):
+        logging.info("No results_schema_versions table exists, which means that this installation "
+                     "is fresh. Upgrading db.")
+        upgrade_db(dburl=dburl)
+        return
+    with engine.begin() as conn:
+        current_revision = conn.execute(
+            'select version_num from results_schema_versions limit 1'
+        ).scalar()
+        logging.info("Database's results schema version is %s", current_revision)
+        triage_head = script_.get_current_head()
+        logging.info("Code's results schema version is %s", triage_head)
+        database_is_ahead = not any(
+            migration.revision == current_revision
+            for migration in script_.walk_revisions()
+        )
+        if database_is_ahead:
+            raise ValueError(
+                f"Your database's results schema version, {current_revision}, is not a known "
+                "revision to this version of Triage. Usually, this happens if you use a branch "
+                "with a new results schema version and upgrade the database to that version. "
+                "To use this version of Triage, you will likely need to check out that branch "
+                f"and downgrade to {triage_head}",
+            )
+        elif current_revision != triage_head:
+            raise ValueError(
+                f"Your database's results schema revision, {current_revision}, is out of date "
+                "for this version of Triage. However, your database can be upgraded to this "
+                "revision. If you would like to upgrade your database from the console, and "
+                "you've installed Triage, you may execute `triage db upgrade`. "
+                "If the `triage` command is unavailable, (because you are running Triage directly "
+                " from a repository checkout), then `manage alembic upgrade head`. "
+                "The database changes may take a long time on a heavily populated database. "
+                "Otherwise, you can also downgrade your Triage version to match your database."
+            )
+
+
 def alembic_config(dburl):
     path = os.path.abspath(__file__)
     dir_path = os.path.dirname(path)

src/triage/experiments/base.py

@@ -127,7 +127,7 @@ def __init__(
         self.save_predictions = save_predictions
         self.skip_validation = skip_validation
         self.db_engine = db_engine
-        results_schema.upgrade_db(db_engine=self.db_engine)
+        results_schema.upgrade_if_clean(dburl=self.db_engine.url)
 
         self.features_schema_name = "features"
         self.materialize_subquery_fromobjs = materialize_subquery_fromobjs