BFD-4810: Implement Prior Auth Sync Strategy

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("{LAST_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/loader.py

@@ -78,7 +78,10 @@ async def _async_load(
             timeout=600,
         ) as pool:
             await pool.wait()
-            return await BatchLoader(
+            loader_cls = (
+                FullSyncBatchLoader if model.should_fully_sync_delete_diff() else BatchLoader
+            )
+            return await loader_cls(
                 fetch_results,
                 model,
                 pool,
@@ -116,7 +119,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,6 +170,7 @@ 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:
@@ -418,6 +422,77 @@ async def _copy_data(
                 )
 
 
+class FullSyncBatchLoader(BatchLoader):
+    async def load(self) -> bool:
+        timestamp = datetime.now(UTC)
+        self.full_load_timer.start()
+        num_rows = 0
+
+        async with self.pool.connection() as conn, conn.cursor(binary=True) as cur:
+            self.progress_start_timer.start()
+            await self._insert_batch_start(cur)
+            self.progress_start_timer.stop()
+
+            full_temp_table = await self._setup_temp_table(cur, "full_temp")
+
+            while True:
+                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 self._copy_data(cur, full_temp_table, results)
+
+            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 True
+
+    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 = (
+            "" if self.load_mode == LoadMode.SYNTHETIC else "WHERE utn NOT LIKE '-%'"
+        )
+        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_data_filter}
+                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.
     # Postgres doesn't allow these in text fields.

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,15 @@ def should_replace() -> bool:
         """Whether to merge or replace data when loading this table."""
         return False
 
+    @staticmethod
+    def should_fully_sync_delete_diff() -> 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
+
     @classmethod
     @abstractmethod
     def fetch_query(
@@ -528,6 +537,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,16 @@ def last_updated_date_column() -> list[str]:
     def model_type() -> ModelType:
         return ModelType.PRIOR_AUTH
 
+    @override
+    @staticmethod
+    def should_fully_sync_delete_diff() -> bool:
+        return True
+
+    @override
+    @classmethod
+    def is_immutable(cls) -> bool:
+        return False
+
     @override
     @classmethod
     def fetch_query(cls, partition: LoadPartition, start_time: datetime, source: Source) -> str:
@@ -98,7 +97,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 > {{LAST_TS}}
             ) 
             SELECT {{COLUMNS}} FROM distinct_prior_auths {prior_auth}
             LEFT JOIN {IDR_PROVIDER_HISTORY_TABLE} {prvdr_att_phy}
@@ -110,5 +109,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;
             """

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

@@ -7,18 +7,14 @@
 from load_partition import LoadPartition
 from model.base_model import (
     ALIAS_PRIOR_AUTH,
-    BATCH_TIMESTAMP,
-    INSERT_EXCLUDE,
     PRIMARY_KEY_ORDER,
-    UPDATE_TIMESTAMP,
     IdrBaseModel,
     ModelType,
     Source,
     transform_default_string,
     transform_null_date_to_max,
     transform_null_date_to_min,
 )
-from settings import MIN_PRIOR_AUTH_LOAD_DATE
 
 
 class IdrPriorAuthItem(IdrBaseModel):
@@ -43,12 +39,6 @@ class IdrPriorAuthItem(IdrBaseModel):
     mr_count_st_dt: Annotated[date, BeforeValidator(transform_null_date_to_min)]
     mr_count_end_dt: Annotated[date, BeforeValidator(transform_null_date_to_max)]
     rrb_excl_ind: Annotated[str, BeforeValidator(transform_default_string)]
-    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
@@ -65,6 +55,16 @@ def last_updated_date_column() -> list[str]:
     def model_type() -> ModelType:
         return ModelType.PRIOR_AUTH
 
+    @override
+    @staticmethod
+    def should_fully_sync_delete_diff() -> bool:
+        return True
+
+    @override
+    @classmethod
+    def is_immutable(cls) -> bool:
+        return False
+
     @override
     @classmethod
     def fetch_query(cls, partition: LoadPartition, start_time: datetime, source: Source) -> str:
@@ -75,8 +75,8 @@ def fetch_query(cls, partition: LoadPartition, start_time: datetime, source: Sou
                 SELECT *, ROW_NUMBER() 
                     OVER (PARTITION BY mbi_num, utn, current_segment ORDER BY mbi_num) as row_order
                 FROM {IDR_PRIOR_AUTH_TABLE}
-                WHERE pa_req_rec_dt > '{MIN_PRIOR_AUTH_LOAD_DATE}'
+                WHERE pa_req_rec_dt > {{LAST_TS}}
             )
             SELECT {{COLUMNS}} FROM distinct_prior_auths {prior_auth}
-            {{WHERE_CLAUSE}} AND row_order = 1;
+            WHERE row_order = 1;
             """

apps/bfd-pipeline-idr/pipeline_utils.py

@@ -87,7 +87,11 @@ def extract_and_load(
             else:
                 logger.info("no previous progress for {} - {}", cls.table(), partition.name)
 
-            data_iter = data_extractor.extract_idr_data(progress, job_start, source)
+            data_iter = (
+                data_extractor.extract_full_idr_data(source)
+                if cls.should_fully_sync_delete_diff()
+                else data_extractor.extract_idr_data(progress, job_start, source)
+            )
             res = loader.load(
                 data_iter,
                 cls,

apps/bfd-pipeline-idr/test_pipeline.py

@@ -545,6 +545,77 @@ def _do_test_pipeline(conn: Connection[DictRow], load_type: LoadType) -> None:
         assert updated_ss_job.completion_time >= ss_clm_ts
 
 
+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'"
+    )
+    assert cur.rowcount == 1
+    rows = cur.fetchmany(6)
+    assert rows[0]["mbi_num"] == "7ZM6HW2AT68"
+    original_updated_ts = rows[0]["bfd_updated_ts"]
+    original_name = rows[0]["name"]
+
+    cur = conn.execute(
+        "select * from idr.prior_auth where mbi_num = '5OH0K85GU23' and utn = '-SC21YQR4UY4LI'"
+    )
+    assert cur.rowcount == 1
+    row = cur.fetchone()
+    assert row is not None
+
+    prauc_table = sql.Identifier("cms_edp_view_cvm_prau_prd", "prauc")
+    conn.execute(
+        t"""
+        UPDATE {prauc_table:i}
+        SET name = 'BITE AID PHARMACY'
+        WHERE mbi_num = '7ZM6HW2AT68'
+        AND utn = '-OTENCJLOQRAKA'
+        """
+    )
+
+    conn.execute(
+        t"""
+        DELETE FROM {prauc_table:i}
+        WHERE mbi_num = '5OH0K85GU23'
+        AND utn = '-SC21YQR4UY4LI'
+        """
+    )
+    conn.commit()
+
+    _advance_time(datetime.now() + timedelta(days=1))
+    run(Source.POSTGRES, LoadMode.SYNTHETIC, load_type)
+
+    # verify that updated rows by upstream were updated
+    cur = conn.execute(
+        "select * from idr.prior_auth where mbi_num = '7ZM6HW2AT68' and utn = '-OTENCJLOQRAKA'"
+    )
+    assert cur.rowcount == 1
+    updated_row = cur.fetchone()
+    assert updated_row is not None
+    assert updated_row["name"] != original_name
+    assert updated_row["bfd_updated_ts"] > original_updated_ts
+
+    # verify that deleted rows by upstream were deleted in header and item level for prior auth
+    cur = conn.execute(
+        "select * from idr.prior_auth where mbi_num = '5OH0K85GU23' and utn = '-SC21YQR4UY4LI'"
+    )
+    assert cur.rowcount == 0
+
+    cur = conn.execute(
+        "select * from idr.prior_auth_item where mbi_num = '5OH0K85GU23' and utn = '-SC21YQR4UY4LI'"
+    )
+    assert cur.rowcount == 0
+
+    # verify that untouched rows by upstream were not updated
+    cur = conn.execute(
+        "select * from idr.prior_auth where mbi_num = '7ZM6HW2AT68' and utn = '-RVUOWAUT5V5QZ'"
+    )
+    rows = cur.fetchmany(2)
+    assert rows[0]["bfd_updated_ts"] < updated_row["bfd_updated_ts"]
+
+
 def _advance_time(timestamp: datetime) -> None:
     new_time = timestamp + timedelta(minutes=1)
     os.environ["BFD_TEST_DATE"] = new_time.isoformat()
@@ -623,6 +694,7 @@ def _test_pipeline_load(postgres_db: tuple[PostgresContainer, str], load_type: L
         _reset_db(conn, sample_dir, postgres)
         _setup_pipeline_environment(conn.info)
         _do_test_pipeline(cast(Connection[DictRow], conn), load_type)
+        _do_test_prior_auth_update_and_delete(cast(Connection[DictRow], conn), load_type)
     logger.remove()