Merge pull request #50 from CAVEconnectome/improved_migration_scripts

feat: use standard template

dynamicannotationdb/migration/alembic/script.py.mako

@@ -7,6 +7,7 @@ Create Date: ${create_date}
 """
 from alembic import op
 import sqlalchemy as sa
+from sqlalchemy.engine import reflection
 
 ${imports if imports else ""}
 
@@ -17,7 +18,18 @@ branch_labels = ${repr(branch_labels)}
 depends_on = ${repr(depends_on)}
 
 
+def get_tables(connection):
+    inspector = reflection.Inspector.from_engine(connection)
+    return inspector.get_table_names()
+
+
+def _table_has_column(connection, table, column):
+    insp = reflection.Inspector.from_engine(connection)
+    return any(column in col["name"] for col in insp.get_columns(table))
+
+
 def upgrade():
+    connection = op.get_bind()
     ${upgrades if upgrades else "pass"}
 
 

dynamicannotationdb/migration/alembic/versions/309cf493a1e2_adding_warning_field.py

@@ -7,35 +7,34 @@
 """
 from alembic import op
 import sqlalchemy as sa
-
-from sqlalchemy.dialects import postgresql
-from sqlalchemy import engine_from_config
 from sqlalchemy.engine import reflection
 
 # revision identifiers, used by Alembic.
-revision = '309cf493a1e2'
-down_revision = '8fdc843fc202'
+revision = "309cf493a1e2"
+down_revision = "8fdc843fc202"
 branch_labels = None
 depends_on = None
 
 
-def _table_has_column(table, column):
-    config = op.get_context().config
-    engine = engine_from_config(
-        config.get_section(config.config_ini_section), prefix="sqlalchemy."
-    )
-
-    insp = reflection.Inspector.from_engine(engine)
-    return any(column in col["name"] for col in insp.get_columns(table))
+def get_tables(connection):
+    inspector = reflection.Inspector.from_engine(connection)
+    return inspector.get_table_names()
 
 
+def _table_has_column(connection, table, column):
+    insp = reflection.Inspector.from_engine(connection)
+    return any(column in col["name"] for col in insp.get_columns(table))
 
 def upgrade():
-    if not _table_has_column('annotation_table_metadata', 'notice_text'):
-        op.add_column('annotation_table_metadata', sa.Column('notice_text', sa.Text(), nullable=True))
+    connection = op.get_bind()
+    if not _table_has_column(connection, "annotation_table_metadata", "notice_text"):
+        op.add_column(
+            "annotation_table_metadata",
+            sa.Column("notice_text", sa.Text(), nullable=True),
+        )
 
 
 def downgrade():
     # ### commands auto generated by Alembic - please adjust! ###
-    op.drop_column('annotation_table_metadata', 'notice_text')
+    op.drop_column("annotation_table_metadata", "notice_text")
     # ### end Alembic commands ###

dynamicannotationdb/migration/alembic/versions/5a1d7c0ad006_add_status_column.py

@@ -7,31 +7,31 @@
 """
 from alembic import op
 import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql
-from sqlalchemy import engine_from_config
 from sqlalchemy.engine import reflection
+from sqlalchemy.dialects import postgresql
 
 # revision identifiers, used by Alembic.
 revision = "5a1d7c0ad006"
 down_revision = "7c79eff751b4"
 branch_labels = None
 depends_on = None
 
-def _table_has_column(table, column):
-    config = op.get_context().config
-    engine = engine_from_config(
-        config.get_section(config.config_ini_section), prefix="sqlalchemy."
-    )
+def get_tables(connection):
+    inspector = reflection.Inspector.from_engine(connection)
+    return inspector.get_table_names()
+
 
-    insp = reflection.Inspector.from_engine(engine)
+def _table_has_column(connection, table, column):
+    insp = reflection.Inspector.from_engine(connection)
     return any(column in col["name"] for col in insp.get_columns(table))
 
 def upgrade():
+    connection = op.get_bind()
     status_enum = postgresql.ENUM(
         "AVAILABLE", "RUNNING", "FAILED", "EXPIRED", name="version_status"
     )
     status_enum.create(op.get_bind())
-    if not _table_has_column("analysisversion", "status"):
+    if not _table_has_column(connection, "analysisversion", "status"):
         op.add_column(
             "analysisversion",
             sa.Column(

dynamicannotationdb/migration/alembic/versions/6e7f580ff680_add_error_msg.py

@@ -7,9 +7,6 @@
 """
 from alembic import op
 import sqlalchemy as sa
