Changes for staging data and testing for non part D claim pruning,

naming refactor for accuracy,
refactor query logic into base_model.py to match current pruning

Author: patrickgmayday

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

@@ -865,3 +865,44 @@ def stale_phase_1_claims_query(
         """,
         (cutoff_date, cutoff_date),
     )
+
+
+def non_latest_non_part_d_claim_items_query(
+    item_table: str,
+    cutoff_date: datetime,
+) -> tuple[str, tuple[datetime]]:
+    part_d_codes = ",".join(str(code) for code in PART_D_CLAIM_TYPE_CODES)
+
+    return (
+        f"""
+            SELECT item.clm_uniq_id, item.bfd_row_id
+            FROM {item_table} item
+            JOIN {IDR_CLAIM_TABLE} clm ON clm.clm_uniq_id = item.clm_uniq_id
+            WHERE clm.clm_ltst_clm_ind = 'N'
+            AND clm.clm_type_cd NOT IN ({part_d_codes})
+            AND clm.clm_idr_ld_dt < %s
+            ORDER BY item.clm_uniq_id, item.bfd_row_id
+            LIMIT {PHASE_1_PRUNE_BATCH_LIMIT}
+        """,
+        (cutoff_date,),
+    )
+
+
+def non_latest_non_part_d_parent_claims_query(
+    claim_table: str,
+    cutoff_date: datetime,
+) -> tuple[str, tuple[datetime]]:
+    part_d_codes = ",".join(str(code) for code in PART_D_CLAIM_TYPE_CODES)
+
+    return (
+        f"""
+            SELECT clm.clm_uniq_id
+            FROM {claim_table} clm
+            WHERE clm.clm_ltst_clm_ind = 'N'
+            AND clm.clm_type_cd NOT IN ({part_d_codes})
+            AND clm.clm_idr_ld_dt < %s
+            ORDER BY clm.clm_uniq_id
+            LIMIT {PHASE_1_PRUNE_BATCH_LIMIT}
+        """,
+        (cutoff_date,),
+    )

apps/bfd-pipeline-idr/pipeline_stages.py

@@ -41,8 +41,8 @@
     prune_bene_lis_cmbnd,
     prune_bene_ma_part_d,
     prune_bene_ma_part_d_rx,
-    prune_non_latest_part_d_ss_claim_items,
-    prune_non_latest_part_d_ss_parent_claims,
+    prune_non_latest_non_part_d_ss_claim_items,
+    prune_non_latest_non_part_d_ss_parent_claims,
     prune_phase_1_ss_claims,
 )
 from settings import enable_prior_auth_ingestion
@@ -178,13 +178,13 @@ def _stage5_prune_obsolete_rows(self) -> Stage[bool]:
                 self.start_time,
             )
             yield functools.partial(
-                prune_non_latest_part_d_ss_claim_items,
+                prune_non_latest_non_part_d_ss_claim_items,
                 model,
                 self.load_mode,
                 self.start_time,
             )
             yield functools.partial(
-                prune_non_latest_part_d_ss_parent_claims,
+                prune_non_latest_non_part_d_ss_parent_claims,
                 model,
                 self.load_mode,
                 self.start_time,

apps/bfd-pipeline-idr/pipeline_utils.py

@@ -15,8 +15,6 @@
     CLAIM_PROFESSIONAL_SS_TABLE,
     DEFAULT_MAX_DATE,
     DEFAULT_PARTITION,
-    IDR_CLAIM_TABLE,
-    PART_D_CLAIM_TYPE_CODES,
     PHASE_1_CUTOFF,
 )
 from extractor import PostgresExtractor, SnowflakeExtractor, Source
@@ -25,6 +23,8 @@
 from model.base_model import (
     LoadMode,
     T,
+    non_latest_non_part_d_claim_items_query,
+    non_latest_non_part_d_parent_claims_query,
     stale_phase_1_claims_query,
 )
 from model.idr_beneficiary_low_income_subsidy_cmbnd import IdrBeneficiaryLowIncomeSubsidyCmbnd
@@ -171,7 +171,7 @@ def prune_phase_1_ss_claims(
     return True
 
 
-def prune_non_latest_part_d_ss_claim_items(
+def prune_non_latest_non_part_d_ss_claim_items(
     cls: type[T],
     load_mode: LoadMode,
     job_start: datetime,
@@ -182,37 +182,29 @@ def prune_non_latest_part_d_ss_claim_items(
         return True
 
     prune_cutoff_date = job_start - timedelta(days=PHASE_1_CUTOFF)
-    part_d_codes = ",".join(str(code) for code in PART_D_CLAIM_TYPE_CODES)
 
     logger.info("pruning non-latest non-Part-D ss claim items older than {}", prune_cutoff_date)
 
-    with psycopg.connect(get_connection_string(load_mode)) as conn, conn.transaction():
+    prune_query, params = non_latest_non_part_d_claim_items_query(item_table, prune_cutoff_date)
+
+    with psycopg.connect(get_connection_string(load_mode)) as conn:
         while True:
-            # Claim items can exist even when the non-latest parent claim was filtered
-            # before final claim-table load, so use the source claim table.
-            res = conn.execute(
-                f"""
-                DELETE FROM {item_table}
-                WHERE (clm_uniq_id, bfd_row_id) IN (
-                    SELECT item.clm_uniq_id, item.bfd_row_id
-                    FROM {item_table} item
-                    JOIN {IDR_CLAIM_TABLE} clm ON clm.clm_uniq_id = item.clm_uniq_id
-                    WHERE clm.clm_ltst_clm_ind = 'N'
-                    AND clm.clm_type_cd NOT IN ({part_d_codes})
-                    AND clm.clm_idr_ld_dt < %s
-                    LIMIT %s
+            with conn.transaction():
+                res = conn.execute(
+                    f"""
+                    DELETE FROM {item_table}
+                    WHERE (clm_uniq_id, bfd_row_id) IN ({prune_query})
+                    """,  # type: ignore
+                    params,
                 )
-                """,  # type: ignore
-                (prune_cutoff_date, PHASE_1_PRUNE_BATCH_LIMIT),
-            )
-            logger.info("pruned {} rows from {}", res.rowcount, item_table)
-            if res.rowcount < PHASE_1_PRUNE_BATCH_LIMIT:
-                break
+                logger.info("pruned {} rows from {}", res.rowcount, item_table)
+                if res.rowcount < PHASE_1_PRUNE_BATCH_LIMIT:
+                    break
 
     return True
 
 
-def prune_non_latest_part_d_ss_parent_claims(
+def prune_non_latest_non_part_d_ss_parent_claims(
     cls: type[T],
     load_mode: LoadMode,
     job_start: datetime,
@@ -222,28 +214,24 @@ def prune_non_latest_part_d_ss_parent_claims(
         return True
 
     prune_cutoff_date = job_start - timedelta(days=PHASE_1_CUTOFF)
-    part_d_codes = ",".join(str(code) for code in PART_D_CLAIM_TYPE_CODES)
 
     logger.info("pruning non-latest non-Part-D ss parent claims older than {}", prune_cutoff_date)
 
-    with psycopg.connect(get_connection_string(load_mode)) as conn, conn.transaction():
+    prune_query, params = non_latest_non_part_d_parent_claims_query(claim_table, prune_cutoff_date)
+
+    with psycopg.connect(get_connection_string(load_mode)) as conn:
         while True:
-            res = conn.execute(
-                f"""
-                DELETE FROM {claim_table}
-                WHERE clm_uniq_id IN (
-                    SELECT clm.clm_uniq_id FROM {claim_table} clm
-                    WHERE clm.clm_ltst_clm_ind = 'N'
-                    AND clm.clm_type_cd NOT IN ({part_d_codes})
-                    AND clm.clm_idr_ld_dt < %s
-                    LIMIT %s
+            with conn.transaction():
+                res = conn.execute(
+                    f"""
+                    DELETE FROM {claim_table}
+                    WHERE clm_uniq_id IN ({prune_query})
+                    """,  # type: ignore
+                    params,
                 )
-                """,  # type: ignore
-                (prune_cutoff_date, PHASE_1_PRUNE_BATCH_LIMIT),
-            )
-            logger.info("pruned {} rows from {}", res.rowcount, claim_table)
-            if res.rowcount < PHASE_1_PRUNE_BATCH_LIMIT:
-                break
+                logger.info("pruned {} rows from {}", res.rowcount, claim_table)
+                if res.rowcount < PHASE_1_PRUNE_BATCH_LIMIT:
+                    break
 
     return True
 

apps/bfd-pipeline-idr/test_pipeline.py

@@ -93,6 +93,42 @@ def _do_test_pipeline(conn: Connection[DictRow], load_type: LoadType) -> None:
         rows = cur.fetchmany(1)
         assert rows[0]["mbi_num"] == "1OX4Y88RV68"
 
+    # Seed a parent claim so the prune job has a row to delete.
+    # CSV covers item pruning because non-latest parents do not load.
+    if load_type == LoadType.INCREMENTAL:
+        claim_table = sql.Identifier("idr", "claim_institutional_ss")
+        cur = conn.execute(
+            "select * from idr.claim_institutional_ss where clm_uniq_id = 123359318723"
+        )
+        assert cur.rowcount == 1
+        row = cur.fetchone()
+        assert row is not None
+
+        non_latest_parent_claim = dict(row)
+        non_latest_parent_claim["clm_uniq_id"] = 999999434801
+        non_latest_parent_claim["clm_ltst_clm_ind"] = "N"
+        non_latest_parent_claim["clm_type_cd"] = 2081
+        # Make the seeded claim old enough to be picked up by the prune job.
+        non_latest_parent_claim["clm_idr_ld_dt"] = datetime(2019, 6, 13, tzinfo=UTC).date()
+
+        if "bfd_row_id" in non_latest_parent_claim:
+            non_latest_parent_claim["bfd_row_id"] = 999999434801
+
+        if "clm_dt_sgntr_sk" in non_latest_parent_claim:
+            non_latest_parent_claim["clm_dt_sgntr_sk"] = 999999434801
+
+        conn.execute(
+            t"""
+            INSERT INTO {claim_table:i} (
+                {sql.SQL(", ").join(sql.Identifier(k) for k in non_latest_parent_claim):q}
+            )
+            VALUES (
+                {sql.SQL(", ").join(non_latest_parent_claim.values()):q}
+            )
+            """
+        )
+        conn.commit()
+
     cur = conn.execute("select max(last_ts) as max_ts from idr.load_progress")
     row = cur.fetchone()
     assert row is not None
@@ -215,22 +251,17 @@ def _do_test_pipeline(conn: Connection[DictRow], load_type: LoadType) -> None:
     rows = cur.fetchmany(1)
     assert rows[0]["clm_uniq_id"] == 113370100080
 
-    # Phase 1 SS (PAC) claims older than 60 days will be pruned on incremental loads
-    if load_type == LoadType.INITIAL:
-        cur = conn.execute("select * from idr.claim_institutional_ss order by clm_uniq_id")
-        assert cur.rowcount == 21
-        rows = cur.fetchmany(1)
-        assert rows[0]["clm_uniq_id"] == 123359318723
-    else:
-        cur = conn.execute("select * from idr.claim_institutional_ss order by clm_uniq_id")
-        assert cur.rowcount == 9
-        rows = cur.fetchmany(1)
-        assert rows[0]["clm_uniq_id"] == 849348853948
-
     # Non-latest non-Part-D SS parent claims are filtered before final claim-table load
     cur = conn.execute("select * from idr.claim_institutional_ss where clm_uniq_id = 999999434800")
     assert cur.rowcount == 0
 
+    # Existing non-latest non-Part-D SS parent claims are pruned on incremental loads
+    if load_type == LoadType.INCREMENTAL:
+        cur = conn.execute(
+            "select * from idr.claim_institutional_ss where clm_uniq_id = 999999434801"
+        )
+        assert cur.rowcount == 0
+
     cur = conn.execute("select * from idr.claim_professional_nch order by clm_uniq_id")
     assert cur.rowcount == 51
     rows = cur.fetchmany(1)
@@ -251,18 +282,6 @@ def _do_test_pipeline(conn: Connection[DictRow], load_type: LoadType) -> None:
     rows = cur.fetchmany(1)
     assert rows[0]["clm_uniq_id"] == 113370100080
 
-    # Phase 1 SS (PAC) claims older than 60 days will be pruned on incremental loads
-    if load_type == LoadType.INITIAL:
-        cur = conn.execute("select * from idr.claim_item_institutional_ss order by clm_uniq_id")
-        assert cur.rowcount == 328
-        rows = cur.fetchmany(1)
-        assert rows[0]["clm_uniq_id"] == 123359318723
-    else:
-        cur = conn.execute("select * from idr.claim_item_institutional_ss order by clm_uniq_id")
-        assert cur.rowcount == 151
-        rows = cur.fetchmany(1)
-        assert rows[0]["clm_uniq_id"] == 849348853948
-
     # Items for non-latest non-Part-D SS claims are pruned on incremental loads
     cur = conn.execute(
         "select * from idr.claim_item_institutional_ss where clm_uniq_id = 999999434800"