31328 - custom schema name

data-tool/flows/common/colin_queries.py

@@ -76,13 +76,15 @@ def get_identifiers_per_batch(mig_batch_id: int) -> str:
     WHERE mcb.mig_batch_id IN ({mig_batch_id})
     """
 
-def unfreeze_identifiers() -> str:
+def unfreeze_identifiers(target_schema: str | None = 'public') -> str:
+    safe_schema = (target_schema or 'public').strip() or 'public'
+    safe_schema = '"' + safe_schema.replace('"', '""') + '"'
     return f"""
-    UPDATE corporation AS c
+    UPDATE {safe_schema}.corporation AS c
     SET corp_frozen_type_cd = NULL
-    FROM mig_group AS mg
-            JOIN mig_batch AS mb ON mb.mig_group_id = mg.id
-            JOIN mig_corp_batch AS mcb ON mcb.mig_batch_id = mb.id
+    FROM {safe_schema}.mig_group AS mg
+            JOIN {safe_schema}.mig_batch AS mb ON mb.mig_group_id = mg.id
+            JOIN {safe_schema}.mig_corp_batch AS mcb ON mcb.mig_batch_id = mb.id
     WHERE c.corp_num = mcb.corp_num
     -- cprd
     and mg.name in ('group_0', 'group_1', 'group_3', 'group_4','gcp_migration_group_test','misc_group')

data-tool/flows/common/query_utils.py

@@ -109,7 +109,9 @@ def prune_candidates_from_account(pruning_corps_list: list) -> str:
     WHERE corp_num IN {in_list}
     """
 
-def get_cutoff_timestamp_query() -> str:
+def get_cutoff_timestamp_query(target_schema: str | None = 'public') -> str:
+    safe_schema = (target_schema or 'public').strip() or 'public'
+    safe_schema = '"' + safe_schema.replace('"', '""') + '"'
     return f"""
-    SELECT extracted_at FROM colin_extract_version
+    SELECT extracted_at FROM {safe_schema}.colin_extract_version
     """