-
-from sqlalchemy.dialects import postgresql
-from sqlalchemy import engine_from_config
 from sqlalchemy.engine import reflection
 
 # revision identifiers, used by Alembic.
@@ -19,17 +16,19 @@
 depends_on = None
 
 
-def _table_has_column(table, column):
-    config = op.get_context().config
-    engine = engine_from_config(
-        config.get_section(config.config_ini_section), prefix="sqlalchemy."
-    )
+def get_tables(connection):
+    inspector = reflection.Inspector.from_engine(connection)
+    return inspector.get_table_names()
+
 
-    insp = reflection.Inspector.from_engine(engine)
+def _table_has_column(connection, table, column):
+    insp = reflection.Inspector.from_engine(connection)
     return any(column in col["name"] for col in insp.get_columns(table))
 
 def upgrade():
-    if not _table_has_column("version_error", "exception"):
+    connection = op.get_bind()
+
+    if not _table_has_column(connection, "version_error", "exception"):
         op.add_column('version_error', sa.Column('exception', sa.String(), nullable=True))
 
 

dynamicannotationdb/migration/alembic/versions/7c79eff751b4_add_parent_version_column.py

@@ -7,8 +7,6 @@
 """
 from alembic import op
 import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql
-from sqlalchemy import engine_from_config
 from sqlalchemy.engine import reflection
 
 # revision identifiers, used by Alembic.
@@ -18,18 +16,18 @@
 depends_on = "ef5c2d7f96d8"
 
 
-def _table_has_column(table, column):
-    config = op.get_context().config
-    engine = engine_from_config(
-        config.get_section(config.config_ini_section), prefix="sqlalchemy."
-    )
+def get_tables(connection):
+    inspector = reflection.Inspector.from_engine(connection)
+    return inspector.get_table_names()
 
-    insp = reflection.Inspector.from_engine(engine)
-    return any(column in col["name"] for col in insp.get_columns(table))
 
+def _table_has_column(connection, table, column):
+    insp = reflection.Inspector.from_engine(connection)
+    return any(column in col["name"] for col in insp.get_columns(table))
 
 def upgrade():
-    if not _table_has_column("analysisversion", "parent_version"):  
+    connection = op.get_bind()
+    if not _table_has_column(connection, "analysisversion", "parent_version"):  
 
         with op.batch_alter_table("analysisversion", schema=None) as batch_op:
             op.add_column(

dynamicannotationdb/migration/alembic/versions/814d72d74e3b_add_version_error_table.py

@@ -7,28 +7,28 @@
 """
 from alembic import op
 import sqlalchemy as sa
-
-from sqlalchemy.dialects import postgresql
-from sqlalchemy import engine_from_config
 from sqlalchemy.engine import reflection
+from sqlalchemy.dialects import postgresql
 
 # revision identifiers, used by Alembic.
 revision = "814d72d74e3b"
 down_revision = "975a79461cab"
 branch_labels = None
 depends_on = None
 
-def get_tables():
-    config = op.get_context().config
-    engine = engine_from_config(
-        config.get_section(config.config_ini_section), prefix="sqlalchemy."
-    )
-    inspector = reflection.Inspector.from_engine(engine)
+def get_tables(connection):
+    inspector = reflection.Inspector.from_engine(connection)
     return inspector.get_table_names()
 
 
+def _table_has_column(connection, table, column):
+    insp = reflection.Inspector.from_engine(connection)
+    return any(column in col["name"] for col in insp.get_columns(table))
+
+
 def upgrade():
