Repository: jupyterhub/jupyterhub

Issue:
Issues using upgrade-db on a fresh database
Database migration works only if you start on an old version of the database. But, consider the case where I create a JupyterHub database today, and then a future release updates the database. Running `jupyterhub upgrade-db` then would fail, because the database existing today doesn't have the `alembic_version` table specifying what the current version is. So, alembic will assume the version is the old version, and try to insert columns that already exist:

```
$ jupyterhub
...
$ jupyterhub upgrade-db
INFO  [alembic.runtime.migration] Context impl SQLiteImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
INFO  [alembic.runtime.migration] Running upgrade  -> 19c0846f6344, base revision for 0.5
INFO  [alembic.runtime.migration] Running upgrade 19c0846f6344 -> eeb276e51423, auth_state
Traceback (most recent call last):
  File "/usr/local/lib/python3.5/dist-packages/sqlalchemy/engine/base.py", line 1182, in _execute_context
    context)
  File "/usr/local/lib/python3.5/dist-packages/sqlalchemy/engine/default.py", line 470, in do_execute
    cursor.execute(statement, parameters)
sqlite3.OperationalError: duplicate column name: auth_state

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/local/bin/alembic", line 11, in <module>
    sys.exit(main())
  File "/usr/local/lib/python3.5/dist-packages/alembic/config.py", line 479, in main
    CommandLine(prog=prog).main(argv=argv)
  File "/usr/local/lib/python3.5/dist-packages/alembic/config.py", line 473, in main
    self.run_cmd(cfg, options)
  File "/usr/local/lib/python3.5/dist-packages/alembic/config.py", line 456, in run_cmd
    **dict((k, getattr(options, k, None)) for k in kwarg)
  File "/usr/local/lib/python3.5/dist-packages/alembic/command.py", line 254, in upgrade
    script.run_env()
  File "/usr/local/lib/python3.5/dist-packages/alembic/script/base.py", line 416, in run_env
    util.load_python_file(self.dir, 'env.py')
  File "/usr/local/lib/python3.5/dist-packages/alembic/util/pyfiles.py", line 93, in load_python_file
    module = load_module_py(module_id, path)
  File "/usr/local/lib/python3.5/dist-packages/alembic/util/compat.py", line 64, in load_module_py
    module_id, path).load_module(module_id)
  File "<frozen importlib._bootstrap_external>", line 388, in _check_name_wrapper
  File "<frozen importlib._bootstrap_external>", line 809, in load_module
  File "<frozen importlib._bootstrap_external>", line 668, in load_module
  File "<frozen importlib._bootstrap>", line 268, in _load_module_shim
  File "<frozen importlib._bootstrap>", line 693, in _load
  File "<frozen importlib._bootstrap>", line 673, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 665, in exec_module
  File "<frozen importlib._bootstrap>", line 222, in _call_with_frames_removed
  File "/usr/local/lib/python3.5/dist-packages/jupyterhub/alembic/env.py", line 70, in <module>
    run_migrations_online()
  File "/usr/local/lib/python3.5/dist-packages/jupyterhub/alembic/env.py", line 65, in run_migrations_online
    context.run_migrations()
  File "<string>", line 8, in run_migrations
  File "/usr/local/lib/python3.5/dist-packages/alembic/runtime/environment.py", line 817, in run_migrations
    self.get_context().run_migrations(**kw)
  File "/usr/local/lib/python3.5/dist-packages/alembic/runtime/migration.py", line 323, in run_migrations
    step.migration_fn(**kw)
  File "/usr/local/lib/python3.5/dist-packages/jupyterhub/alembic/versions/eeb276e51423_auth_state.py", line 21, in upgrade
    op.add_column('users', sa.Column('auth_state', JSONDict))
  File "<string>", line 8, in add_column
  File "<string>", line 3, in add_column
  File "/usr/local/lib/python3.5/dist-packages/alembic/operations/ops.py", line 1551, in add_column
    return operations.invoke(op)
  File "/usr/local/lib/python3.5/dist-packages/alembic/operations/base.py", line 318, in invoke
    return fn(self, operation)
  File "/usr/local/lib/python3.5/dist-packages/alembic/operations/toimpl.py", line 123, in add_column
    schema=schema
  File "/usr/local/lib/python3.5/dist-packages/alembic/ddl/impl.py", line 172, in add_column
    self._exec(base.AddColumn(table_name, column, schema=schema))
  File "/usr/local/lib/python3.5/dist-packages/alembic/ddl/impl.py", line 118, in _exec
    return conn.execute(construct, *multiparams, **params)
  File "/usr/local/lib/python3.5/dist-packages/sqlalchemy/engine/base.py", line 945, in execute
    return meth(self, multiparams, params)
  File "/usr/local/lib/python3.5/dist-packages/sqlalchemy/sql/ddl.py", line 68, in _execute_on_connection
    return connection._execute_ddl(self, multiparams, params)
  File "/usr/local/lib/python3.5/dist-packages/sqlalchemy/engine/base.py", line 1002, in _execute_ddl
    compiled
  File "/usr/local/lib/python3.5/dist-packages/sqlalchemy/engine/base.py", line 1189, in _execute_context
    context)
  File "/usr/local/lib/python3.5/dist-packages/sqlalchemy/engine/base.py", line 1402, in _handle_dbapi_exception
    exc_info
  File "/usr/local/lib/python3.5/dist-packages/sqlalchemy/util/compat.py", line 203, in raise_from_cause
    reraise(type(exception), exception, tb=exc_tb, cause=cause)
  File "/usr/local/lib/python3.5/dist-packages/sqlalchemy/util/compat.py", line 186, in reraise
    raise value.with_traceback(tb)
  File "/usr/local/lib/python3.5/dist-packages/sqlalchemy/engine/base.py", line 1182, in _execute_context
    context)
  File "/usr/local/lib/python3.5/dist-packages/sqlalchemy/engine/default.py", line 470, in do_execute
    cursor.execute(statement, parameters)
sqlalchemy.exc.OperationalError: (sqlite3.OperationalError) duplicate column name: auth_state [SQL: 'ALTER TABLE users ADD COLUMN auth_state TEXT']
[E 2017-06-02 19:55:16.017 JupyterHub app:1527]
    Traceback (most recent call last):
      File "/usr/local/lib/python3.5/dist-packages/jupyterhub/app.py", line 1525, in launch_instance_async
        yield self.start()
      File "/usr/local/lib/python3.5/dist-packages/jupyterhub/app.py", line 1436, in start
        self.subapp.start()
      File "/usr/local/lib/python3.5/dist-packages/jupyterhub/app.py", line 185, in start
        dbutil.upgrade(hub.db_url)
      File "/usr/local/lib/python3.5/dist-packages/jupyterhub/dbutil.py", line 80, in upgrade
        ['alembic', '-c', alembic_ini, 'upgrade', revision]
      File "/usr/lib/python3.5/subprocess.py", line 581, in check_call
        raise CalledProcessError(retcode, cmd)
    subprocess.CalledProcessError: Command '['alembic', '-c', '/tmp/tmpkqvnuq0a/alembic.ini', 'upgrade', 'head']' returned non-zero exit status 1
```

I got around this in nbgrader by creating the `alembic_version` table the first time the database is initialized:

https://github.com/jupyter/nbgrader/blob/master/nbgrader/api.py#L1045


Discussion:
Good catch @jhamrick.
I'll try to get this one this week. I think we need to:

1. call ansible.tag on startup creating a new db, to prevent this in the future. But be careful to only do this if no tag is not already set
2. handle this bug: detect existing db without ansible tag data and trigger tag for 0.7 version

Resolving patch:
diff --git a/jupyterhub/alembic/versions/3ec6993fe20c_encrypted_auth_state.py b/jupyterhub/alembic/versions/3ec6993fe20c_encrypted_auth_state.py
--- a/jupyterhub/alembic/versions/3ec6993fe20c_encrypted_auth_state.py
+++ b/jupyterhub/alembic/versions/3ec6993fe20c_encrypted_auth_state.py
@@ -1,4 +1,11 @@
-"""encrypted_auth_state
+"""0.8 changes
+
+- encrypted auth_state
+- remove proxy/hub data from db
+
+OAuth data was also added in this revision,
+but no migration to do because they are entirely new tables,
+which will be created on launch.
 
 Revision ID: 3ec6993fe20c
 Revises: af4cbdb2d13c
@@ -12,7 +19,9 @@
 branch_labels = None
 depends_on = None
 
-import warnings
+import logging
+
+logger = logging.getLogger('alembic')
 
 from alembic import op
 import sqlalchemy as sa
@@ -20,15 +29,36 @@
 
 
 def upgrade():
+    # proxy/table info is no longer in the database
+    op.drop_table('proxies')
+    op.drop_table('hubs')
+
+    # drop some columns no longer in use
     try:
         op.drop_column('users', 'auth_state')
-    except sa.exc.OperationalError as e:
-        # sqlite3 can't drop columns
-        warnings.warn("Failed to drop column: %s" % e)
+        op.drop_column('users', '_server_id')
+    except sa.exc.OperationalError:
+        # this won't be a problem moving forward, but downgrade will fail
+        if op.get_context().dialect.name == 'sqlite':
+            logger.warning("sqlite cannot drop columns. Leaving unused old columns in place.")
+        else:
+            raise
+
     op.add_column('users', sa.Column('encrypted_auth_state', sa.types.LargeBinary))
 
 
 def downgrade():
+    # drop all the new tables
+    engine = op.get_bind().engine
+    for table in ('oauth_clients',
+                  'oauth_codes',
+                  'oauth_access_tokens',
+                  'spawners'):
+        if engine.has_table(table):
+            op.drop_table(table)
+
     op.drop_column('users', 'encrypted_auth_state')
+
     op.add_column('users', sa.Column('auth_state', JSONDict))
-    
+    op.add_column('users',  sa.Column('_server_id', sa.Integer, sa.ForeignKey('servers.id')))
+
diff --git a/jupyterhub/app.py b/jupyterhub/app.py
--- a/jupyterhub/app.py
+++ b/jupyterhub/app.py
@@ -185,6 +185,7 @@ def _backup_db_file(self, db_file):
     def start(self):
         hub = JupyterHub(parent=self)
         hub.load_config_file(hub.config_file)
+        self.log = hub.log
         if (hub.db_url.startswith('sqlite:///')):
             db_file = hub.db_url.split(':///', 1)[1]
             self._backup_db_file(db_file)
@@ -862,6 +863,8 @@ def init_db(self):
                 "to upgrade your JupyterHub database schema",
             ]))
             self.exit(1)
+        except orm.DatabaseSchemaMismatch as e:
+            self.exit(e)
 
     def init_hub(self):
         """Load the Hub config into the database"""
diff --git a/jupyterhub/orm.py b/jupyterhub/orm.py
--- a/jupyterhub/orm.py
+++ b/jupyterhub/orm.py
@@ -7,6 +7,9 @@
 import enum
 import json
 
+import alembic.config
+import alembic.command
+from alembic.script import ScriptDirectory
 from tornado.log import app_log
 
 from sqlalchemy.types import TypeDecorator, TEXT, LargeBinary
@@ -21,6 +24,7 @@
 from sqlalchemy.sql.expression import bindparam
 from sqlalchemy import create_engine, Table
 
