Output Parquet files as well as SQLite in PUDL ETL

src/pudl/analysis/allocate_gen_fuel.py

@@ -316,7 +316,7 @@ def gen_fuel_by_gen_esc_owner(
     for freq in ["AS", "MS"]
     for allocated_net_gen_asset in allocate_gen_fuel_asset_factory(
         freq=freq,
-        io_manager_key="pudl_sqlite_io_manager",
+        io_manager_key="pudl_io_manager",
     )
 ]
 

src/pudl/analysis/mcoe.py

@@ -56,7 +56,7 @@ def mcoe_asset_factory(
             "bga": AssetIn(key="core_eia860__assn_boiler_generator"),
         },
         compute_kind="Python",
-        io_manager_key="pudl_sqlite_io_manager",
+        io_manager_key="pudl_io_manager",
     )
     def hr_by_unit_asset(gen: pd.DataFrame, bga: pd.DataFrame) -> pd.DataFrame:
         return heat_rate_by_unit(gen_fuel_by_energy_source=gen, bga=bga)
@@ -69,7 +69,7 @@ def hr_by_unit_asset(gen: pd.DataFrame, bga: pd.DataFrame) -> pd.DataFrame:
             "gens": AssetIn(key="_out_eia__yearly_generators"),
         },
         compute_kind="Python",
-        io_manager_key="pudl_sqlite_io_manager",
+        io_manager_key="pudl_io_manager",
     )
     def hr_by_gen_asset(
         bga: pd.DataFrame, hr_by_unit: pd.DataFrame, gens: pd.DataFrame
@@ -86,7 +86,7 @@ def hr_by_gen_asset(
             "frc": AssetIn(key=f"out_eia923__{agg_freqs[freq]}_fuel_receipts_costs"),
         },
         compute_kind="Python",
-        io_manager_key="pudl_sqlite_io_manager",
+        io_manager_key="pudl_io_manager",
     )
     def fc_asset(
         hr_by_gen: pd.DataFrame, gens: pd.DataFrame, frc: pd.DataFrame
@@ -102,7 +102,7 @@ def fc_asset(
             "gens": AssetIn(key="_out_eia__yearly_generators"),
         },
         compute_kind="Python",
-        io_manager_key="pudl_sqlite_io_manager",
+        io_manager_key="pudl_io_manager",
     )
     def cf_asset(gens: pd.DataFrame, gen: pd.DataFrame) -> pd.DataFrame:
         return capacity_factor(gens=gens, gen=gen, freq=freq)
@@ -157,7 +157,7 @@ def cf_asset(gens: pd.DataFrame, gen: pd.DataFrame) -> pd.DataFrame:
                 ),
             ),
         },
-        io_manager_key="pudl_sqlite_io_manager",
+        io_manager_key="pudl_io_manager",
     )
     def mcoe_asset(
         context, fuel_cost: pd.DataFrame, capacity_factor: pd.DataFrame
@@ -179,7 +179,7 @@ def mcoe_asset(
             ),
             "gens": AssetIn(key="_out_eia__yearly_generators"),
         },
