copy indexes onto temp table for pa only + some clean up

Author: mel1-G

apps/bfd-pipeline-idr/load-synthetic-env.sh

@@ -46,4 +46,4 @@ if [[ -n "$1" ]]; then
     args+=('--seed-from' "$1")
 fi
 
-IDR_ENABLE_DATE_PARTITIONS=0 IDR_ENABLE_PRIOR_AUTH=1 uv run pipeline.py "${args[@]}"
+IDR_ENABLE_DATE_PARTITIONS=0 uv run pipeline.py "${args[@]}"

apps/bfd-pipeline-idr/loader.py

@@ -333,7 +333,7 @@ async def _mark_batch_complete(self, cur: psycopg.AsyncCursor) -> None:
         )
 
     async def _setup_temp_table(
-        self, cur: psycopg.AsyncCursor[Any], suffix: str | None = None
+        self, cur: psycopg.AsyncCursor[Any], suffix: str | None = None, copy_indexes: bool = False
     ) -> str:
         # Load each batch into a temp table
         # This is necessary because we want to use COPY to quickly
@@ -345,8 +345,9 @@ async def _setup_temp_table(
         # For simplicity's sake, we'll create our temp tables using the existing schema and
         # just drop the columns we need to ignore.
         full_tablename = f"{self.temp_table}_{suffix or ''}"
+        copy_indexes_option = "INCLUDING INDEXES" if copy_indexes else ""
         await cur.execute(
-            f'CREATE TEMPORARY TABLE "{full_tablename}" (LIKE {self.table}) '  # type: ignore
+            f'CREATE TEMPORARY TABLE "{full_tablename}" (LIKE {self.table} {copy_indexes_option}) '  # type: ignore
             "ON COMMIT DROP"
         )
         # Created/updated columns don't need to be loaded from the source.
@@ -450,7 +451,7 @@ async def load(self) -> bool:
 
         async with self.pool.connection() as conn, conn.cursor(binary=True) as cur:
             await self._record_batch_start(conn, cur, commit=False)
-            full_temp_table = await self._setup_temp_table(cur, "full_temp")
+            full_temp_table = await self._setup_temp_table(cur, "full_temp", True)
 
             num_rows = await self._stage_all_batches(
                 functools.partial(self._copy_data, cur, full_temp_table)

apps/bfd-pipeline-idr/run-db.sh

@@ -21,7 +21,6 @@ function do_load() {
   BFD_DB_PASSWORD="$DB_PASSWORD" \
   BFD_DB_ENDPOINT="$DB_ENDPOINT" \
   IDR_ENABLE_DATE_PARTITIONS=0 \
-  IDR_ENABLE_PRIOR_AUTH=1 \
   uv run pipeline.py \
   --source postgres \
   --load-mode synthetic \

apps/bfd-pipeline-idr/settings.py

@@ -27,7 +27,7 @@ def bfd_test_date() -> datetime | None:
 
 
 def enable_prior_auth_ingestion() -> bool:
-    return _parse_bool_default_false("IDR_ENABLE_PRIOR_AUTH")
+    return _parse_bool_default_true("IDR_ENABLE_PRIOR_AUTH")
 
 
 ENABLE_DATE_PARTITIONS = _parse_bool_default_true("IDR_ENABLE_DATE_PARTITIONS")

apps/bfd-pipeline-idr/test_pipeline.py

@@ -26,7 +26,6 @@
 from model.base_model import LoadMode, Source
 from pipeline import run
 from pydantic_utils import fields
-from settings import enable_prior_auth_ingestion
 
 # ryuk throws a 500 or 404 error for some reason
 # seems to have issues with podman https://github.com/testcontainers/testcontainers-python/issues/753
@@ -82,16 +81,15 @@ def _do_test_pipeline(conn: Connection[DictRow], load_type: LoadType) -> None:
     rows = cur.fetchmany(1)
     assert rows[0]["bene_xref_efctv_sk"] == 353816021
 
-    if enable_prior_auth_ingestion():
-        cur = conn.execute("select * from idr.prior_auth order by mbi_num")
-        assert cur.rowcount == 21
-        rows = cur.fetchmany(1)
-        assert rows[0]["mbi_num"] == "1OX4Y88RV68"
+    cur = conn.execute("select * from idr.prior_auth order by mbi_num")
+    assert cur.rowcount == 21
+    rows = cur.fetchmany(1)
+    assert rows[0]["mbi_num"] == "1OX4Y88RV68"
 
-        cur = conn.execute("select * from idr.prior_auth_item order by mbi_num")
-        assert cur.rowcount == 64
-        rows = cur.fetchmany(1)
-        assert rows[0]["mbi_num"] == "1OX4Y88RV68"
+    cur = conn.execute("select * from idr.prior_auth_item order by mbi_num")
+    assert cur.rowcount == 64
+    rows = cur.fetchmany(1)
+    assert rows[0]["mbi_num"] == "1OX4Y88RV68"
 
     cur = conn.execute("select max(last_ts) as max_ts from idr.load_progress")
     row = cur.fetchone()
@@ -546,9 +544,6 @@ def _do_test_pipeline(conn: Connection[DictRow], load_type: LoadType) -> None:
 
 
 def _do_test_prior_auth_update_and_delete(conn: Connection[DictRow], load_type: LoadType) -> None:
-    if not enable_prior_auth_ingestion():
-        return
-
     cur = conn.execute(
         "select * from idr.prior_auth where mbi_num = '7ZM6HW2AT68' and utn = '-OTENCJLOQRAKA'"
     )
@@ -676,7 +671,6 @@ def _setup_pipeline_environment(info: psycopg.ConnectionInfo) -> None:
     os.environ["BFD_TEST_DATE"] = "2023-04-02"
     os.environ["IDR_PER_BATCH_MIN_CONNECTIONS"] = "1"
     os.environ["IDR_PER_BATCH_MAX_CONNECTIONS"] = "1"
-    os.environ["IDR_ENABLE_PRIOR_AUTH"] = "1"
 
 
 @pytest.fixture(scope="module")

apps/bfd-server-ng/src/test/java/gov/cms/bfd/server/ng/IntegrationTestConfiguration.java

@@ -132,7 +132,6 @@ private void runPython(PostgreSQLContainer<?> container, Instant date, String...
     env.put("BFD_TEST_DATE", date.toString());
     env.put("IDR_PER_BATCH_MIN_CONNECTIONS", "1");
     env.put("IDR_PER_BATCH_MAX_CONNECTIONS", "1");
-    env.put("IDR_ENABLE_PRIOR_AUTH", "1");
     // Suppress noisy output unless the log level is explicitly set
     env.putIfAbsent("IDR_LOG_LEVEL", "WARNING");