Merge branch 'master' into BFD-4764

Author: alex-dzeda

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

@@ -0,0 +1,25 @@
+DROP VIEW idr.beneficiary_part_c_and_d_enrollment;
+CREATE VIEW idr.beneficiary_part_c_and_d_enrollment AS
+SELECT
+    e.bene_sk,
+    e.bene_enrlmt_pgm_type_cd,
+    e.bene_enrlmt_bgn_dt,
+    e.bene_enrlmt_end_dt,
+    e.bene_cntrct_num,
+    e.bene_pbp_num,
+    e.cntrct_pbp_sk,
+    e.bene_cvrg_type_cd,
+    e.bene_enrlmt_emplr_sbsdy_sw,
+    COALESCE(rx.bene_enrlmt_pdp_rx_info_bgn_dt, DATE '9999-12-31') AS bene_enrlmt_pdp_rx_info_bgn_dt,
+    rx.bene_pdp_enrlmt_mmbr_id_num,
+    rx.bene_pdp_enrlmt_grp_num,
+    rx.bene_pdp_enrlmt_prcsr_num,
+    rx.bene_pdp_enrlmt_bank_id_num
+FROM idr.beneficiary_ma_part_d_enrollment e
+     LEFT JOIN idr.beneficiary_ma_part_d_enrollment_rx rx
+       ON e.bene_sk = rx.bene_sk
+       AND e.bene_enrlmt_bgn_dt = rx.bene_enrlmt_bgn_dt
+       AND e.cntrct_pbp_sk = rx.cntrct_pbp_sk
+       AND e.bene_enrlmt_pgm_type_cd in ('2', '3')
+       AND rx.idr_trans_obslt_ts >= '9999-12-31'
+WHERE e.idr_trans_obslt_ts >= '9999-12-31';

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

@@ -0,0 +1 @@
+CREATE INDEX ON idr.beneficiary_ma_part_d_enrollment(idr_trans_obslt_ts) WHERE idr_trans_obslt_ts < '9999-12-31';

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

@@ -0,0 +1 @@
+CREATE INDEX ON idr.beneficiary_ma_part_d_enrollment_rx(idr_trans_obslt_ts) WHERE idr_trans_obslt_ts < '9999-12-31';

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

@@ -3,10 +3,7 @@
 
 from pydantic import BeforeValidator
 
