Adding table size checks to hash migration (#5310)

ethyca · Sep 21, 2024 · e864642 · e864642
1 parent 44a5095
commit e864642
Show file tree

Hide file tree

Showing 2 changed files with 110 additions and 36 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -15,7 +15,12 @@ The types of changes are:
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
-## [Unreleased](https://github.com/ethyca/fides/compare/2.45.1...main)
+## [Unreleased](https://github.com/ethyca/fides/compare/2.45.2...main)
+
+## [2.45.2](https://github.com/ethyca/fides/compare/2.45.1...2.45.2)
+
+### Fixed
+- Updated the hash migration script to only run on tables with less than 1 million rows. [#5310](https://github.com/ethyca/fides/pull/5310)
 
 ## [2.45.1](https://github.com/ethyca/fides/compare/2.45.0...2.45.1)
 

diff --git a/src/fides/api/alembic/migrations/versions/784f145ec892_add_is_hash_migrated_column.py b/src/fides/api/alembic/migrations/versions/784f145ec892_add_is_hash_migrated_column.py
@@ -8,6 +8,7 @@
 
 import sqlalchemy as sa
 from alembic import op
+from loguru import logger
 
 # revision identifiers, used by Alembic.
 revision = "784f145ec892"
@@ -21,57 +22,125 @@ def upgrade():
         "currentprivacypreferencev2",
         sa.Column("is_hash_migrated", sa.Boolean(), server_default="f", nullable=False),
     )
-    op.create_index(
-        "idx_currentprivacypreferencev2_unmigrated",
-        "currentprivacypreferencev2",
-        ["is_hash_migrated"],
-        unique=False,
-        postgresql_where=sa.text("is_hash_migrated IS false"),
-    )
     op.add_column(
         "custom_privacy_request_field",
         sa.Column("is_hash_migrated", sa.Boolean(), server_default="f", nullable=False),
     )
-    op.create_index(
-        "idx_custom_privacy_request_field_unmigrated",
-        "custom_privacy_request_field",
-        ["is_hash_migrated"],
-        unique=False,
-        postgresql_where=sa.text("is_hash_migrated IS false"),
-    )
     op.add_column(
         "privacypreferencehistory",
         sa.Column("is_hash_migrated", sa.Boolean(), server_default="f", nullable=False),
     )
-    op.create_index(
-        "idx_privacypreferencehistory_unmigrated",
-        "privacypreferencehistory",
-        ["is_hash_migrated"],
-        unique=False,
-        postgresql_where=sa.text("is_hash_migrated IS false"),
-    )
     op.add_column(
         "providedidentity",
         sa.Column("is_hash_migrated", sa.Boolean(), server_default="f", nullable=False),
     )
-    op.create_index(
-        "idx_providedidentity_unmigrated",
-        "providedidentity",
-        ["is_hash_migrated"],
-        unique=False,
-        postgresql_where=sa.text("is_hash_migrated IS false"),
-    )
     op.add_column(
         "servednoticehistory",
         sa.Column("is_hash_migrated", sa.Boolean(), server_default="f", nullable=False),
     )
-    op.create_index(
-        "idx_servednoticehistory_unmigrated",
-        "servednoticehistory",
-        ["is_hash_migrated"],
-        unique=False,
-        postgresql_where=sa.text("is_hash_migrated IS false"),
-    )
+
+    connection = op.get_bind()
+
+    # only run the index creation if the tables have less than 1 million rows
+    currentprivacypreferencev2_count = connection.execute(
+        sa.text("SELECT COUNT(*) FROM currentprivacypreferencev2")
+    ).scalar()
+    if currentprivacypreferencev2_count < 1000000:
+        op.create_index(
+            "idx_currentprivacypreferencev2_unmigrated",
+            "currentprivacypreferencev2",
+            ["is_hash_migrated"],
+            unique=False,
+            postgresql_where=sa.text("is_hash_migrated IS false"),
+        )
+    else:
+        logger.warning(
+            "The currentprivacypreferencev2 table has more than 1 million rows, "
+            "skipping index creation. Be sure to manually run "
+            "'CREATE INDEX CONCURRENTLY idx_currentprivacypreferencev2_unmigrated "
+            "ON currentprivacypreferencev2 (is_hash_migrated) "
+            "WHERE is_hash_migrated IS false'"
+        )
+
+    custom_privacy_request_field_count = connection.execute(
+        sa.text("SELECT COUNT(*) FROM custom_privacy_request_field")
+    ).scalar()
+    if custom_privacy_request_field_count < 1000000:
+        op.create_index(
+            "idx_custom_privacy_request_field_unmigrated",
+            "custom_privacy_request_field",
+            ["is_hash_migrated"],
+            unique=False,
+            postgresql_where=sa.text("is_hash_migrated IS false"),
+        )
+    else:
+        logger.warning(
+            "The custom_privacy_request_field table has more than 1 million rows, "
+            "skipping index creation. Be sure to manually run "
+            "'CREATE INDEX CONCURRENTLY idx_custom_privacy_request_field_unmigrated "
+            "ON custom_privacy_request_field (is_hash_migrated) "
+            "WHERE is_hash_migrated IS false'"
+        )
+
+    privacypreferencehistory_count = connection.execute(
+        sa.text("SELECT COUNT(*) FROM privacypreferencehistory")
+    ).scalar()
+    if privacypreferencehistory_count < 1000000:
+        op.create_index(
+            "idx_privacypreferencehistory_unmigrated",
+            "privacypreferencehistory",
+            ["is_hash_migrated"],
+            unique=False,
+            postgresql_where=sa.text("is_hash_migrated IS false"),
+        )
+    else:
+        logger.warning(
+            "The privacypreferencehistory table has more than 1 million rows, "
+            "skipping index creation. Be sure to manually run "
+            "'CREATE INDEX CONCURRENTLY idx_privacypreferencehistory_unmigrated "
+            "ON privacypreferencehistory (is_hash_migrated) "
+            "WHERE is_hash_migrated IS false'"
+        )
+
+    providedidentity_count = connection.execute(
+        sa.text("SELECT COUNT(*) FROM providedidentity")
+    ).scalar()
+    if providedidentity_count < 1000000:
+        op.create_index(
+            "idx_providedidentity_unmigrated",
+            "providedidentity",
+            ["is_hash_migrated"],
+            unique=False,
+            postgresql_where=sa.text("is_hash_migrated IS false"),
+        )
+    else:
+        logger.warning(
+            "The providedidentity table has more than 1 million rows, "
+            "skipping index creation. Be sure to manually run "
+            "'CREATE INDEX CONCURRENTLY idx_providedidentity_unmigrated "
+            "ON providedidentity (is_hash_migrated) "
+            "WHERE is_hash_migrated IS false'"
+        )
+
+    servednoticehistory_count = connection.execute(
+        sa.text("SELECT COUNT(*) FROM servednoticehistory")
+    ).scalar()
+    if servednoticehistory_count < 1000000:
+        op.create_index(
+            "idx_servednoticehistory_unmigrated",
+            "servednoticehistory",
+            ["is_hash_migrated"],
+            unique=False,
+            postgresql_where=sa.text("is_hash_migrated IS false"),
+        )
+    else:
+        logger.warning(
+            "The servednoticehistory table has more than 1 million rows, "
+            "skipping index creation. Be sure to manually run "
+            "'CREATE INDEX CONCURRENTLY idx_servednoticehistory_unmigrated "
+            "ON servednoticehistory (is_hash_migrated) "
+            "WHERE is_hash_migrated IS false'"
+        )
 
 
 def downgrade():