-        io_manager_key="pudl_sqlite_io_manager",
+        io_manager_key="pudl_io_manager",
         compute_kind="Python",
         config_schema={
             "all_gens": Field(

src/pudl/analysis/plant_parts_eia.py

@@ -380,7 +380,7 @@ def out_eia__yearly_generators_by_ownership(
 
 
 @asset(
-    io_manager_key="pudl_sqlite_io_manager",
+    io_manager_key="pudl_io_manager",
     compute_kind="Python",
 )
 def out_eia__yearly_plant_parts(

src/pudl/analysis/record_linkage/eia_ferc1_record_linkage.py

@@ -276,7 +276,7 @@ def run_matching_model(features_train, features_all, y_df):
 @op(
     out={
         "out_pudl__yearly_assn_eia_ferc1_plant_parts": Out(
-            io_manager_key="pudl_sqlite_io_manager"
+            io_manager_key="pudl_io_manager"
         )
     }
 )

src/pudl/analysis/service_territory.py

@@ -456,7 +456,7 @@ def dagster_compile_geoms(
     ass
     for entity in list(ENTITY_TYPE)
     for ass in compiled_geoms_asset_factory(
-        entity_type=entity, io_manager_key="pudl_sqlite_io_manager"
+        entity_type=entity, io_manager_key="pudl_io_manager"
     )
 ]
 

src/pudl/analysis/state_demand.py

@@ -565,7 +565,7 @@ def total_state_sales_eia861(
 
 
 @asset(
-    io_manager_key="pudl_sqlite_io_manager",
+    io_manager_key="pudl_io_manager",
     compute_kind="Python",
     config_schema={
         "mean_overlaps": Field(

src/pudl/etl/__init__.py

@@ -15,7 +15,7 @@
     epacems_io_manager,
     ferc1_dbf_sqlite_io_manager,
     ferc1_xbrl_sqlite_io_manager,
-    pudl_sqlite_io_manager,
+    pudl_mixed_format_io_manager,
 )
 from pudl.resources import dataset_settings, datastore, ferc_to_sqlite_settings
 from pudl.settings import EtlSettings
@@ -95,7 +95,7 @@
 
 default_resources = {
     "datastore": datastore,
-    "pudl_sqlite_io_manager": pudl_sqlite_io_manager,
+    "pudl_io_manager": pudl_mixed_format_io_manager,
     "ferc1_dbf_sqlite_io_manager": ferc1_dbf_sqlite_io_manager,
     "ferc1_xbrl_sqlite_io_manager": ferc1_xbrl_sqlite_io_manager,
     "dataset_settings": dataset_settings,

src/pudl/etl/check_foreign_keys.py

@@ -3,11 +3,12 @@
 import sys
 
 import click
-from dagster import build_init_resource_context
+import pandas as pd
+import sqlalchemy as sa
 from dotenv import load_dotenv
 
 import pudl
-from pudl.io_managers import pudl_sqlite_io_manager
+from pudl.workspace.setup import PudlPaths
 
 logger = pudl.logging_helpers.get_logger(__name__)
 
@@ -31,7 +32,17 @@
         ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], case_sensitive=False
     ),
 )
-def pudl_check_fks(logfile: pathlib.Path, loglevel: str):
+@click.option(
+    "--db_path",
+    help="Path to PUDL SQLite database where foreign keys should be checked.",
+    type=click.Path(
+        exists=False,
+        resolve_path=True,
+        path_type=pathlib.Path,
+    ),
+    default=None,
+)
+def pudl_check_fks(logfile: pathlib.Path, loglevel: str, db_path: pathlib.Path):
     """Check that foreign key constraints in the PUDL database are respected.
 
     Dagster manages the dependencies between various assets in our ETL pipeline,
@@ -50,15 +61,154 @@ def pudl_check_fks(logfile: pathlib.Path, loglevel: str):
     # Display logged output from the PUDL package:
     pudl.logging_helpers.configure_root_logger(logfile=logfile, loglevel=loglevel)
 
-    context = build_init_resource_context()
-    io_manager = pudl_sqlite_io_manager(context)
+    # Using PudlPaths to get default value for CLI causes validation issues
+    if not db_path:
+        db_path = PudlPaths().output_dir / "pudl.sqlite"
 
-    database_path = io_manager.base_dir / f"{io_manager.db_name}.sqlite"
-    logger.info(f"Checking foreign key constraints in {database_path}")
+    engine = sa.create_engine(f"sqlite:///{db_path}")
 
-    io_manager.check_foreign_keys()
+    check_foreign_keys(engine)
     return 0
 
 
+class ForeignKeyError(sa.exc.SQLAlchemyError):
+    """Raised when data in a database violates a foreign key constraint."""
+
+    def __init__(
+        self, child_table: str, parent_table: str, foreign_key: str, rowids: list[int]
+    ):
+        """Initialize a new ForeignKeyError object.
+
+        Args:
+            child_table: The table that a foreign key constraint is applied to.
+            parent_table: The table that a foreign key constraint refers to.
+            foreign_key: Comma seperated string of key(s) that make up foreign key.
+            rowids: Row(s) of child_table where constraint failed.
+        """
+        self.child_table = child_table
+        self.parent_table = parent_table
+        self.foreign_key = foreign_key
+        self.rowids = rowids
+
+    def __str__(self):
+        """Create string representation of ForeignKeyError object."""
+        return (
+            f"Foreign key error for table: {self.child_table} -- {self.parent_table} "
+            f"{self.foreign_key} -- on rows {self.rowids}\n"
+        )
+
+    def __eq__(self, other):
+        """Compare a ForeignKeyError with another object."""
+        if isinstance(other, ForeignKeyError):
+            return (
+                (self.child_table == other.child_table)
+                and (self.parent_table == other.parent_table)
+                and (self.foreign_key == other.foreign_key)
+                and (self.rowids == other.rowids)
+            )
+        return False
+
+
+class ForeignKeyErrors(sa.exc.SQLAlchemyError):  # noqa: N818
+    """Raised when data in a database violate multiple foreign key constraints."""
+
+    def __init__(self, fk_errors: list[ForeignKeyError]):
+        """Initialize a new ForeignKeyErrors object."""
+        self.fk_errors = fk_errors
+
+    def __str__(self):
+        """Create string representation of ForeignKeyErrors object."""
+        fk_errors = [str(x) for x in self.fk_errors]
+        return "\n".join(fk_errors)
+
+    def __iter__(self):
+        """Iterate over the fk errors."""
+        return self.fk_errors
+
+    def __getitem__(self, idx):
+        """Index the fk errors."""
+        return self.fk_errors[idx]
+
+
+def _get_fk_list(engine: sa.Engine, table: str) -> pd.DataFrame:
+    """Retrieve a dataframe of foreign keys for a table.
+
+    Description from the SQLite Docs: 'This pragma returns one row for each foreign
+    key constraint created by a REFERENCES clause in the CREATE TABLE statement of
+    table "table-name".'
+
+    The PRAGMA returns one row for each field in a foreign key constraint. This
+    method collapses foreign keys with multiple fields into one record for
+    readability.
+    """
+    with engine.begin() as con:
+        table_fks = pd.read_sql_query(f"PRAGMA foreign_key_list({table});", con)
+
+    # Foreign keys with multiple fields are reported in separate records.
+    # Combine the multiple fields into one string for readability.
+    # Drop duplicates so we have one FK for each table and foreign key id
+    table_fks["fk"] = table_fks.groupby("table")["to"].transform(
+        lambda field: "(" + ", ".join(field) + ")"
+    )
+    table_fks = table_fks[["id", "table", "fk"]].drop_duplicates()
+
+    # Rename the fields so we can easily merge with the foreign key errors.
+    table_fks = table_fks.rename(columns={"id": "fkid", "table": "parent"})
+    table_fks["table"] = table
+    return table_fks
+
+
+def check_foreign_keys(engine: sa.Engine):
+    """Check foreign key relationships in the database.
+
+    The order assets are loaded into the database will not satisfy foreign key
+    constraints so we can't enable foreign key constraints. However, we can
+    check for foreign key failures once all of the data has been loaded into
+    the database using the `foreign_key_check` and `foreign_key_list` PRAGMAs.
+
+    You can learn more about the PRAGMAs in the `SQLite docs
+    <https://www.sqlite.org/pragma.html#pragma_foreign_key_check>`__.
+
+    Raises:
+        ForeignKeyErrors: if data in the database violate foreign key constraints.
+    """
+    logger.info(f"Running foreign key check on {engine.name} database.")
+    with engine.begin() as con:
+        fk_errors = pd.read_sql_query("PRAGMA foreign_key_check;", con)
+
+    if not fk_errors.empty:
+        # Merge in the actual FK descriptions
+        tables_with_fk_errors = fk_errors.table.unique().tolist()
+        table_foreign_keys = pd.concat(
+            [_get_fk_list(engine, table) for table in tables_with_fk_errors]
+        )
+
+        fk_errors_with_keys = fk_errors.merge(
+            table_foreign_keys,
+            how="left",
+            on=["parent", "fkid", "table"],
+            validate="m:1",
+        )
+
+        errors = []
+        # For each foreign key error, raise a ForeignKeyError
+        for (
+            table_name,
+            parent_name,
+            parent_fk,
+        ), parent_fk_df in fk_errors_with_keys.groupby(["table", "parent", "fk"]):
+            errors.append(
+                ForeignKeyError(
+                    child_table=table_name,
+                    parent_table=parent_name,
+                    foreign_key=parent_fk,
+                    rowids=parent_fk_df["rowid"].values,
+                )
+            )
+        raise ForeignKeyErrors(errors)
+
+    logger.info("Success! No foreign key constraint errors found.")
+
+
 if __name__ == "__main__":
     sys.exit(pudl_check_fks())

src/pudl/etl/cli.py

@@ -88,6 +88,7 @@ def get_pudl_etl_job():
 @click.option(
     "--gcs-cache-path",
     type=str,
+    default="",
     help=(
         "Load cached inputs from Google Cloud Storage if possible. This is usually "
         "much faster and more reliable than downloading from Zenodo directly. The "
@@ -128,7 +129,7 @@ def pudl_etl(
 
     dataset_settings_config = etl_settings.datasets.model_dump()
     process_epacems = True
-    if etl_settings.datasets.epacems is None:
+    if etl_settings.datasets.epacems is None or etl_settings.datasets.epacems.disabled:
         process_epacems = False
         # Dagster config expects values for the epacems settings even though
         # the CEMS assets will not be executed. Fill in the config dictionary

src/pudl/etl/eia_bulk_elec_assets.py

@@ -7,7 +7,7 @@
 
 
 @asset(
-    io_manager_key="pudl_sqlite_io_manager",
+    io_manager_key="pudl_io_manager",
     required_resource_keys={"datastore"},
 )
 def core_eia__yearly_fuel_receipts_costs_aggs(context):

src/pudl/etl/glue_assets.py

@@ -17,7 +17,7 @@
 
 @multi_asset(
     outs={
-        table_name: AssetOut(io_manager_key="pudl_sqlite_io_manager")
+        table_name: AssetOut(io_manager_key="pudl_io_manager")
         for table_name in Package.get_etl_group_tables("glue")
         #  do not load core_epa__assn_eia_epacamd glue assets bc they are stand-alone assets below.
         if "core_epa__assn_eia_epacamd" not in table_name
@@ -76,9 +76,7 @@ def raw_pudl__assn_eia_epacamd(context) -> pd.DataFrame:
     return pd.concat(year_matches, ignore_index=True)
 
 
-@asset(
-    required_resource_keys={"dataset_settings"}, io_manager_key="pudl_sqlite_io_manager"
-)
+@asset(required_resource_keys={"dataset_settings"}, io_manager_key="pudl_io_manager")
 def core_epa__assn_eia_epacamd(
     context,
     raw_pudl__assn_eia_epacamd: pd.DataFrame,
@@ -262,7 +260,7 @@ def correct_epa_eia_plant_id_mapping(df: pd.DataFrame) -> pd.DataFrame:
 ##############################
 
 
-@asset(io_manager_key="pudl_sqlite_io_manager")
+@asset(io_manager_key="pudl_io_manager")
 def core_epa__assn_eia_epacamd_subplant_ids(
     _core_epa__assn_eia_epacamd_unique: pd.DataFrame,
     core_eia860__scd_generators: pd.DataFrame,

src/pudl/etl/static_assets.py

@@ -40,7 +40,7 @@ def _read_static_encoding_tables(
 
 @multi_asset(
     outs={
-        table_name: AssetOut(io_manager_key="pudl_sqlite_io_manager")
+        table_name: AssetOut(io_manager_key="pudl_io_manager")
         for table_name in Package.get_etl_group_tables("static_pudl")
     },
     required_resource_keys={"dataset_settings", "datastore"},
@@ -62,7 +62,7 @@ def static_pudl_tables(context):
 
 @multi_asset(
     outs={
-        table_name: AssetOut(io_manager_key="pudl_sqlite_io_manager")
+        table_name: AssetOut(io_manager_key="pudl_io_manager")
         for table_name in Package.get_etl_group_tables("static_eia")
     },
 )
@@ -76,7 +76,7 @@ def static_eia_tables():
 
 @multi_asset(
     outs={
-        table_name: AssetOut(io_manager_key="pudl_sqlite_io_manager")
+        table_name: AssetOut(io_manager_key="pudl_io_manager")
         for table_name in Package.get_etl_group_tables("static_ferc1")
     },
 )