Fixed Alembic revision scripts that were causing the test cases to fail

cfpb · Dec 4, 2023 · 13b10c6 · 13b10c6
1 parent 804dbc1
commit 13b10c6
Show file tree

Hide file tree

Showing 5 changed files with 119 additions and 72 deletions.
diff --git a/db_revisions/env.py b/db_revisions/env.py
@@ -86,11 +86,15 @@ def run_migrations_online() -> None:
             context.config.get_section(context.config.config_ini_section, {}),
             prefix="sqlalchemy.",
             poolclass=pool.NullPool,
+            render_as_batch=True,
         )
 
     with connectable.connect() as connection:
         context.configure(
-            connection=connection, target_metadata=target_metadata, version_table_schema=target_metadata.schema
+            connection=connection,
+            target_metadata=target_metadata,
+            version_table_schema=target_metadata.schema,
+            render_as_batch=True,
         )
 
         with context.begin_transaction():

diff --git a/db_revisions/versions/045aa502e050_update_financial_institutions_table.py b/db_revisions/versions/045aa502e050_update_financial_institutions_table.py
@@ -10,6 +10,8 @@
 from alembic import op
 import sqlalchemy as sa
 
+from db_revisions.utils import table_exists
+
 
 # revision identifiers, used by Alembic.
 revision: str = "045aa502e050"
@@ -19,54 +21,85 @@
 
 
 def upgrade() -> None:
-    op.add_column("financial_institutions", sa.Column("tax_id", sa.String(length=9), nullable=False))
-    op.add_column("financial_institutions", sa.Column("rssd_id", sa.Integer(), nullable=False))
-    op.add_column(
-        "financial_institutions", sa.Column("primary_federal_regulator_id", sa.String(length=4), nullable=False)
-    )
-    op.add_column("financial_institutions", sa.Column("hmda_institution_type_id", sa.String(), nullable=False))
-    op.add_column("financial_institutions", sa.Column("sbl_institution_type_id", sa.String(), nullable=False))
-    op.add_column("financial_institutions", sa.Column("hq_address_street_1", sa.String(), nullable=False))
-    op.add_column("financial_institutions", sa.Column("hq_address_street_2", sa.String(), nullable=False))
-    op.add_column("financial_institutions", sa.Column("hq_address_city", sa.String(), nullable=False))
-    op.add_column("financial_institutions", sa.Column("hq_address_state", sa.String(length=2), nullable=False))
-    op.add_column("financial_institutions", sa.Column("hq_address_zip", sa.String(length=5), nullable=False))
-    op.add_column("financial_institutions", sa.Column("parent_lei", sa.String(length=20), nullable=False))
-    op.add_column("financial_institutions", sa.Column("parent_legal_name", sa.String(), nullable=False))
-    op.add_column("financial_institutions", sa.Column("parent_rssd_id", sa.Integer(), nullable=False))
-    op.add_column("financial_institutions", sa.Column("top_holder_lei", sa.String(length=20), nullable=False))
-    op.add_column("financial_institutions", sa.Column("top_holder_legal_name", sa.String(), nullable=False))
-    op.add_column("financial_institutions", sa.Column("top_holder_rssd_id", sa.Integer(), nullable=False))
-    op.create_index(
-        op.f("ix_financial_institutions_hmda_institution_type_id"),
-        "financial_institutions",
-        ["hmda_institution_type_id"],
-        unique=False,
-    )
-    op.create_index(
-        op.f("ix_financial_institutions_hq_address_state"), "financial_institutions", ["hq_address_state"], unique=False
-    )
-    op.create_index(
-        op.f("ix_financial_institutions_primary_federal_regulator_id"),
-        "financial_institutions",
-        ["primary_federal_regulator_id"],
-        unique=False,
-    )
-    op.create_index(
-        op.f("ix_financial_institutions_sbl_institution_type_id"),
-        "financial_institutions",
-        ["sbl_institution_type_id"],
-        unique=False,
-    )
-    op.create_unique_constraint(None, "financial_institutions", ["rssd_id"])
-    op.create_unique_constraint(None, "financial_institutions", ["tax_id"])
-    op.create_foreign_key(None, "financial_institutions", "federal_regulator", ["primary_federal_regulator_id"], ["id"])
-    op.create_foreign_key(None, "financial_institutions", "address_state", ["hq_address_state"], ["code"])
-    op.create_foreign_key(None, "financial_institutions", "hmda_institution_type", ["hmda_institution_type_id"], ["id"])
-    op.create_foreign_key(None, "financial_institutions", "sbl_institution_type", ["sbl_institution_type_id"], ["id"])
+    if table_exists("financial_institutions"):
+        op.add_column("financial_institutions", sa.Column("tax_id", sa.String(length=9), nullable=False))
+        op.add_column("financial_institutions", sa.Column("rssd_id", sa.Integer(), nullable=False))
+        op.add_column(
+            "financial_institutions", sa.Column("primary_federal_regulator_id", sa.String(length=4), nullable=False)
+        )
+        op.add_column("financial_institutions", sa.Column("hmda_institution_type_id", sa.String(), nullable=False))
+        op.add_column("financial_institutions", sa.Column("sbl_institution_type_id", sa.String(), nullable=False))
+        op.add_column("financial_institutions", sa.Column("hq_address_street_1", sa.String(), nullable=False))
+        op.add_column("financial_institutions", sa.Column("hq_address_street_2", sa.String(), nullable=False))
+        op.add_column("financial_institutions", sa.Column("hq_address_city", sa.String(), nullable=False))
+        op.add_column("financial_institutions", sa.Column("hq_address_state", sa.String(length=2), nullable=False))
+        op.add_column("financial_institutions", sa.Column("hq_address_zip", sa.String(length=5), nullable=False))
+        op.add_column("financial_institutions", sa.Column("parent_lei", sa.String(length=20), nullable=False))
+        op.add_column("financial_institutions", sa.Column("parent_legal_name", sa.String(), nullable=False))
+        op.add_column("financial_institutions", sa.Column("parent_rssd_id", sa.Integer(), nullable=False))
+        op.add_column("financial_institutions", sa.Column("top_holder_lei", sa.String(length=20), nullable=False))
+        op.add_column("financial_institutions", sa.Column("top_holder_legal_name", sa.String(), nullable=False))
+        op.add_column("financial_institutions", sa.Column("top_holder_rssd_id", sa.Integer(), nullable=False))
+        op.create_index(
+            op.f("ix_financial_institutions_hmda_institution_type_id"),
+            "financial_institutions",
+            ["hmda_institution_type_id"],
+            unique=False,
+        )
+        op.create_index(
+            op.f("ix_financial_institutions_hq_address_state"),
+            "financial_institutions",
+            ["hq_address_state"],
+            unique=False,
+        )
+        op.create_index(
+            op.f("ix_financial_institutions_primary_federal_regulator_id"),
+            "financial_institutions",
+            ["primary_federal_regulator_id"],
+            unique=False,
+        )
+        op.create_index(
+            op.f("ix_financial_institutions_sbl_institution_type_id"),
+            "financial_institutions",
+            ["sbl_institution_type_id"],
+            unique=False,
+        )
+        with op.batch_alter_table("financial_institutions") as batch_op:
+            batch_op.create_unique_constraint("unique_association", ["rssd_id', 'tax_id"])
+            batch_op.create_foreign_key(
+                "fk_federal_regulator_financial_institutions",
+                "federal_regulator",
+                ["primary_federal_regulator_id"],
+                ["id"],
+            )
+            batch_op.create_foreign_key(
+                "fk_address_state_financial_institutions", "address_state", ["hq_address_state"], ["code"]
+            )
+            batch_op.create_foreign_key(
+                "fk_hmda_institution_type_financial_institutions",
+                "hmda_institution_type",
+                ["hmda_institution_type_id"],
+                ["id"],
+            )
+            batch_op.create_foreign_key(
+                "fk_sbl_institution_type_financial_institutions",
+                "fsbl_institution_type",
+                ["sbl_institution_type_id"],
+                ["id"],
+            )
 
 
 def downgrade() -> None:
+    op.drop_constraint(
+        constraint_name="fk_federal_regulator_financial_institutions", table_name="financial_institutions"
+    )
+    op.drop_constraint(constraint_name="fk_address_state_financial_institutions", table_name="financial_institutions")
+    op.drop_constraint(
+        constraint_name="fk_hmda_institution_type_financial_institutions", table_name="financial_institutions"
+    )
+    op.drop_constraint(
+        constraint_name="fk_sbl_institution_type_financial_institutions", table_name="financial_institutions"
+    )
     op.drop_column("financial_institutions", "top_holder_rssd_id")
     op.drop_column("financial_institutions", "top_holder_legal_name")
     op.drop_column("financial_institutions", "top_holder_lei")

