[DEV-14868] DEFCGold to CSVModel

Author: dpb-bah

brus_backend_common/models/__init__.py

@@ -1,6 +1,6 @@
 from brus_backend_common.models.lakehouse_model import LakeHouseCurrentMigration, ExternalDataLoadDate
 
-from brus_backend_common.models.reference import DEFCSilver, DEFCGroup, DEFCBronze
+from brus_backend_common.models.reference import DEFCBronze, DEFCGold, DEFCGroup
 from brus_backend_common.config import CONFIG
 
 LAKEHOUSE_BUCKETS = {
@@ -13,7 +13,7 @@
 LAKEHOUSE_MODEL_CLASSES = [
     DEFCBronze,
     DEFCGroup,
-    DEFCSilver,
+    DEFCGold,
     LakeHouseCurrentMigration,
     ExternalDataLoadDate,
 ]

brus_backend_common/models/reference.py

@@ -8,7 +8,7 @@
     TimestampType,
 )
 from brus_backend_common.config import CONFIG
-from brus_backend_common.models.lakehouse_model import DeltaModel, CSVModel, LakeHouseDatabase
+from brus_backend_common.models.lakehouse_model import CSVModel, LakeHouseDatabase
 
 
 class DEFCBronze(CSVModel):
@@ -31,7 +31,7 @@ class DEFCBronze(CSVModel):
 
 class DEFCGroup(CSVModel):
     BUCKET_NAME = CONFIG.REFERENCE_S3_BUCKET
-    DATABASE_NAME = LakeHouseDatabase.SILVER
+    DATABASE_NAME = LakeHouseDatabase.GOLD
     TABLE_NAME = "defc_mapping"
     DESCRIPTION = "Internal CSV to dynamically group DEFCs together"
     CSV_NAME = "DEFC_MAPPING.csv"
@@ -47,15 +47,15 @@ class DEFCGroup(CSVModel):
     )
 
 
-class DEFCSilver(DeltaModel):
+class DEFCGold(CSVModel):
     BUCKET_NAME = CONFIG.REFERENCE_S3_BUCKET
-    DATABASE_NAME = LakeHouseDatabase.SILVER
+    DATABASE_NAME = LakeHouseDatabase.GOLD
     TABLE_NAME = "defc"
     DESCRIPTION = "DEFC data after initial processing"
+    CSV_NAME = "def_codes.csv"
     PK = "defc_id"
     UNIQUE_CONSTRAINTS = ["code"]
     MIGRATION_HISTORY = []
-
     STRUCTURE = StructType(
         [
             StructField("created_at", TimestampType(), True),

…kend_common/scripts/loaders/load_defc.py …kend_common/scripts/loaders/defc_gold.pybrus_backend_common/scripts/loaders/load_defc.py renamed to brus_backend_common/scripts/loaders/defc_gold.py brus_backend_common/scripts/loaders/load_defc.py renamed to brus_backend_common/scripts/loaders/defc_gold.py

@@ -8,10 +8,9 @@
 import re
 from datetime import datetime
 
-from brus_backend_common.models import DEFCBronze, DEFCSilver, DEFCGroup
+from brus_backend_common.models import DEFCBronze, DEFCGroup, DEFCGold
 from brus_backend_common.models.lakehouse_model import update_external_data_load_date
 from brus_backend_common.helpers.aws import _get_boto3
-from brus_backend_common.helpers.spark import SparkScriptSession
 from brus_backend_common.helpers.pandas import check_dataframe_diff
 from brus_backend_common.helpers.scripts import (
     clean_data,
@@ -203,107 +202,106 @@ def main(local_file: str | None = None, force_reload: bool = False, metrics_json
 
     s3 = _get_boto3("client", "s3")
 
-    with SparkScriptSession() as spark:
-        raw_model = DEFCBronze()
-        if not raw_model.exists():
-            raise ValueError(f"{raw_model.TABLE_REF} doesn't exist. Use create_migrate_delta_table beforehand.")
-
-        group_model = DEFCGroup()
-        if not group_model.exists():
-            raise ValueError(f"{group_model.TABLE_REF} doesn't exist. Use create_migrate_delta_table beforehand.")
-
-        int_model = DEFCSilver(spark=spark)
-        if not int_model.exists():
-            raise ValueError(f"{int_model.TABLE_REF} doesn't exist. Use create_migrate_delta_table beforehand.")
-
-        start_time = datetime.now()
-        metrics_json["start_time"] = str(start_time)
-
-        logger.info("Parsing DEFC data")
-        try:
-            if not local_file:
-                raw_data = raw_model.to_pandas_df(dtype=str, na_filter=False)
-            else:
-                raw_data = pd.read_csv(local_file, dtype=str, na_filter=False)
-        except pd.errors.EmptyDataError:
-            metrics_json["blank_file"] = True
-            metrics_json["exit_code"] = 4  # exit code chosen arbitrarily, to indicate distinct failure states
-            return metrics_json
-        headers = set([header.upper() for header in list(raw_data)])
-
-        if not VALID_HEADERS.issubset(headers):
-            logger.error("Missing required headers. Required headers include: %s" % str(VALID_HEADERS))
-            metrics_json["exit_code"] = 4
-            return metrics_json
-        metrics_json["records_received"] = len(raw_data)
-        # Creating a dataframe of the export csv first and then copying columns to match the database
-        raw_data = raw_data.rename(columns={"DEFC_CODE": "DEFC", "DEFC_TITLE": "Public Law"})
-
-        group_model_df = group_model.to_pandas_df()
-        group_mapping = group_model_df.groupby("group")["code"].agg(list).to_dict()
-
-        raw_data = apply_defc_derivations(raw_data, group_mapping)
-
-        raw_data = add_defc_outliers(raw_data, group_mapping)
-
-        # Clear any lingering np.nan's
-        raw_data = raw_data.replace({np.nan: None})
-
-        logger.info("Checking for differences in DEFC data")
-        defc_mapping = {
-            "defc": "code",
-            "public_law": "public_laws",
-            "public_law_short_title": "public_law_short_titles",
-            "group_name": "group",
-            "urls": "urls",
-            "is_valid": "is_valid",
-            "earliest_public_law_enactment_date": "earliest_pl_action_date",
-        }
-        data = clean_data(raw_data, defc_mapping, {})
-        diff_found = check_dataframe_diff(data, int_model.to_pandas_df(), ["defc_id"], ["code"], date_format="%Y-%m-%d")
-        if force_reload or diff_found:
-
-            # The only diff should be whenever a new code is added. Noting it here
-            if diff_found:
-                incoming_defcs = list(data["code"])
-                curr_defcs = list(int_model.to_pandas_df()["code"])
-                diff_defcs = list(set(incoming_defcs) - set(curr_defcs))
-                metrics_json["new_defc"] = diff_defcs
-                logger.info(f"Difference found: {diff_defcs}")
-
-            logger.info("Overwriting new DEFC data to Broker")
-            int_model.save(data)
-
-            update_external_data_load_date(int_model, start_time, datetime.now())
-            logger.info("{} records inserted to DEFC".format(len(data)))
-
-            # convert the arrays to pipe-delimited strings
-            defc_delim = "|"
-            array_cols = ["Public Law", "Public Law Short Title", "URLs"]
-            for array_col in array_cols:
-                raw_data[array_col] = raw_data[array_col].apply(lambda value: defc_delim.join(value))
-
-            header_order = [
-                "DEFC",
-                "Public Law",
-                "Public Law Short Title",
-                "Group Name",
-                "URLs",
-                "Is Valid",
-                "Earliest Public Law Enactment Date",
-            ]
-            raw_data = raw_data[header_order]
-            export_name = "def_codes.csv"
-            logger.info("Exporting loaded DEFC file to {}".format(export_name))
-            raw_data.to_csv(export_name, index=0)
-
-            s3.upload_file(export_name, CONFIG.PUBLIC_FILES_BUCKET, export_name)
-
-            os.remove(export_name)
+    raw_model = DEFCBronze()
+    if not raw_model.exists():
+        raise ValueError(f"{raw_model.TABLE_REF} doesn't exist. Use create_migrate_delta_table beforehand.")
+
+    group_model = DEFCGroup()
+    if not group_model.exists():
+        raise ValueError(f"{group_model.TABLE_REF} doesn't exist. Use create_migrate_delta_table beforehand.")
+
+    gold_model = DEFCGold()
+    if not gold_model.exists():
+        raise ValueError(f"{gold_model.TABLE_REF} doesn't exist. Use create_migrate_delta_table beforehand.")
+
+    start_time = datetime.now()
+    metrics_json["start_time"] = str(start_time)
+
+    logger.info("Parsing DEFC data")
+    try:
+        if not local_file:
+            raw_data = raw_model.to_pandas_df(dtype=str, na_filter=False)
         else:
-            logger.info("No differences found, skipping defc table reload.")
+            raw_data = pd.read_csv(local_file, dtype=str, na_filter=False)
+    except pd.errors.EmptyDataError:
+        metrics_json["blank_file"] = True
+        metrics_json["exit_code"] = 4  # exit code chosen arbitrarily, to indicate distinct failure states
+        return metrics_json
+    headers = set([header.upper() for header in list(raw_data)])
+
+    if not VALID_HEADERS.issubset(headers):
+        logger.error("Missing required headers. Required headers include: %s" % str(VALID_HEADERS))
+        metrics_json["exit_code"] = 4
+        return metrics_json
+    metrics_json["records_received"] = len(raw_data)
+    # Creating a dataframe of the export csv first and then copying columns to match the database
+    raw_data = raw_data.rename(columns={"DEFC_CODE": "DEFC", "DEFC_TITLE": "Public Law"})
+
+    group_model_df = group_model.to_pandas_df()
+    group_mapping = group_model_df.groupby("group")["code"].agg(list).to_dict()
+
+    raw_data = apply_defc_derivations(raw_data, group_mapping)
+
+    raw_data = add_defc_outliers(raw_data, group_mapping)
+
+    # Clear any lingering np.nan's
+    raw_data = raw_data.replace({np.nan: None})
+
+    logger.info("Checking for differences in DEFC data")
+    defc_mapping = {
+        "defc": "code",
+        "public_law": "public_laws",
+        "public_law_short_title": "public_law_short_titles",
+        "group_name": "group",
+        "urls": "urls",
+        "is_valid": "is_valid",
+        "earliest_public_law_enactment_date": "earliest_pl_action_date",
+    }
+    data = clean_data(raw_data, defc_mapping, {})
+    diff_found = check_dataframe_diff(data, gold_model.to_pandas_df(), ["defc_id"], ["code"], date_format="%Y-%m-%d")
+    if force_reload or diff_found:
+
+        # The only diff should be whenever a new code is added. Noting it here
+        if diff_found:
+            incoming_defcs = list(data["code"])
+            curr_defcs = list(gold_model.to_pandas_df()["code"])
+            diff_defcs = list(set(incoming_defcs) - set(curr_defcs))
+            metrics_json["new_defc"] = diff_defcs
+            logger.info(f"Difference found: {diff_defcs}")
+
+        logger.info("Overwriting new DEFC data to Broker")
+        gold_model.save(data)
+
+        update_external_data_load_date(gold_model, start_time, datetime.now())
+        logger.info("{} records inserted to DEFC".format(len(data)))
+
+        # convert the arrays to pipe-delimited strings
+        defc_delim = "|"
+        array_cols = ["Public Law", "Public Law Short Title", "URLs"]
+        for array_col in array_cols:
+            raw_data[array_col] = raw_data[array_col].apply(lambda value: defc_delim.join(value))
+
+        header_order = [
+            "DEFC",
+            "Public Law",
+            "Public Law Short Title",
+            "Group Name",
+            "URLs",
+            "Is Valid",
+            "Earliest Public Law Enactment Date",
+        ]
+        raw_data = raw_data[header_order]
+        export_name = "def_codes.csv"
+        logger.info("Exporting loaded DEFC file to {}".format(export_name))
+        raw_data.to_csv(export_name, index=0)
+
+        s3.upload_file(export_name, CONFIG.PUBLIC_FILES_BUCKET, export_name)
+
+        os.remove(export_name)
+    else:
+        logger.info("No differences found, skipping defc table reload.")
 
-        total_defc_count = int_model.count()
+    total_defc_count = gold_model.count()
 
     metrics_json["total_defc_count"] = total_defc_count
 

…mmon/tests/integration/test_load_defc.py …ts/integration/loaders/test_defc_gold.pybrus_backend_common/tests/integration/test_load_defc.py renamed to brus_backend_common/tests/integration/loaders/test_defc_gold.py brus_backend_common/tests/integration/test_load_defc.py renamed to brus_backend_common/tests/integration/loaders/test_defc_gold.py

@@ -1,12 +1,10 @@
 import os
 import pytest
-from typing import Generator, List
-
-from pyspark.sql import SparkSession
+from typing import List
 
 from brus_backend_common.config import _SRC_ROOT_DIR
 from brus_backend_common.helpers.aws import _get_boto3
-from brus_backend_common.scripts.loaders import load_defc
+from brus_backend_common.scripts.loaders import defc_gold
 from brus_backend_common.models import LAKEHOUSE_MODELS
 
 
@@ -26,7 +24,7 @@ def raw_defc_file():
 @pytest.fixture(scope="function")
 def raw_defc_mapping_file():
     # Mimic placing the raw DEFC mapping file in the expected location (directly or copied from another bucket)
-    defc_mapping_model = LAKEHOUSE_MODELS["silver.defc_mapping"]()
+    defc_mapping_model = LAKEHOUSE_MODELS["gold.defc_mapping"]()
     s3_client = _get_boto3("client", "s3")
     csv_file_path = os.path.join(_SRC_ROOT_DIR, "tests", "integration", "data", "defc_groups.csv")
     s3_client.upload_file(csv_file_path, defc_mapping_model.BUCKET_NAME, defc_mapping_model.RELATIVE_CSV_PATH)
@@ -39,42 +37,40 @@ def raw_defc_mapping_file():
 def test_load_defc(
     raw_defc_file: str,
     raw_defc_mapping_file: str,
-    spark: SparkSession,
     setup_teardown_buckets: List[str],
-    hive_unittest_metastore_db: Generator[str | None, None, None],
     external_data_load_dates: str,
 ):
-    # RAW DEFC
-    raw_defc_model = LAKEHOUSE_MODELS["bronze.defc"]()
+    # Bronze DEFC
+    defc_bronze_model = LAKEHOUSE_MODELS["bronze.defc"]()
 
-    assert raw_defc_model.exists()
-    df = raw_defc_model.to_pandas_df()
+    assert defc_bronze_model.exists()
+    df = defc_bronze_model.to_pandas_df()
     assert df is not None and not df.empty
     assert df.loc[df["DEFC_CODE"] == "S", "DEFC_TITLE"].values[0] == "Disaster PL 116-260"
 
     # DEFC Mapping
-    defc_mapping_model = LAKEHOUSE_MODELS["silver.defc_mapping"]()
+    defc_mapping_model = LAKEHOUSE_MODELS["gold.defc_mapping"]()
 
     assert defc_mapping_model.exists()
     df = defc_mapping_model.to_pandas_df()
     assert df is not None and not df.empty
     assert df.loc[df["code"] == "L", "group"].values[0] == "covid_19"
 
-    # INT DEFC
-    int_defc_model = LAKEHOUSE_MODELS["silver.defc"](spark=spark)
-    int_defc_model.initialize()
+    # Gold DEFC
+    defc_gold_model = LAKEHOUSE_MODELS["gold.defc"]()
+    defc_gold_model.initialize(recreate=True)
 
-    load_defc.main()
+    defc_gold.main()
 
-    assert int_defc_model.exists()
-    df = int_defc_model.to_pandas_df()
+    assert defc_gold_model.exists()
+    df = defc_gold_model.to_pandas_df()
     assert df is not None and not df.empty
     assert df.loc[df["code"] == "L", "public_laws"].values[0] == "Emergency P.L. 116-123"
 
     # Confirming the external load date was updated
-    edld_model = LAKEHOUSE_MODELS["bronze.external_data_load_date"]()
+    edld_model = LAKEHOUSE_MODELS["gold.external_data_load_date"]()
 
     assert edld_model.exists()
     df = edld_model.to_pandas_df()
     assert df is not None and not df.empty
-    assert not df.loc[df["name"] == "silver.defc"].empty
+    assert not df.loc[df["name"] == "gold.defc"].empty