fix: fix and improve alembic migrations (#176)

* fix: improve and fix alembic migrations

* fix: new migration

* refactor: ruff

api/alembic/versions/c3795dd0d722_init_db.py → api/alembic/versions/6edab3f23907_init_db.py

@@ -1,8 +1,8 @@
-"""init_db
+"""init db
 
-Revision ID: c3795dd0d722
+Revision ID: 6edab3f23907
 Revises: 
-Create Date: 2024-07-18 11:42:03.862912
+Create Date: 2024-10-16 13:50:44.743062
 
 """
 from typing import Sequence, Union, Text
@@ -12,7 +12,7 @@
 from app.db.tables.commons.json_encoded_dict import JSONEncodedDict
 
 # revision identifiers, used by Alembic.
-revision: str = 'c3795dd0d722'
+revision: str = '6edab3f23907'
 down_revision: Union[str, None] = None
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
@@ -37,9 +37,8 @@ def upgrade() -> None:
     sa.Column('CREATED_AT', sa.TIMESTAMP(timezone=True), nullable=False),
     sa.Column('UPDATED_AT', sa.TIMESTAMP(timezone=True), nullable=False),
     sa.Column('DELETED', sa.BOOLEAN(), nullable=False),
-    sa.PrimaryKeyConstraint('ID'),
-    sa.UniqueConstraint('UUID'),
-    schema='public'
+    sa.PrimaryKeyConstraint('ID', name=op.f('pk_model')),
+    sa.UniqueConstraint('UUID', name=op.f('uq_model_UUID'))
     )
     op.create_table('current_dataset',
     sa.Column('UUID', sa.UUID(), nullable=False),
@@ -48,21 +47,17 @@ def upgrade() -> None:
     sa.Column('DATE', sa.TIMESTAMP(timezone=True), nullable=False),
     sa.Column('CORRELATION_ID_COLUMN', sa.VARCHAR(), nullable=True),
     sa.Column('STATUS', sa.VARCHAR(), nullable=False),
-    sa.ForeignKeyConstraint(['MODEL_UUID'], ['public.model.UUID'], ),
-    sa.PrimaryKeyConstraint('UUID'),
-    sa.UniqueConstraint('UUID'),
-    schema='public'
+    sa.ForeignKeyConstraint(['MODEL_UUID'], ['model.UUID'], name=op.f('fk_current_dataset_MODEL_UUID_model')),
+    sa.PrimaryKeyConstraint('UUID', name=op.f('pk_current_dataset'))
     )
     op.create_table('reference_dataset',
     sa.Column('UUID', sa.UUID(), nullable=False),
     sa.Column('MODEL_UUID', sa.UUID(), nullable=False),
     sa.Column('PATH', sa.VARCHAR(), nullable=False),
     sa.Column('DATE', sa.TIMESTAMP(timezone=True), nullable=False),
     sa.Column('STATUS', sa.VARCHAR(), nullable=False),
-    sa.ForeignKeyConstraint(['MODEL_UUID'], ['public.model.UUID'], ),
-    sa.PrimaryKeyConstraint('UUID'),
-    sa.UniqueConstraint('UUID'),
-    schema='public'
+    sa.ForeignKeyConstraint(['MODEL_UUID'], ['model.UUID'], name=op.f('fk_reference_dataset_MODEL_UUID_model')),
+    sa.PrimaryKeyConstraint('UUID', name=op.f('pk_reference_dataset'))
     )
     op.create_table('current_dataset_metrics',
     sa.Column('UUID', sa.UUID(), nullable=False),
@@ -71,30 +66,26 @@ def upgrade() -> None:
     sa.Column('DATA_QUALITY', JSONEncodedDict(astext_type=Text()), nullable=True),
     sa.Column('DRIFT', JSONEncodedDict(astext_type=Text()), nullable=True),
     sa.Column('STATISTICS', JSONEncodedDict(astext_type=Text()), nullable=True),
-    sa.ForeignKeyConstraint(['CURRENT_UUID'], ['public.current_dataset.UUID'], ),
-    sa.PrimaryKeyConstraint('UUID'),
-    sa.UniqueConstraint('UUID'),
-    schema='public'
+    sa.ForeignKeyConstraint(['CURRENT_UUID'], ['current_dataset.UUID'], name=op.f('fk_current_dataset_metrics_CURRENT_UUID_current_dataset')),
+    sa.PrimaryKeyConstraint('UUID', name=op.f('pk_current_dataset_metrics'))
     )
     op.create_table('reference_dataset_metrics',
     sa.Column('UUID', sa.UUID(), nullable=False),
     sa.Column('REFERENCE_UUID', sa.UUID(), nullable=False),
     sa.Column('MODEL_QUALITY', JSONEncodedDict(astext_type=Text()), nullable=True),
     sa.Column('DATA_QUALITY', JSONEncodedDict(astext_type=Text()), nullable=True),
     sa.Column('STATISTICS', JSONEncodedDict(astext_type=Text()), nullable=True),
-    sa.ForeignKeyConstraint(['REFERENCE_UUID'], ['public.reference_dataset.UUID'], ),
-    sa.PrimaryKeyConstraint('UUID'),
-    sa.UniqueConstraint('UUID'),
-    schema='public'
+    sa.ForeignKeyConstraint(['REFERENCE_UUID'], ['reference_dataset.UUID'], name=op.f('fk_reference_dataset_metrics_REFERENCE_UUID_reference_dataset')),
+    sa.PrimaryKeyConstraint('UUID', name=op.f('pk_reference_dataset_metrics'))
     )
     # ### end Alembic commands ###
 
 
 def downgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.drop_table('reference_dataset_metrics', schema='public')
-    op.drop_table('current_dataset_metrics', schema='public')
-    op.drop_table('reference_dataset', schema='public')
-    op.drop_table('current_dataset', schema='public')
-    op.drop_table('model', schema='public')
+    op.drop_table('reference_dataset_metrics')
+    op.drop_table('current_dataset_metrics')
+    op.drop_table('reference_dataset')
+    op.drop_table('current_dataset')
+    op.drop_table('model')
     # ### end Alembic commands ###

api/app/db/database.py

@@ -15,7 +15,26 @@ class Reflected(DeferredReflection):
     __abstract__ = True
 
 
-BaseTable = declarative_base(metadata=MetaData(schema=get_config().db_config.db_schema))
+# https://github.com/sqlalchemy/alembic/discussions/1532
+# https://alembic.sqlalchemy.org/en/latest/naming.html
+naming_convention = {
+    'ix': 'ix_%(column_0_label)s',
+    'uq': 'uq_%(table_name)s_%(column_0_name)s',
+    'ck': 'ck_%(table_name)s_%(constraint_name)s',
+    'fk': 'fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s',
+    'pk': 'pk_%(table_name)s',
+}
+
+# https://github.com/sqlalchemy/alembic/discussions/1351
+# If the schema is the default, Alembic needs None otherwise migrations are messed up
+schema_name = (
+    None
+    if get_config().db_config.db_schema == 'public'
+    else get_config().db_config.db_schema
+)
+BaseTable = declarative_base(
+    metadata=MetaData(schema=schema_name, naming_convention=naming_convention)
+)
 
 
 class Database:

api/app/db/tables/current_dataset_metrics_table.py

@@ -13,7 +13,6 @@ class CurrentDatasetMetrics(Reflected, BaseTable, BaseDAO):
     uuid = Column(
         'UUID',
         UUID(as_uuid=True),
-        unique=True,
         nullable=False,
         default=uuid4,
         primary_key=True,

api/app/db/tables/current_dataset_table.py

@@ -13,7 +13,6 @@ class CurrentDataset(Reflected, BaseTable, BaseDAO):
     uuid = Column(
         'UUID',
         UUID(as_uuid=True),
-        unique=True,
         nullable=False,
         default=uuid4,
         primary_key=True,

api/app/db/tables/reference_dataset_metrics_table.py

@@ -13,7 +13,6 @@ class ReferenceDatasetMetrics(Reflected, BaseTable, BaseDAO):
     uuid = Column(
         'UUID',
         UUID(as_uuid=True),
-        unique=True,
         nullable=False,
         default=uuid4,
         primary_key=True,

api/app/db/tables/reference_dataset_table.py

@@ -13,7 +13,6 @@ class ReferenceDataset(Reflected, BaseTable, BaseDAO):
     uuid = Column(
         'UUID',
         UUID(as_uuid=True),
-        unique=True,
         nullable=False,
         default=uuid4,
         primary_key=True,