-    tables = get_tables()
+    connection = op.get_bind()
+    tables = get_tables(connection)
     if "version_error" not in tables:
         op.create_table(
             "version_error",

dynamicannotationdb/migration/alembic/versions/8fdc843fc202_adding_permission_and_last_modified.py

@@ -7,66 +7,82 @@
 """
 from alembic import op
 import sqlalchemy as sa
-
-from sqlalchemy.dialects import postgresql
-from sqlalchemy import engine_from_config
 from sqlalchemy.engine import reflection
+from sqlalchemy.dialects import postgresql
+import logging
 
 # revision identifiers, used by Alembic.
 revision = "8fdc843fc202"
 down_revision = "6e7f580ff680"
 branch_labels = None
 depends_on = None
 
-def _table_has_column(table, column):
-    config = op.get_context().config
-    engine = engine_from_config(
-        config.get_section(config.config_ini_section), prefix="sqlalchemy."
-    )
 
-    insp = reflection.Inspector.from_engine(engine)
-    return any(column in col["name"] for col in insp.get_columns(table))
+def get_tables(connection):
+    inspector = reflection.Inspector.from_engine(connection)
+    return inspector.get_table_names()
 
 
+def _table_has_column(connection, table, column):
+    insp = reflection.Inspector.from_engine(connection)
+    return any(column in col["name"] for col in insp.get_columns(table))
+
 def upgrade():
+    logging.basicConfig(level=logging.INFO)
+    connection = op.get_bind()
+
     permission_enum = postgresql.ENUM(
         "PRIVATE", "GROUP", "PUBLIC", name="readwrite_permission"
     )
-    permission_enum.create(op.get_bind())   
+    try:
+        permission_enum.create(connection)
+    except Exception as e:
+        logging.info(f"Enum already exists: {e}")
 
-    if not _table_has_column("annotation_table_metadata", "write_permission"):
+    if not _table_has_column(
+        connection, "annotation_table_metadata", "write_permission"
+    ):
         op.add_column(
             "annotation_table_metadata",
             sa.Column(
                 "write_permission",
-                postgresql.ENUM("PRIVATE", "GROUP", "PUBLIC", name="readwrite_permission"),
+                postgresql.ENUM(
+                    "PRIVATE", "GROUP", "PUBLIC", name="readwrite_permission"
+                ),
                 nullable=True,
             ),
         )
         op.execute("UPDATE annotation_table_metadata SET write_permission = 'PRIVATE'")
         op.alter_column("annotation_table_metadata", "write_permission", nullable=False)
 
-    if not _table_has_column("annotation_table_metadata", "read_permission"):
+    if not _table_has_column(
+        connection, "annotation_table_metadata", "read_permission"
+    ):
         op.add_column(
             "annotation_table_metadata",
             sa.Column(
                 "read_permission",
-                postgresql.ENUM("PRIVATE", "GROUP", "PUBLIC", name="readwrite_permission"),
+                postgresql.ENUM(
+                    "PRIVATE", "GROUP", "PUBLIC", name="readwrite_permission"
+                ),
                 nullable=True,
             ),
         )
         op.execute("UPDATE annotation_table_metadata SET read_permission = 'PUBLIC'")
         op.alter_column("annotation_table_metadata", "read_permission", nullable=False)
 
-    if not _table_has_column("annotation_table_metadata", "last_modified"):
+    if not _table_has_column(connection, "annotation_table_metadata", "last_modified"):
         op.add_column(
             "annotation_table_metadata",
             sa.Column("last_modified", sa.DateTime(), nullable=True),
         )
 
-        op.execute("UPDATE annotation_table_metadata SET last_modified = current_timestamp")
+        op.execute(
+            "UPDATE annotation_table_metadata SET last_modified = current_timestamp"
+        )
         op.alter_column("annotation_table_metadata", "last_modified", nullable=False)
-
+
+
 def downgrade():
     # ### commands auto generated by Alembic - please adjust! ###
     op.drop_column("annotation_table_metadata", "last_modified")

dynamicannotationdb/migration/alembic/versions/975a79461cab_add_is_merged.py

@@ -7,9 +7,6 @@
 """
 from alembic import op
 import sqlalchemy as sa
-
-from sqlalchemy.dialects import postgresql
-from sqlalchemy import engine_from_config
 from sqlalchemy.engine import reflection
 
 # revision identifiers, used by Alembic.
@@ -18,18 +15,20 @@
 branch_labels = None
 depends_on = None
 
-def _table_has_column(table, column):
-    config = op.get_context().config
-    engine = engine_from_config(
-        config.get_section(config.config_ini_section), prefix="sqlalchemy."
-    )
 
-    insp = reflection.Inspector.from_engine(engine)
-    return any(column in col["name"] for col in insp.get_columns(table))
+def get_tables(connection):
+    inspector = reflection.Inspector.from_engine(connection)
+    return inspector.get_table_names()
 
 
+def _table_has_column(connection, table, column):
+    insp = reflection.Inspector.from_engine(connection)
+    return any(column in col["name"] for col in insp.get_columns(table))
+
 def upgrade():
-    if not _table_has_column("analysisversion", "is_merged"):
+    connection = op.get_bind()
+
+    if not _table_has_column(connection,"analysisversion", "is_merged"):
         op.add_column(
             "analysisversion",
             sa.Column("is_merged", sa.Boolean(), nullable=True, default=True),