+from .dbutil import _temp_alembic_ini
 from .utils import (
     random_port,
     new_token, hash_token, compare_token,
@@ -431,6 +435,73 @@ class OAuthClient(Base):
     redirect_uri = Column(Unicode(1023))
 
 
+class DatabaseSchemaMismatch(Exception):
+    """Exception raised when the database schema version does not match
+
+    the current version of JupyterHub.
+    """
+
+def check_db_revision(engine):
+    """Check the JupyterHub database revision
+
+    After calling this function, an alembic tag is guaranteed to be stored in the db.
+
+    - Checks the alembic tag and raises a ValueError if it's not the current revision
+    - If no tag is stored (Bug in Hub prior to 0.8),
+      guess revision based on db contents and tag the revision.
+    - Empty databases are tagged with the current revision
+    """
+    # Check database schema version
+    current_table_names = set(engine.table_names())
+    my_table_names = set(Base.metadata.tables.keys())
+
+    with _temp_alembic_ini(engine.url) as ini:
+        cfg = alembic.config.Config(ini)
+        scripts = ScriptDirectory.from_config(cfg)
+        head = scripts.get_heads()[0]
+        base = scripts.get_base()
+
+        if not my_table_names.intersection(current_table_names):
+            # no tables have been created, stamp with current revision
+            app_log.debug("Stamping empty database with alembic revision %s", head)
+            alembic.command.stamp(cfg, head)
+            return
+
+        if 'alembic_version' not in current_table_names:
+            # Has not been tagged or upgraded before.
+            # we didn't start tagging revisions correctly except during `upgrade-db`
+            # until 0.8
+            # This should only occur for databases created prior to JupyterHub 0.8
+            msg_t = "Database schema version not found, guessing that JupyterHub %s created this database."
+            if 'spawners' in current_table_names:
+                # 0.8
+                app_log.warning(msg_t, '0.8.dev')
+                rev = head
+            elif 'services' in current_table_names:
+                # services is present, tag for 0.7
+                app_log.warning(msg_t, '0.7.x')
+                rev = 'af4cbdb2d13c'
+            else:
+                # it's old, mark as first revision
+                app_log.warning(msg_t, '0.6 or earlier')
+                rev = base
+            app_log.debug("Stamping database schema version %s", rev)
+            alembic.command.stamp(cfg, rev)
+
+    # check database schema version
+    # it should always be defined at this point
+    alembic_revision = engine.execute('SELECT version_num FROM alembic_version').first()[0]
+    if alembic_revision == head:
+        app_log.debug("database schema version found: %s", alembic_revision)
+        pass
+    else:
+        raise DatabaseSchemaMismatch("Found database schema version {found} != {head}. "
+        "Backup your database and run `jupyterhub upgrade-db`"
+        " to upgrade to the latest schema.".format(
+            found=alembic_revision,
+            head=head,
+        ))
+
 def new_session_factory(url="sqlite:///:memory:", reset=False, **kwargs):
     """Create a new session at url"""
     if url.startswith('sqlite'):
@@ -446,6 +517,9 @@ def new_session_factory(url="sqlite:///:memory:", reset=False, **kwargs):
     engine = create_engine(url, **kwargs)
     if reset:
         Base.metadata.drop_all(engine)
+
+    # check the db revision (will raise, pointing to `upgrade-db` if version doesn't match)
+    check_db_revision(engine)
     Base.metadata.create_all(engine)
 
     session_factory = sessionmaker(bind=engine)


Test patch:
diff --git a/jupyterhub/tests/test_db.py b/jupyterhub/tests/test_db.py
--- a/jupyterhub/tests/test_db.py
+++ b/jupyterhub/tests/test_db.py
@@ -2,8 +2,6 @@
 import os
 import shutil
 
-from sqlalchemy.exc import OperationalError
-
 import pytest
 from pytest import raises
 
@@ -32,7 +30,7 @@ def test_upgrade_entrypoint(tmpdir):
     tmpdir.chdir()
     tokenapp = NewToken()
     tokenapp.initialize(['kaylee'])
-    with raises(OperationalError):
+    with raises(SystemExit):
         tokenapp.start()
 
     sqlite_files = glob(os.path.join(str(tmpdir), 'jupyterhub.sqlite*'))