diff --git a/db_revisions/versions/1f6c33f20a2e_create_address_state_table.py b/db_revisions/versions/1f6c33f20a2e_create_address_state_table.py
@@ -10,6 +10,8 @@
 from alembic import op
 import sqlalchemy as sa
 
+from db_revisions.utils import table_exists
+
 
 # revision identifiers, used by Alembic.
 revision: str = "1f6c33f20a2e"
@@ -19,14 +21,16 @@
 
 
 def upgrade() -> None:
-    op.create_table(
-        "address_state",
-        sa.Column("code", sa.String(length=2), nullable=False),
-        sa.Column("name", sa.String(), nullable=False),
-        sa.Column("event_time", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
-        sa.PrimaryKeyConstraint("code"),
-        sa.UniqueConstraint("name"),
-    )
+    if not table_exists("address_state"):
+        op.create_table(
+            "address_state",
+            sa.Column("code", sa.String(length=2), nullable=False),
+            sa.Column("name", sa.String(), nullable=False),
+            sa.Column("event_time", sa.DateTime(), server_default=sa.func.now(), nullable=False),
+            sa.PrimaryKeyConstraint("code"),
+            sa.UniqueConstraint("name"),
+        )
+        op.create_index(op.f("ix_address_state_code"), "address_state", ["code"], unique=False)
 
 
 def downgrade() -> None:

diff --git a/db_revisions/versions/56ef0b5cd2d4_create_sbl_institution_type_table.py b/db_revisions/versions/56ef0b5cd2d4_create_sbl_institution_type_table.py
@@ -10,6 +10,8 @@
 from alembic import op
 import sqlalchemy as sa
 
+from db_revisions.utils import table_exists
+
 
 # revision identifiers, used by Alembic.
 revision: str = "56ef0b5cd2d4"
@@ -19,15 +21,16 @@
 
 
 def upgrade() -> None:
-    op.create_table(
-        "sbl_institution_type",
-        sa.Column("id", sa.String(), nullable=False),
-        sa.Column("name", sa.String(), nullable=False),
-        sa.Column("event_time", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
-        sa.PrimaryKeyConstraint("id"),
-        sa.UniqueConstraint("name"),
-    )
-    op.create_index(op.f("ix_sbl_institution_type_id"), "sbl_institution_type", ["id"], unique=False)
+    if not table_exists("sbl_institution_type"):
+        op.create_table(
+            "sbl_institution_type",
+            sa.Column("id", sa.String(), nullable=False),
+            sa.Column("name", sa.String(), nullable=False),
+            sa.Column("event_time", sa.DateTime(), server_default=sa.func.now(), nullable=False),
+            sa.PrimaryKeyConstraint("id"),
+            sa.UniqueConstraint("name"),
+        )
+        op.create_index(op.f("ix_sbl_institution_type_id"), "sbl_institution_type", ["id"], unique=False)
 
 
 def downgrade() -> None:

diff --git a/db_revisions/versions/8b1ba6a3275b_create_hmda_institution_type_table.py b/db_revisions/versions/8b1ba6a3275b_create_hmda_institution_type_table.py
@@ -10,6 +10,8 @@
 from alembic import op
 import sqlalchemy as sa
 
+from db_revisions.utils import table_exists
+
 
 # revision identifiers, used by Alembic.
 revision: str = "8b1ba6a3275b"
@@ -19,15 +21,16 @@
 
 
 def upgrade() -> None:
-    op.create_table(
-        "hmda_institution_type",
-        sa.Column("id", sa.String(), nullable=False),
-        sa.Column("name", sa.String(), nullable=False),
-        sa.Column("event_time", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
-        sa.PrimaryKeyConstraint("id"),
-        sa.UniqueConstraint("name"),
-    )
-    op.create_index(op.f("ix_hmda_institution_type_id"), "hmda_institution_type", ["id"], unique=False)
+    if not table_exists("hmda_institution_type"):
+        op.create_table(
+            "hmda_institution_type",
+            sa.Column("id", sa.String(), nullable=False),
+            sa.Column("name", sa.String(), nullable=False),
+            sa.Column("event_time", sa.DateTime(), server_default=sa.func.now(), nullable=False),
+            sa.PrimaryKeyConstraint("id"),
+            sa.UniqueConstraint("name"),
+        )
+        op.create_index(op.f("ix_hmda_institution_type_id"), "hmda_institution_type", ["id"], unique=False)
 
 
 def downgrade() -> None: