BFD-4167: Initial IDR pipeline terraservice and fixes for prod data

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

@@ -34,6 +34,7 @@ CREATE TABLE idr.beneficiary(
 );
 
 CREATE INDEX ON idr.beneficiary(bene_mbi_id);
+CREATE INDEX ON idr.beneficiary(bene_xref_efctv_sk_computed);
 
 CREATE TABLE idr.beneficiary_history(
     bene_sk BIGINT NOT NULL,
@@ -148,10 +149,10 @@ CREATE TABLE idr.beneficiary_xref (
     bene_kill_cred_cd VARCHAR(1) NOT NULL,
     idr_insrt_ts TIMESTAMPTZ NOT NULL,
     idr_updt_ts TIMESTAMPTZ NOT NULL,
-    src_rec_ctre_ts TIMESTAMPTZ NOT NULL,
+    src_rec_crte_ts TIMESTAMPTZ NOT NULL,
     bfd_created_ts TIMESTAMPTZ NOT NULL,
     bfd_updated_ts TIMESTAMPTZ NOT NULL,
-    PRIMARY KEY(bene_sk, bene_hicn_num, src_rec_ctre_ts)
+    PRIMARY KEY(bene_sk, bene_hicn_num, src_rec_crte_ts)
 );
 
 CREATE TABLE idr.contract_pbp_number (
@@ -166,7 +167,8 @@ CREATE TABLE idr.load_progress(
     id INT GENERATED ALWAYS AS IDENTITY,
     table_name TEXT NOT NULL UNIQUE,
     last_ts TIMESTAMPTZ NOT NULL,
-    batch_completion_ts TIMESTAMPTZ NOT NULL
+    batch_start_ts TIMESTAMPTZ NOT NULL,
+    batch_complete_ts TIMESTAMPTZ NOT NULL
 );
 
 CREATE TABLE idr.claim (
@@ -205,7 +207,9 @@ CREATE TABLE idr.claim (
     clm_rndrg_prvdr_npi_num VARCHAR(10) NOT NULL,
     clm_rndrg_prvdr_last_name VARCHAR(60) NOT NULL,
     prvdr_blg_prvdr_npi_num VARCHAR(10) NOT NULL,
+    prvdr_rfrg_prvdr_npi_num VARCHAR(10) NOT NULL,
     clm_disp_cd VARCHAR(2) NOT NULL,
+    clm_ric_cd VARCHAR(1) NOT NULL,
     clm_sbmt_chrg_amt NUMERIC NOT NULL,
     clm_blood_pt_frnsh_qty INT NOT NULL,
     clm_nch_prmry_pyr_cd VARCHAR(1) NOT NULL,
@@ -381,3 +385,32 @@ HAVING COUNT(DISTINCT bene_xref_efctv_sk) > 1;
 
 -- required to refresh view with CONCURRENTLY
 CREATE UNIQUE INDEX ON idr.overshare_mbis (bene_mbi_id);
+
+CREATE OR REPLACE FUNCTION idr.refresh_overshare_mbis()
+    RETURNS VOID AS $$
+    DECLARE comment_sql TEXT;
+BEGIN
+    -- Using "concurrently" will make the refresh slower, but it will not block any reads
+    -- on the view while the refresh is in progress
+    REFRESH MATERIALIZED VIEW CONCURRENTLY idr.overshare_mbis;
+    -- There's no implicit way to know when a materialized view was last updated
+    -- add a comment on the object in case we need to verify that it's being updated as expected
+    comment_sql := 'COMMENT ON MATERIALIZED VIEW idr.overshare_mbis is '
+        || quote_literal('{"last_refreshed": "' || now() || '"}');
+    EXECUTE comment_sql;
+END;
+$$
+LANGUAGE plpgsql
+
+-- Only the owner of the view may refresh it, we need to set "security definer" so the function
+-- can execute in the context of the creator
+SECURITY DEFINER;
+-- search_path is the order in which schemas are searched when a name is referenced with no schema specified
+-- Postgres recommends setting this on functions marked as "security definer" to prevent malicious users from
+-- creating an object that shadows an existing one on a globally writable schema
+ALTER FUNCTION idr.refresh_overshare_mbis() SET search_path = idr;
+-- Execute privilege is granted to PUBLIC by default
+REVOKE ALL ON FUNCTION idr.refresh_overshare_mbis() FROM PUBLIC;
+
+-- This only needs to be executed by the pipeline
+-- GRANT EXECUTE ON FUNCTION idr.refresh_overshare_mbis() TO api_pipeline_svcs;

apps/bfd-pipeline/bfd-pipeline-idr/extractor.py

@@ -1,14 +1,18 @@
+import logging
 import os
+import time
 from abc import ABC, abstractmethod
 from collections.abc import Iterator, Mapping
 from datetime import date, datetime
 
 import psycopg
 import snowflake.connector
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
 from psycopg.rows import class_row
-from snowflake.connector import DictCursor
+from snowflake.connector import DictCursor, ProgrammingError, SnowflakeConnection
+from snowflake.connector.network import ReauthenticationRequest, RetryRequest
 
-from constants import DEFAULT_MAX_DATE
 from model import LoadProgress, T
 from timer import Timer
 
@@ -17,6 +21,8 @@
 cursor_fetch_timer = Timer("cursor_fetch")
 transform_timer = Timer("transform")
 
+logger = logging.getLogger(__name__)
+
 type DbType = str | float | int | bool | date | datetime
 
 
@@ -39,17 +45,21 @@ class Extractor(ABC):
     def extract_many(self, cls: type[T], sql: str, params: dict[str, DbType]) -> Iterator[list[T]]:
         pass
 
-    def get_query(self, cls: type[T], is_historical: bool) -> str:
-        query = cls.fetch_query(is_historical)
+    def get_query(self, cls: type[T], is_historical: bool, start_time: datetime) -> str:
+        query = cls.fetch_query(is_historical, start_time)
         columns = ",".join(cls.column_aliases())
         columns_raw = ",".join(cls.columns_raw())
         return query.replace("{COLUMNS}", columns).replace("{COLUMNS_NO_ALIAS}", columns_raw)
 
-    def extract_idr_data(self, cls: type[T], progress: LoadProgress | None) -> Iterator[list[T]]:
+    def extract_idr_data(
+        self, cls: type[T], progress: LoadProgress | None, start_time: datetime
+    ) -> Iterator[list[T]]:
         is_historical = progress is None or progress.is_historical()
-        fetch_query = self.get_query(cls, is_historical)
+        fetch_query = self.get_query(cls, is_historical, start_time)
         batch_timestamp_col = cls.batch_timestamp_col_alias(is_historical)
         update_timestamp_col = cls.update_timestamp_col_alias()
+
+        logger.info("extracting %s", cls.table())
         if progress is None:
             idr_query_timer.start()
             # No saved progress, process the whole table from the beginning
@@ -64,12 +74,14 @@ def extract_idr_data(self, cls: type[T], progress: LoadProgress | None) -> Itera
             idr_query_timer.stop()
             return res
 
-        previous_batch_complete = progress.batch_completion_ts != DEFAULT_MAX_DATE
-        op = ">" if previous_batch_complete else ">="
+        previous_batch_complete = progress.batch_complete_ts >= progress.batch_start_ts
+        compare_timestamp = progress.batch_start_ts if previous_batch_complete else progress.last_ts
+
         idr_query_timer.start()
-        # Saved progress found, start processing from where we left
+        # Saved progress found, start processing from where we left off
         update_clause = (
-            f"OR {update_timestamp_col} IS NOT NULL AND {update_timestamp_col} {op} %(timestamp)s"
+            f"""AND ({update_timestamp_col} IS NULL 
+                OR {update_timestamp_col} >= %(timestamp)s)"""
             if update_timestamp_col is not None
             else ""
         )
@@ -78,13 +90,15 @@ def extract_idr_data(self, cls: type[T], progress: LoadProgress | None) -> Itera
             fetch_query.replace(
                 "{WHERE_CLAUSE}",
                 f"""
-                    WHERE 
-                        ({update_timestamp_col} IS NOT NULL 
-                            AND {batch_timestamp_col} {op} %(timestamp)s {update_clause})
-                        AND {batch_timestamp_col} {op} '{get_min_transaction_date()}' 
+                    WHERE
+                        (
+                            {batch_timestamp_col} >= %(timestamp)s
+                            {update_clause}
+                        )
+                        AND {batch_timestamp_col} >= '{get_min_transaction_date()}' 
                     """,
             ).replace("{ORDER_BY}", f"ORDER BY {batch_timestamp_col}"),
-            {"timestamp": progress.last_ts},
+            {"timestamp": compare_timestamp},
         )
         idr_query_timer.stop()
         return res
@@ -115,40 +129,64 @@ def extract_single(self, cls: type[T], sql: str, params: dict[str, DbType]) -> T
 class SnowflakeExtractor(Extractor):
     def __init__(self, batch_size: int) -> None:
         super().__init__()
-        self.conn = snowflake.connector.connect(  # type: ignore
+
+        self.conn = SnowflakeExtractor._connect()
+        self.batch_size = batch_size
+
+    @staticmethod
+    def _connect() -> SnowflakeConnection:
+        private_key = serialization.load_pem_private_key(
+            os.environ["IDR_PRIVATE_KEY"].encode(), password=None, backend=default_backend()
+        )
+        private_key_bytes = private_key.private_bytes(
+            encoding=serialization.Encoding.DER,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+        return snowflake.connector.connect(  # type: ignore
             user=os.environ["IDR_USERNAME"],
-            password=os.environ["IDR_PASSWORD"],
+            private_key=private_key_bytes,
             account=os.environ["IDR_ACCOUNT"],
             warehouse=os.environ["IDR_WAREHOUSE"],
             database=os.environ["IDR_DATABASE"],
             schema=os.environ["IDR_SCHEMA"],
         )
-        self.batch_size = batch_size
 
     def extract_many(self, cls: type[T], sql: str, params: dict[str, DbType]) -> Iterator[list[T]]:
         cur = None
-        try:
-            cursor_execute_timer.start()
-            cur = self.conn.cursor(DictCursor)
-            cur.execute(sql, params)
-            cursor_execute_timer.stop()
-
-            cursor_fetch_timer.start()
-            # fetchmany can return list[dict] or list[tuple] but we'll only use
-            # queries that return dicts
-            batch: list[dict[str, DbType]] = cur.fetchmany(self.batch_size)  # type: ignore[assignment]
-            cursor_fetch_timer.stop()
-
-            while len(batch) > 0:  # type: ignore
-                transform_timer.start()
-                data = [cls(**{k.lower(): v for k, v in row.items()}) for row in batch]
-                transform_timer.stop()
-
-                yield data
+        max_attempts = 5
+        for attempt in range(max_attempts):
+            try:
+                cursor_execute_timer.start()
+                cur = self.conn.cursor(DictCursor)
+                cur.execute(sql, params)
+                cursor_execute_timer.stop()
 
                 cursor_fetch_timer.start()
-                batch = cur.fetchmany(self.batch_size)  # type: ignore[assignment]
+                # fetchmany can return list[dict] or list[tuple] but we'll only use
+                # queries that return dicts
+                batch: list[dict[str, DbType]] = cur.fetchmany(self.batch_size)  # type: ignore[assignment]
                 cursor_fetch_timer.stop()
-        finally:
-            if cur:
-                cur.close()
+
+                while len(batch) > 0:  # type: ignore
+                    transform_timer.start()
+                    data = [cls(**{k.lower(): v for k, v in row.items()}) for row in batch]
+                    transform_timer.stop()
+
+                    yield data
+
+                    cursor_fetch_timer.start()
+                    batch = cur.fetchmany(self.batch_size)  # type: ignore[assignment]
+                    cursor_fetch_timer.stop()
+                return
+            except (ReauthenticationRequest, RetryRequest, ProgrammingError) as ex:
+                logger.warning("received transient error, retrying...", exc_info=ex)
+                if attempt == max_attempts - 1:
+                    logger.error("max attempts exceeded")
+                    raise ex
+                self.conn = SnowflakeExtractor._connect()
+                time.sleep(1)
+
+            finally:
+                if cur:
+                    cur.close()

apps/bfd-pipeline/bfd-pipeline-idr/load-credentials.sh

@@ -1,23 +1,23 @@
 #!/usr/bin/env bash
 
-IDR_USERNAME="$(aws ssm get-parameter --name /bfd/3913-prod/pipeline/sensitive/idr_username --with-decryption --query "Parameter.Value" --output text)"
+IDR_USERNAME="$(aws ssm get-parameter --name /bfd/${BFD_ENV}/idr-pipeline/sensitive/idr_username --with-decryption --query "Parameter.Value" --output text)"
 export IDR_USERNAME
-IDR_PASSWORD="$(aws ssm get-parameter --name /bfd/3913-prod/pipeline/sensitive/idr_password --with-decryption --query "Parameter.Value" --output text)"
-export IDR_PASSWORD
-IDR_ACCOUNT="$(aws ssm get-parameter --name /bfd/3913-prod/pipeline/sensitive/idr_account --with-decryption --query "Parameter.Value" --output text)"
+IDR_PRIVATE_KEY="$(aws ssm get-parameter --name /bfd/${BFD_ENV}/idr-pipeline/sensitive/idr_private_key --with-decryption --query "Parameter.Value" --output text)"
+export IDR_PRIVATE_KEY
+IDR_ACCOUNT="$(aws ssm get-parameter --name /bfd/${BFD_ENV}/idr-pipeline/sensitive/idr_account --with-decryption --query "Parameter.Value" --output text)"
 export IDR_ACCOUNT
-IDR_WAREHOUSE="$(aws ssm get-parameter --name /bfd/3913-prod/pipeline/sensitive/idr_warehouse --with-decryption --query "Parameter.Value" --output text)"
+IDR_WAREHOUSE="$(aws ssm get-parameter --name /bfd/${BFD_ENV}/idr-pipeline/sensitive/idr_warehouse --with-decryption --query "Parameter.Value" --output text)"
 export IDR_WAREHOUSE
-IDR_DATABASE="$(aws ssm get-parameter --name /bfd/3913-prod/pipeline/sensitive/idr_database --with-decryption --query "Parameter.Value" --output text)"
+IDR_DATABASE="$(aws ssm get-parameter --name /bfd/${BFD_ENV}/idr-pipeline/sensitive/idr_database --with-decryption --query "Parameter.Value" --output text)"
 export IDR_DATABASE
-IDR_SCHEMA="$(aws ssm get-parameter --name /bfd/3913-prod/pipeline/sensitive/idr_schema --with-decryption --query "Parameter.Value" --output text)"
+IDR_SCHEMA="$(aws ssm get-parameter --name /bfd/${BFD_ENV}/idr-pipeline/sensitive/idr_schema --with-decryption --query "Parameter.Value" --output text)"
 export IDR_SCHEMA
 
-BFD_DB_USERNAME="$(aws ssm get-parameter --name /bfd/3913-prod/pipeline/sensitive/db/username --with-decryption --query "Parameter.Value" --output text)"
+BFD_DB_USERNAME="$(aws ssm get-parameter --name /bfd/${BFD_ENV}/idr-pipeline/sensitive/db/username --with-decryption --query "Parameter.Value" --output text)"
 export BFD_DB_USERNAME
-BFD_DB_PASSWORD="$(aws ssm get-parameter --name /bfd/3913-prod/pipeline/sensitive/db/password --with-decryption --query "Parameter.Value" --output text)"
+BFD_DB_PASSWORD="$(aws ssm get-parameter --name /bfd/${BFD_ENV}/idr-pipeline/sensitive/db/password --with-decryption --query "Parameter.Value" --output text)"
 export BFD_DB_PASSWORD
 
-db_cluster="$(aws ssm get-parameter --name /bfd/3913-prod/common/nonsensitive/rds_cluster_identifier --with-decryption --query "Parameter.Value" --output text)"
+db_cluster="bfd-${BFD_ENV}-aurora-cluster"
 BFD_DB_ENDPOINT="$(aws rds describe-db-clusters --db-cluster-identifier $db_cluster --query "DBClusters[0].Endpoint" --output text)"
 export BFD_DB_ENDPOINT