BFD-4810: Implement Prior Auth Sync Strategy

apps/bfd-db-migrator-ng/migrations/V1_38_0__add_utn_pattern_idx.sql

@@ -0,0 +1,2 @@
+CREATE INDEX ON idr.prior_auth (utn) WHERE utn NOT LIKE '-%';
+CREATE INDEX ON idr.prior_auth_item (utn) WHERE utn NOT LIKE '-%';

apps/bfd-pipeline-idr/extractor.py

@@ -192,6 +192,14 @@ def extract_idr_data(
             {"timestamp": compare_timestamp},
         )
 
+    def extract_full_idr_data(self, source: Source) -> Iterator[list[T]]:
+        start_time = self.cls.model_type().min_transaction_date
+        fetch_query = self.get_query(start_time, source)
+        logger.info("extracting full {}", self.cls.table())
+        return self.extract_many(
+            fetch_query.replace("{MIN_TS}", "%(timestamp)s"), {"timestamp": start_time}
+        )
+
     def _transform(self, batch: list[dict[str, DbType]]) -> list[T]:
         self.transform_timer.start()
         res = self.type_adapter.validate_python(

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

@@ -1,8 +1,9 @@
+import functools
 import itertools
 import operator
 from collections.abc import Awaitable, Callable, Iterator, Sequence
 from datetime import UTC, datetime
-from typing import Any, cast
+from typing import Any, Generic, cast, override
 
 import anyio
 import psycopg
@@ -78,7 +79,8 @@ async def _async_load(
             timeout=600,
         ) as pool:
             await pool.wait()
-            return await BatchLoader(
+            loader_cls = FullSyncBatchLoader if model.should_delete_missing() else BatchLoader
+            return await loader_cls(
                 fetch_results,
                 model,
                 pool,
@@ -91,7 +93,7 @@ async def _async_load(
             ).load()
 
 
-class BatchLoader:
+class BatchLoader(Generic[T]):  # noqa: UP046
     def __init__(
         self,
         fetch_results: Iterator[list[T]],
@@ -116,7 +118,7 @@ def __init__(
         self.batch_start = datetime.now(UTC)
         self.insert_cols = list(model.insert_keys())
         self.insert_cols.sort()
-        self.immutable = not model.update_timestamp_col()
+        self.immutable = model.is_immutable()
         self.meta_keys = (
             ["bfd_created_ts"] if self.immutable else ["bfd_created_ts", "bfd_updated_ts"]
         )
@@ -167,32 +169,20 @@ def __init__(
         self.full_batch_timer = Timer("full_batch", model, partition)
         self.full_load_timer = Timer("full_load", model, partition)
         self.load_type = load_type
+        self.load_mode = load_mode
         self.enable_load_progress = should_track_load_progress(load_mode)
 
     async def load(self) -> bool:
         timestamp = datetime.now(UTC)
-
         self.full_load_timer.start()
         async with self.pool.connection() as conn, conn.cursor(binary=True) as cur:
-            self.progress_start_timer.start()
-            await self._insert_batch_start(cur)
-            await conn.commit()
-            self.progress_start_timer.stop()
+            await self._record_batch_start(conn, cur, commit=True)
 
-        data_loaded = False
-        num_rows = 0
         batch_num = 1
-        while True:
-            self.idr_query_timer.start()
-            # We unfortunately need to use a while true loop here since we need to wrap the
-            # iterator with the timer calls.
-            results = next(self.fetch_results, None)
-            self.idr_query_timer.stop()
-            if not results:
-                break
 
+        async def _process_batch(results: list[T]) -> None:
+            nonlocal batch_num
             self.full_batch_timer.start()
-            data_loaded = True
             logger.info(
                 "{}-{}-{}: loading next {} results concurrently {} row(s) at a time",
                 self.table,
@@ -201,8 +191,6 @@ async def load(self) -> bool:
                 len(results),
                 PER_BATCH_CONCURRENT_ROWS,
             )
-            num_rows += len(results)
-
             self.sort_batch_timer.start()
             results.sort(key=operator.attrgetter(*self.ordered_pkeys))
             self.sort_batch_timer.stop()
@@ -258,6 +246,9 @@ async def _wrap_batch_chunk(
             batch_num += 1
             self.full_batch_timer.stop()
 
+        num_rows = await self._stage_all_batches(_process_batch)
+        data_loaded = num_rows > 0
+
         # Wait until the background worker signals that all pending loading tasks are completed
         # for the current partition before marking it totally complete
         self.worker_client.wait_until_done(self.model, self.partition)
@@ -342,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
@@ -354,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.
@@ -417,6 +409,105 @@ async def _copy_data(
                     [_remove_null_bytes(getattr(row, k)) for k in self.insert_cols]
                 )
 
+    async def _record_batch_start(
+        self, conn: psycopg.AsyncConnection, cur: psycopg.AsyncCursor[Any], commit: bool
+    ) -> None:
+        self.progress_start_timer.start()
+        await self._insert_batch_start(cur)
+        if commit:
+            await conn.commit()
+        self.progress_start_timer.stop()
+
+    def _next_batch(self) -> list[T] | None:
+        self.idr_query_timer.start()
+        results = next(self.fetch_results, None)
+        self.idr_query_timer.stop()
+        return results
+
+    async def _stage_all_batches(self, process_batch: Callable[[list[T]], Awaitable[None]]) -> int:
+        num_rows = 0
+
+        while True:
+            # We unfortunately need to use a while true loop here since we need to wrap the
+            # iterator with the timer calls.
+            self.idr_query_timer.start()
+            results = next(self.fetch_results, None)
+            self.idr_query_timer.stop()
+            if not results:
+                break
+
+            num_rows += len(results)
+            await process_batch(results)
+
+        return num_rows
+
+
+class FullSyncBatchLoader(BatchLoader[T]):
+    @override
+    async def load(self) -> bool:
+        timestamp = datetime.now(UTC)
+        self.full_load_timer.start()
+        data_loaded = False
+
+        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", True)
+
+            num_rows = await self._stage_all_batches(
+                functools.partial(self._copy_data, cur, full_temp_table)
+            )
+            data_loaded = num_rows > 0
+            logger.info(
+                "{}-{}: staged {} row(s) for full sync",
+                self.table,
+                self.partition.name,
+                num_rows,
+            )
+
+            self.insert_batch_timer.start()
+            updated_keys = await self._upsert(cur, full_temp_table, timestamp)
+            deleted_count = await self._delete_missing(cur, full_temp_table)
+            self.insert_batch_timer.stop()
+
+            logger.info(
+                "{}-{}: upserted {} new/changed row(s), deleted {} row(s) no longer present "
+                "upstream",
+                self.table,
+                self.partition.name,
+                len(updated_keys),
+                deleted_count,
+            )
+            await self._mark_batch_complete(cur)
+
+        self.full_load_timer.stop()
+        logger.info(
+            "{}-{}: finished full sync",
+            self.table,
+            self.partition.name,
+        )
+        return data_loaded
+
+    async def _delete_missing(self, cur: psycopg.AsyncCursor[Any], temp_tablename: str) -> int:
+        # We have to exclude our synthetic data that also exists in prod from deletion
+        synthetic_data_filter = self.model.synthetic_data_filter()
+        synthetic_where_clause = (
+            f"WHERE {synthetic_data_filter}"
+            if synthetic_data_filter and self.load_mode != LoadMode.SYNTHETIC
+            else ""
+        )
+        result = await cur.execute(  # type: ignore
+            f'''
+            DELETE FROM {self.table}
+            WHERE ({self.primary_keys_str}) IN (
+                SELECT {self.primary_keys_str} FROM {self.table}
+                {synthetic_where_clause}
+                EXCEPT
+                SELECT {self.primary_keys_str} FROM "{temp_tablename}"
+            )
+            '''  # type: ignore
+        )
+        return result.rowcount  # type: ignore
+
 
 def _remove_null_bytes(val: DbType) -> DbType:
     # Some IDR strings have null bytes.
@@ -431,5 +522,5 @@ def _remove_null_bytes(val: DbType) -> DbType:
 
 
 def should_track_load_progress(load_mode: LoadMode) -> bool:
-    # Whether to read/write load progress, which is diabled for synthetic and testing loads.
+    # Whether to read/write load progress, which is disabled for synthetic and testing loads.
     return load_mode == LoadMode.PROD or force_load_progress()

apps/bfd-pipeline-idr/mock-idr.sql

@@ -871,7 +871,5 @@ CREATE TABLE cms_edp_view_cvm_prau_prd.prauc (
     mr_count_end_dt DATE,
     att_phy_npi VARCHAR(10) NOT NULL,
     rrb_excl_ind VARCHAR(1),
-    idr_insrt_ts TIMESTAMPTZ,
-    idr_updt_ts TIMESTAMPTZ,
     PRIMARY KEY(mbi_num, utn, current_segment)
 );

apps/bfd-pipeline-idr/model/base_model.py

@@ -490,6 +490,20 @@ def should_replace() -> bool:
         """Whether to merge or replace data when loading this table."""
         return False
 
+    @staticmethod
+    def should_delete_missing() -> bool:
+        """Whether upstream data deletion requires manual cleanup on our end.
+
+        Upstream data can be deleted with no indicator like an obsolete timestamp, requiring
+        us to delete it on our end.
+        """
+        return False
+
+    @staticmethod
+    def synthetic_data_filter() -> str:
+        """Expression used to exclude synthetic data from being deleted in FullSyncBatchLoader."""
+        return ""
+
     @classmethod
     @abstractmethod
     def fetch_query(
@@ -528,6 +542,10 @@ def batch_timestamp_col(cls, is_historical: bool) -> list[str]:
     def update_timestamp_col(cls) -> list[str]:
         return cls._extract_meta_keys(UPDATE_TIMESTAMP)
 
+    @classmethod
+    def is_immutable(cls) -> bool:
+        return not cls.update_timestamp_col()
+
     @classmethod
     def batch_id_col_alias(cls) -> str | None:
         col = cls._single_or_default(BATCH_ID)

apps/bfd-pipeline-idr/model/idr_prior_auth.py

@@ -10,11 +10,8 @@
     ALIAS_PRVDR_ATT_PHY,
     ALIAS_PRVDR_ORDER_REFER,
     ALIAS_PRVDR_RENDER,
-    BATCH_TIMESTAMP,
     EXPR,
-    INSERT_EXCLUDE,
     PRIMARY_KEY_ORDER,
-    UPDATE_TIMESTAMP,
     IdrBaseModel,
     ModelType,
     Source,
@@ -24,7 +21,6 @@
     transform_null_date_to_max,
     transform_null_date_to_min,
 )
-from settings import MIN_PRIOR_AUTH_LOAD_DATE
 
 
 class IdrPriorAuth(IdrBaseModel):
@@ -62,13 +58,6 @@ class IdrPriorAuth(IdrBaseModel):
         BeforeValidator(transform_default_string),
     ]
     bfd_att_phy_npi_type: Annotated[int | None, {EXPR: provider_npi_type_expr(ALIAS_PRVDR_ATT_PHY)}]
-    # TBD: might have to change the insert & update ts once IDR adds those
-    idr_insrt_ts: Annotated[datetime, {BATCH_TIMESTAMP: True, INSERT_EXCLUDE: True}]
-    idr_updt_ts: Annotated[
-        datetime,
-        {UPDATE_TIMESTAMP: True, INSERT_EXCLUDE: True},
-        BeforeValidator(transform_null_date_to_min),
-    ]
 
     @override
     @staticmethod
@@ -85,6 +74,21 @@ def last_updated_date_column() -> list[str]:
     def model_type() -> ModelType:
         return ModelType.PRIOR_AUTH
 
+    @override
+    @staticmethod
+    def should_delete_missing() -> bool:
+        return True
+
+    @override
+    @classmethod
+    def is_immutable(cls) -> bool:
+        return False
+
+    @override
+    @staticmethod
+    def synthetic_data_filter() -> str:
+        return "utn NOT LIKE '-%'"
+
     @override
     @classmethod
     def fetch_query(cls, partition: LoadPartition, start_time: datetime, source: Source) -> str:
@@ -98,7 +102,7 @@ def fetch_query(cls, partition: LoadPartition, start_time: datetime, source: Sou
                 SELECT *, ROW_NUMBER() 
                     OVER (PARTITION BY mbi_num, utn ORDER BY current_segment) as row_order
                 FROM {IDR_PRIOR_AUTH_TABLE}
-                WHERE pa_req_rec_dt > '{MIN_PRIOR_AUTH_LOAD_DATE}'
+                WHERE pa_req_rec_dt > {{MIN_TS}}
             ) 
             SELECT {{COLUMNS}} FROM distinct_prior_auths {prior_auth}
             LEFT JOIN {IDR_PROVIDER_HISTORY_TABLE} {prvdr_att_phy}
@@ -110,5 +114,5 @@ def fetch_query(cls, partition: LoadPartition, start_time: datetime, source: Sou
             LEFT JOIN {IDR_PROVIDER_HISTORY_TABLE} {prvdr_render}
                 ON {prvdr_render}.prvdr_npi_num = {prior_auth}.render_npi
                 AND {prvdr_render}.prvdr_hstry_obslt_dt >= '{DEFAULT_MAX_DATE}'
-            {{WHERE_CLAUSE}} AND row_order = 1;
+            WHERE row_order = 1;
             """