Start EIA-176 pipelines: company data

src/pudl/etl/__init__.py

@@ -32,6 +32,7 @@
 default_assets = (
     *load_assets_from_modules([eia_bulk_elec_assets], group_name="eia_bulk_elec"),
     *load_assets_from_modules([epacems_assets], group_name="epacems"),
+    *load_assets_from_modules([pudl.extract.eia176], group_name="raw_eia176"),
     *load_assets_from_modules([pudl.extract.eia860], group_name="raw_eia860"),
     *load_assets_from_modules([pudl.transform.eia860], group_name="_core_eia860"),
     *load_assets_from_modules([pudl.extract.eia861], group_name="raw_eia861"),

src/pudl/extract/__init__.py

@@ -8,6 +8,7 @@
 :mod:`pudl.transform` subpackage.
 """
 from . import (
+    eia176,
     eia860,
     eia860m,
     eia861,

src/pudl/extract/csv.py

@@ -0,0 +1,106 @@
+"""Extractor for CSV data."""
+from csv import DictReader
+from importlib import resources
+
+import pandas as pd
+from dagster import AssetsDefinition, OpDefinition, graph_asset, op
+
+import pudl.logging_helpers
+from pudl.workspace.datastore import Datastore
+
+logger = pudl.logging_helpers.get_logger(__name__)
+
+
+class CsvExtractor:
+    """Generalized class for extracting dataframes from CSV files.
+
+    When subclassing from this generic extractor, one should implement dataset specific
+    logic in the following manner:
+
+    2. Set DATASET class attribute. This is used to load metadata from package_data/{dataset} subdirectory.
+
+    The extraction logic is invoked by calling extract() method of this class.
+    """
+
+    DATASET = None
+
+    def __init__(self, datastore: Datastore):
+        """Create a new instance of CsvExtractor.
+
+        This can be used for retrieving data from CSV files.
+
+        Args:
+            datastore: provides access to raw files on disk.
+        """
+        self._zipfile = datastore.get_zipfile_resource(self.DATASET)
+        self._table_file_map = {
+            row["table"]: row["filename"]
+            for row in self._open_csv_resource("table_file_map.csv")
+        }
+
+    def _open_csv_resource(self, base_filename: str) -> DictReader:
+        """Open the given resource file as :class:`csv.DictReader`."""
+        csv_path = resources.files(f"pudl.package_data.{self.DATASET}") / base_filename
+        return DictReader(csv_path.open())
+
+    def read_source(self, filename: str) -> pd.DataFrame:
+        """Read the data from the CSV source file and return as a dataframe."""
+        logger.info(f"Extracting {filename} from CSV into pandas DataFrame.")
+        with self._zipfile.open(filename) as f:
+            df = pd.read_csv(f)
+        return df
+
+    def extract(self) -> dict[str, pd.DataFrame]:
+        """Extracts a dictionary of table names and dataframes from CSV source files."""
+        data = {}
+        for table in self._table_file_map:
+            filename = self._table_file_map[table]
+            df = self.read_source(filename)
+            data[table] = df
+        return data
+
+
+def extractor_factory(extractor_cls: type[CsvExtractor], name: str) -> OpDefinition:
+    """Construct a Dagster op that extracts data given an extractor class.
+
+    Args:
+        extractor_cls: Class of type :class:`CsvExtractor` used to extract the data.
+        name: Name of a CSV-based dataset (e.g. "eia176").
+    """
+
+    def extract(context) -> dict[str, pd.DataFrame]:
+        """A function that extracts data from a CSV file.
+
+        This function will be decorated with a Dagster op and returned.
+
+        Args:
+            context: Dagster keyword that provides access to resources and config.
+
+        Returns:
+            A dictionary of DataFrames extracted from CSV, keyed by table name.
+        """
+        ds = context.resources.datastore
+        return extractor_cls(ds).extract()
+
+    return op(
+        required_resource_keys={"datastore", "dataset_settings"},
+        name=f"extract_single_{name}_year",
+    )(extract)
+
+
+def raw_df_factory(extractor_cls: type[CsvExtractor], name: str) -> AssetsDefinition:
+    """Return a dagster graph asset to extract a set of raw DataFrames from CSV files.
+
+    Args:
+        extractor_cls: The dataset-specific CSV extractor used to extract the data.
+            Needs to correspond to the dataset identified by ``name``.
+        name: Name of a CSV-based dataset (e.g. "eia176"). Currently this must be
+            one of the attributes of :class:`pudl.settings.EiaSettings`
+    """
+    extractor = extractor_factory(extractor_cls, name)
+
+    def raw_dfs() -> dict[str, pd.DataFrame]:
+        """Produce a dictionary of extracted EIA dataframes."""
+        return extractor()
+
+    return graph_asset(name=f"{name}_raw_dfs")(raw_dfs)

src/pudl/extract/dbf.py

@@ -159,7 +159,7 @@ def get_table_schema(self, table_name: str) -> DbfTableSchema:
         table_columns = self.get_db_schema()[table_name]
         dbf = self.get_table_dbf(table_name)
         dbf_fields = [field for field in dbf.fields if field.name != "_NullFlags"]
-        if len(table_columns) != len(table_columns):
+        if len(dbf_fields) != len(table_columns):
             return ValueError(
                 f"Number of DBF fields in {table_name} does not match what was "
                 f"found in the DBC index file for {self.partition}."

src/pudl/extract/eia176.py

@@ -0,0 +1,46 @@
+"""Extract EIA Form 176 data from CSVs.
+
+The EIA Form 176 archive also contains CSVs for EIA Form 191 and EIA Form 757.
+"""
+
+from dagster import AssetOut, Output, multi_asset
+
+from pudl.extract.csv import CsvExtractor, raw_df_factory
+
+DATASET = "eia176"
+
+
+class Eia176CsvExtractor(CsvExtractor):
+    """Extractor for EIA Form 176 data."""
+
+    DATASET = DATASET
+
+
+# TODO (davidmudrauskas): Add this information to the metadata
+raw_table_names = (f"raw_{DATASET}__company",)
+
+eia176_raw_dfs = raw_df_factory(Eia176CsvExtractor, name=DATASET)
+
+
+@multi_asset(
+    outs={table_name: AssetOut() for table_name in sorted(raw_table_names)},
+    required_resource_keys={"datastore", "dataset_settings"},
+)
+def extract_eia176(context, eia176_raw_dfs):
+    """Extract EIA-176 data from CSV source and return dataframes.
+
+    Args:
+        context: dagster keyword that provides access to resources and config.
+
+    Returns:
+        A tuple of extracted EIA dataframes.
+    """
+    eia176_raw_dfs = {
+        f"raw_{DATASET}__" + table_name: df for table_name, df in eia176_raw_dfs.items()
+    }
+    eia176_raw_dfs = dict(sorted(eia176_raw_dfs.items()))
+
+    return (
+        Output(output_name=table_name, value=df)
+        for table_name, df in eia176_raw_dfs.items()
+    )

src/pudl/extract/eia191.py

@@ -0,0 +1 @@
+"""Extract EIA Form 191 data from CSVs."""

src/pudl/extract/eia757.py

@@ -0,0 +1 @@
+"""Extract EIA Form 757 data from CSVs."""

src/pudl/extract/ferc1.py

@@ -301,7 +301,7 @@ def add_missing_respondents(self):
 
         # Write missing respondents back into SQLite.
         with self.sqlite_engine.begin() as conn:
-            conn.execute(
+            conn.extract(
                 self.sqlite_meta.tables["f1_respondent_id"].insert().values(records)
             )
 

src/pudl/package_data/eia176/table_file_map.csv

@@ -0,0 +1,2 @@
+table,filename
+company,all_company_176.csv

src/pudl/workspace/datastore.py

@@ -169,6 +169,8 @@ class ZenodoDoiSettings(BaseSettings):
     # Sandbox DOIs are provided for reference
     censusdp1tract: ZenodoDoi = "10.5281/zenodo.4127049"
     # censusdp1tract: ZenodoDoi = "10.5072/zenodo.674992"
+    eia176: ZenodoDoi = "10.5281/zenodo.7682358"
+    # eia176: ZenodoDoi - "10.5072/zenodo.1166385"
     eia860: ZenodoDoi = "10.5281/zenodo.10067566"
     # eia860: ZenodoDoi = "10.5072/zenodo.1222854"
     eia860m: ZenodoDoi = "10.5281/zenodo.8188017"

test/integration/etl_test.py

@@ -68,6 +68,17 @@ def test_ferc1_xbrl2sqlite(ferc1_engine_xbrl, ferc1_xbrl_taxonomy_metadata):
     )
 
 
+class TestCsvExtractor:
+    """Verify that we can lead CSV files as provided via the datastore."""
+
+    def test_extract_eia176(self, pudl_datastore_fixture):
+        """Spot check extraction of eia176 csv files."""
+        extractor = pudl.extract.eia176.Eia176CsvExtractor(pudl_datastore_fixture)
+        table = "company"
+        if table not in extractor.extract():
+            raise AssertionError(f"table {table} not found in datastore")
+
+
 class TestExcelExtractor:
     """Verify that we can lead excel files as provided via the datastore."""
 

test/unit/extract/csv_test.py

@@ -0,0 +1,41 @@
+"""Unit tests for pudl.extract.csv module."""
+from unittest.mock import MagicMock, patch
+
+from pudl.extract.csv import CsvExtractor
+
+TABLE_NAME = "company"
+
+FILENAME = "all_company_176.csv"
+TABLE_FILE_MAP = {TABLE_NAME: FILENAME}
+
+DATASET = "eia176"
+
+
+class FakeCsvExtractor(CsvExtractor):
+    DATASET = DATASET
+
+
+def get_csv_extractor():
+    datastore = MagicMock()
+    return FakeCsvExtractor(datastore)
+
+
+@patch("pudl.extract.csv.pd")
+def test_csv_extractor_read_source(mock_pd):
+    extractor = get_csv_extractor()
+    res = extractor.read_source(FILENAME)
+    mock_zipfile = extractor._zipfile
+    mock_zipfile.open.assert_called_once_with(FILENAME)
+    f = mock_zipfile.open.return_value.__enter__.return_value
+    mock_pd.read_csv.assert_called_once_with(f)
+    df = mock_pd.read_csv()
+    assert df == res
+
+
+def test_csv_extractor_extract():
+    extractor = get_csv_extractor()
+    df = MagicMock()
+    with patch.object(CsvExtractor, "read_source", return_value=df) as mock_read_source:
+        raw_dfs = extractor.extract()
+    mock_read_source.assert_called_once_with(FILENAME)
+    assert {TABLE_NAME: df} == raw_dfs