-from constants import (
-    DEFAULT_MAX_DATE,
-    IDR_BENE_MA_PART_D_TABLE,
-)
+from constants import IDR_BENE_MA_PART_D_TABLE
 from load_partition import LoadPartition
 from model.base_model import (
     ALIAS_HSTRY,
@@ -76,7 +73,6 @@ def fetch_query(cls, partition: LoadPartition, start_time: datetime, source: Sou
                 {deceased_bene_filter(hstry, start_time)}
                 AND {hstry}.bene_sk = enrlmt.bene_sk
             )
-            AND idr_trans_obslt_ts >= '{DEFAULT_MAX_DATE}'
             AND bene_enrlmt_pgm_type_cd != '~'
             {{ORDER_BY}}
             {{LIMIT}}

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

@@ -3,10 +3,7 @@
 
 from pydantic import BeforeValidator
 
-from constants import (
-    DEFAULT_MAX_DATE,
-    IDR_BENE_MA_PART_D_RX_TABLE,
-)
+from constants import IDR_BENE_MA_PART_D_RX_TABLE
 from load_partition import LoadPartition
 from model.base_model import (
     ALIAS_HSTRY,
@@ -70,7 +67,6 @@ def fetch_query(cls, partition: LoadPartition, start_time: datetime, source: Sou
                     {deceased_bene_filter(hstry, start_time)}
                     AND {hstry}.bene_sk = enrlmt_rx.bene_sk
                 )
-                AND idr_trans_obslt_ts >= '{DEFAULT_MAX_DATE}'
                 {{ORDER_BY}}
                 {{LIMIT}}
             """

apps/bfd-pipeline-idr/pipeline_stages.py

@@ -36,7 +36,13 @@
 from model.idr_prior_auth import IdrPriorAuth
 from model.idr_prior_auth_item import IdrPriorAuthItem
 from parallel_executor import ParallelStagesExecutor, Stage
-from pipeline_utils import extract_and_load, prune_bene_lis_cmbnd, prune_phase_1_ss_claims
+from pipeline_utils import (
+    extract_and_load,
+    prune_bene_lis_cmbnd,
+    prune_bene_ma_part_d,
+    prune_bene_ma_part_d_rx,
+    prune_phase_1_ss_claims,
+)
 from settings import enable_prior_auth_ingestion
 
 type NodePartitionedModelInput = tuple[type[IdrBaseModel], LoadPartition | None]
@@ -107,7 +113,7 @@ async def start(self) -> bool:
                         self._stage2_do_claims_and_benes_tbls(),
                         self._stage3_do_parent_claims_tbls(),
                         self._stage4_do_beneficiary(),
-                        self._stage5_do_phase_1_prune(),
+                        self._stage5_prune_obsolete_rows(),
                     ],
                 )
             )
@@ -143,7 +149,7 @@ def _stage4_do_beneficiary(self) -> Stage[bool]:
             self._gen_partitioned_node_inputs(self._filter_tables(BENE_TABLES))
         )
 
-    def _stage5_do_phase_1_prune(self) -> Stage[bool]:
+    def _stage5_prune_obsolete_rows(self) -> Stage[bool]:
         if self.load_type == LoadType.INITIAL:
             return
 
@@ -152,6 +158,16 @@ def _stage5_do_phase_1_prune(self) -> Stage[bool]:
             self.load_mode,
         )
 
+        yield functools.partial(
+            prune_bene_ma_part_d,
+            self.load_mode,
+        )
+
+        yield functools.partial(
+            prune_bene_ma_part_d_rx,
+            self.load_mode,
+        )
+
         for model in self._filter_tables(CLAIM_SS_TABLES):
             yield functools.partial(
                 prune_phase_1_ss_claims,

apps/bfd-pipeline-idr/pipeline_utils.py

@@ -26,8 +26,10 @@
     stale_phase_1_claims_query,
 )
 from model.idr_beneficiary_low_income_subsidy_cmbnd import IdrBeneficiaryLowIncomeSubsidyCmbnd
+from model.idr_beneficiary_ma_part_d_enrollment import IdrBeneficiaryMaPartDEnrollment
+from model.idr_beneficiary_ma_part_d_enrollment_rx import IdrBeneficiaryMaPartDEnrollmentRx
 from model.load_progress import LoadProgress
-from settings import BENEFICIARY_PRUNE_BATCH_LIMIT
+from settings import BENEFICIARY_PART_D_PRUNE_BATCH_LIMIT, BENEFICIARY_PRUNE_BATCH_LIMIT
 
 
 def get_progress(
@@ -189,3 +191,68 @@ def prune_bene_lis_cmbnd(
                 break
 
     return True
+
+
+def prune_bene_ma_part_d(
+    load_mode: LoadMode,
+) -> bool:
+    bene_table = IdrBeneficiaryMaPartDEnrollment.table()
+
+    logger.info("pruning obsolete part d beneficiaries", DEFAULT_MAX_DATE)
+
+    with psycopg.connect(get_connection_string(load_mode)) as conn, conn.transaction():
+        while True:
+            res = conn.execute(
+                f"""
+                DELETE FROM {bene_table}
+                WHERE (bene_sk, bene_enrlmt_bgn_dt, bene_enrlmt_pgm_type_cd) IN (
+                    SELECT bene_sk, bene_enrlmt_bgn_dt, bene_enrlmt_pgm_type_cd
+                    FROM {bene_table}
+                    WHERE idr_trans_obslt_ts < %s
+                    LIMIT %s
+                )
+                """,  # type: ignore
+                (DEFAULT_MAX_DATE, BENEFICIARY_PART_D_PRUNE_BATCH_LIMIT),
+            )
+            logger.info("pruned {} rows from {}", res.rowcount, bene_table)
+            if res.rowcount < BENEFICIARY_PART_D_PRUNE_BATCH_LIMIT:
+                break
+
+    return True
+
+
+def prune_bene_ma_part_d_rx(
+    load_mode: LoadMode,
+) -> bool:
+    bene_table = IdrBeneficiaryMaPartDEnrollmentRx.table()
+
+    logger.info("pruning obsolete part d rx beneficiaries", DEFAULT_MAX_DATE)
+
+    with psycopg.connect(get_connection_string(load_mode)) as conn, conn.transaction():
+        while True:
+            res = conn.execute(
+                f"""
+                DELETE FROM {bene_table}
+                WHERE (bene_sk, 
+                       bene_cntrct_num, 
+                       bene_pbp_num, 
+                       bene_enrlmt_bgn_dt, 
+                       bene_enrlmt_pdp_rx_info_bgn_dt
+                    ) IN (
+                    SELECT bene_sk, 
+                           bene_cntrct_num, 
+                           bene_pbp_num, 
+                           bene_enrlmt_bgn_dt, 
+                           bene_enrlmt_pdp_rx_info_bgn_dt
+                    FROM {bene_table}
+                    WHERE idr_trans_obslt_ts < %s
+                    LIMIT %s
+                )
+                """,  # type: ignore
+                (DEFAULT_MAX_DATE, BENEFICIARY_PART_D_PRUNE_BATCH_LIMIT),
+            )
+            logger.info("pruned {} rows from {}", res.rowcount, bene_table)
+            if res.rowcount < BENEFICIARY_PART_D_PRUNE_BATCH_LIMIT:
+                break
+
+    return True

apps/bfd-pipeline-idr/settings.py

@@ -100,6 +100,11 @@ def enable_prior_auth_ingestion() -> bool:
 """Number of minimum connections to hold in the pool concurrently per-batch for non-LOCAL loads.
 Defaults to 20."""
 
+BENEFICIARY_PART_D_PRUNE_BATCH_LIMIT = int(
+    getenv("IDR_BENEFICIARY_PART_D_PRUNE_BATCH_LIMIT", "100000")
+)
+"""Maximum rows to delete per prune statement for Part D beneficiary records."""
+
 PHASE_1_PRUNE_BATCH_LIMIT = int(getenv("PHASE_1_PRUNE_BATCH_LIMIT", "10_000"))
 """The maximum batch size for pruning old claims (phase 1 claims from shared systems) on
 INCREMENTAL loads. Defaults to 10000."""

apps/bfd-pipeline-idr/test_pipeline.py

@@ -168,15 +168,27 @@ def _do_test_pipeline(conn: Connection[DictRow], load_type: LoadType) -> None:
     # only a future record exists for this contract
     assert rows[6]["cntrct_pbp_bgn_dt"].strftime("%Y-%m-%d") == "2026-12-01"
 
-    cur = conn.execute("select * from idr.beneficiary_ma_part_d_enrollment order by bene_sk")
-    assert cur.rowcount == 3
-    rows = cur.fetchmany(1)
-    assert rows[0]["bene_sk"] == 353816020
+    if load_type == LoadType.INITIAL:
+        cur = conn.execute("select * from idr.beneficiary_ma_part_d_enrollment order by bene_sk")
+        assert cur.rowcount == 4
+        rows = cur.fetchmany(1)
+        assert rows[0]["bene_sk"] == 353816020
+    else:
+        cur = conn.execute("select * from idr.beneficiary_ma_part_d_enrollment order by bene_sk")
+        assert cur.rowcount == 3
+        rows = cur.fetchmany(1)
+        assert rows[0]["bene_sk"] == 353816020
 
-    cur = conn.execute("select * from idr.beneficiary_ma_part_d_enrollment_rx order by bene_sk")
-    assert cur.rowcount == 2
-    rows = cur.fetchmany(1)
-    assert rows[0]["bene_sk"] == 353816020
+    if load_type == LoadType.INITIAL:
+        cur = conn.execute("select * from idr.beneficiary_ma_part_d_enrollment_rx order by bene_sk")
+        assert cur.rowcount == 3
+        rows = cur.fetchmany(1)
+        assert rows[0]["bene_sk"] == 353816020
+    else:
+        cur = conn.execute("select * from idr.beneficiary_ma_part_d_enrollment_rx order by bene_sk")
+        assert cur.rowcount == 2
+        rows = cur.fetchmany(1)
+        assert rows[0]["bene_sk"] == 353816020
 
     cur = conn.execute("select * from idr.beneficiary_low_income_subsidy order by bene_sk")
     assert cur.rowcount == 2

apps/bfd-pipeline-idr/test_samples1/SYNTHETIC_BENE_MAPD_ENRLMT.csv

@@ -2,3 +2,4 @@ BENE_SK,CNTRCT_PBP_SK,IDR_LTST_TRANS_FLG,BENE_CNTRCT_NUM,BENE_PBP_NUM,BENE_CVRG_
 547437476,993028062567,Y,S0001,003,11,2,1,2019-02-18,9999-12-31,2019-02-18T00:00:00.000000,2019-02-18T00:00:00.000000,2019-11-11T00:00:00.000000,9999-12-31T00:00:00.000000
 441149422,408933975817,Y,G1234,002,3,1,~,2018-02-04,9999-12-31,2018-02-04T00:00:00.000000,2018-02-04T00:00:00.000000,2018-02-04T00:00:00.000000,9999-12-31T00:00:00.000000
 353816020,761144926385,Y,H1234,003,11,3,Y,2020-09-03,9999-12-31,2020-09-03T00:00:00.000000,2020-09-03T00:00:00.000000,2020-09-03T00:00:00.000000,9999-12-31T00:00:00.000000
+353816021,993028062568,N,T0001,005,11,2,1,2020-06-01,2023-01-31,2020-06-01T00:00:00.000000,2020-06-01T00:00:00.000000,2020-06-01T00:00:00.000000,2023-01-15T00:00:00.000000