clean up + partial idx for delete

Author: mel1-G

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

@@ -197,7 +197,7 @@ def extract_full_idr_data(self, source: Source) -> Iterator[list[T]]:
         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}
+            fetch_query.replace("{MIN_TS}", "%(timestamp)s"), {"timestamp": start_time}
         )
 
     def _transform(self, batch: list[dict[str, DbType]]) -> list[T]:

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,9 +79,7 @@ async def _async_load(
             timeout=600,
         ) as pool:
             await pool.wait()
-            loader_cls = (
-                FullSyncBatchLoader if model.should_fully_sync_delete_diff() else BatchLoader
-            )
+            loader_cls = FullSyncBatchLoader if model.should_delete_missing() else BatchLoader
             return await loader_cls(
                 fetch_results,
                 model,
@@ -94,7 +93,7 @@ async def _async_load(
             ).load()
 
 
-class BatchLoader:
+class BatchLoader(Generic[T]):  # noqa: UP046
     def __init__(
         self,
         fetch_results: Iterator[list[T]],
@@ -175,28 +174,15 @@ def __init__(
 
     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,
@@ -205,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()
@@ -262,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)
@@ -421,29 +408,54 @@ 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
 
-class FullSyncBatchLoader(BatchLoader):
+            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()
-        num_rows = 0
+        data_loaded = False
 
         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()
-
+            await self._record_batch_start(conn, cur, commit=False)
             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)
-
+            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,
@@ -472,19 +484,22 @@ async def load(self) -> bool:
             self.table,
             self.partition.name,
         )
-        return True
+        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 = (
-            "" if self.load_mode == LoadMode.SYNTHETIC else "WHERE utn NOT LIKE '-%'"
+        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_data_filter}
+                {synthetic_where_clause}
                 EXCEPT
                 SELECT {self.primary_keys_str} FROM "{temp_tablename}"
             )
@@ -506,5 +521,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/model/base_model.py

@@ -491,14 +491,19 @@ def should_replace() -> bool:
         return False
 
     @staticmethod
-    def should_fully_sync_delete_diff() -> bool:
+    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(

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

@@ -76,14 +76,19 @@ def model_type() -> ModelType:
 
     @override
     @staticmethod
-    def should_fully_sync_delete_diff() -> bool:
+    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:
@@ -97,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 > {{LAST_TS}}
+                WHERE pa_req_rec_dt > {{MIN_TS}}
             ) 
             SELECT {{COLUMNS}} FROM distinct_prior_auths {prior_auth}
             LEFT JOIN {IDR_PROVIDER_HISTORY_TABLE} {prvdr_att_phy}

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

@@ -57,14 +57,19 @@ def model_type() -> ModelType:
 
     @override
     @staticmethod
-    def should_fully_sync_delete_diff() -> bool:
+    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:
@@ -75,7 +80,7 @@ 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 > {{LAST_TS}}
+                WHERE pa_req_rec_dt > {{MIN_TS}}
             )
             SELECT {{COLUMNS}} FROM distinct_prior_auths {prior_auth}
             WHERE row_order = 1;

apps/bfd-pipeline-idr/pipeline_utils.py

@@ -89,7 +89,7 @@ def extract_and_load(
 
             data_iter = (
                 data_extractor.extract_full_idr_data(source)
-                if cls.should_fully_sync_delete_diff()
+                if cls.should_delete_missing()
                 else data_extractor.extract_idr_data(progress, job_start, source)
             )
             res = loader.load(