alembic: fix in transaction recipe

- Fixes creation of sequence in transaction table,
  to make it able to drop

- Adds a utility function to drop the alembic_version table

Signed-off-by: Dinos Kousidis <konstantinos.kousidis@cern.ch>

invenio_db/alembic/dbdbc1b19cf2_create_transaction_table.py

@@ -25,6 +25,8 @@
 
 import sqlalchemy as sa
 from alembic import op
+from sqlalchemy.schema import Sequence, CreateSequence, \
+    DropSequence
 
 # revision identifiers, used by Alembic.
 revision = 'dbdbc1b19cf2'
@@ -40,10 +42,18 @@ def upgrade():
         sa.Column('issued_at', sa.DateTime(), nullable=True),
         sa.Column('id', sa.BigInteger(), nullable=False),
         sa.Column('remote_addr', sa.String(length=50), nullable=True),
-        sa.PrimaryKeyConstraint('id')
     )
+    op.create_primary_key('pk_transaction', 'transaction', ['id'])
+    # if op._proxy.migration_context.dialect.supports_sequences:
+    op.execute(CreateSequence(Sequence('transaction_id_seq')))
 
 
 def downgrade():
     """Downgrade database."""
     op.drop_table('transaction')
+    # if op._proxy.migration_context.dialect.supports_sequences:
+    #     try:
+    #         from psycopg2 import ProgrammingError
+    op.execute(DropSequence(Sequence('transaction_id_seq')))
+        # except (ProgrammingError, ImportError) as e:
+        #     pass

invenio_db/cli.py

@@ -32,14 +32,11 @@
 from click import _termui_impl
 from flask import current_app
 from flask.cli import with_appcontext
-from sqlalchemy.exc import InternalError
-from sqlalchemy.schema import DropTable
-from sqlalchemy.sql.ddl import SchemaDropper
 from sqlalchemy_utils.functions import create_database, database_exists, \
     drop_database
 from werkzeug.local import LocalProxy
 
-from .utils import create_alembic_version_table
+from .utils import create_alembic_version_table, drop_alembic_version_table
 
 _db = LocalProxy(lambda: current_app.extensions['sqlalchemy'].db)
 
@@ -86,19 +83,12 @@ def create(verbose):
 def drop(verbose):
     """Drop tables."""
     click.secho('Dropping all tables!', fg='red', bold=True)
-    schema_dropper = SchemaDropper(_db.engine.dialect, _db.engine.connect())
     with click.progressbar(reversed(_db.metadata.sorted_tables)) as bar:
         for table in bar:
             if verbose:
                 click.echo(' Dropping table {0}'.format(table))
-            try:
                 table.drop(bind=_db.engine, checkfirst=True)
-            except InternalError:
-                schema_dropper.connection.execute(DropTable(table))
-    if _db.engine.dialect.has_table(_db.engine, 'alembic_version'):
-        alembic_version = _db.Table('alembic_version', _db.metadata,
-                                    autoload_with=_db.engine)
-        alembic_version.drop(bind=_db.engine)
+        drop_alembic_version_table()
     click.secho('Dropped all tables!', fg='green')
 
 

invenio_db/utils.py

@@ -27,9 +27,12 @@
 
 from flask import current_app
 from sqlalchemy.engine import reflection
+from werkzeug.local import LocalProxy
 
 from .shared import db
 
+_db = LocalProxy(lambda: current_app.extensions['sqlalchemy'].db)
+
 
 def rebuild_encrypted_properties(old_key, model, properties):
     """Rebuild a model's EncryptedType properties when the SECRET_KEY is changed.
@@ -78,3 +81,11 @@ def create_alembic_version_table():
         alembic.migration_context._ensure_version_table()
         for head in alembic.script_directory.revision_map._real_heads:
             alembic.migration_context.stamp(alembic.script_directory, head)
+
+
+def drop_alembic_version_table():
+    """Drop alembic_version table."""
+    if _db.engine.dialect.has_table(_db.engine, 'alembic_version'):
+        alembic_version = _db.Table('alembic_version', _db.metadata,
+                                    autoload_with=_db.engine)
+        alembic_version.drop(bind=_db.engine)

tests/test_db.py

@@ -40,6 +40,7 @@
 
 from invenio_db import InvenioDB, shared
 from invenio_db.cli import db as db_cmd
+from invenio_db.utils import drop_alembic_version_table
 
 
 class MockEntryPoint(EntryPoint):
@@ -396,6 +397,12 @@ def test_db_create_alembic_upgrade(app, db):
                                    obj=script_info)
             assert result.exit_code == 0
             assert len(db.engine.table_names()) == 0
+
+            ext.alembic.upgrade()
+            db.drop_all()
+            drop_alembic_version_table()
+            assert len(db.engine.table_names()) == 0
+
         finally:
             drop_database(str(db.engine.url))
             remove_versioning(manager=ext.versioning_manager)