refactor(elt-common): Improve BaseExtract config typing

elt-common/src/elt_common/extract.py

@@ -1,12 +1,12 @@
 import dataclasses as dc
+import datetime as dt
 import importlib.util
 import json
 import sys
 from abc import ABC, abstractmethod
-import datetime as dt
 from pathlib import Path
 from types import ModuleType
-from typing import TYPE_CHECKING, Callable, Iterator, Optional, get_args
+from typing import TYPE_CHECKING, Callable, ClassVar, Iterator, Optional, get_args
 
 from pydantic_settings import BaseSettings
 
@@ -104,22 +104,22 @@ class ResourceProperties:
     watermark_column: Optional[str] = None
 
 
-class BaseExtract(ABC):
+class BaseExtract[C: BaseSettings](ABC):
     """Base class for ingest Extract classes"""
 
-    config_cls: type[BaseSettings] = BaseSettings
+    config_cls: ClassVar[type[BaseSettings]] = BaseSettings
     """Class used to provide configuration options.
 
     Override this in subclasses to provide custom configuration.
 
     Intended to be used with pydantic-settings.
     """
 
-    def __init__(self, config: BaseSettings):
+    def __init__(self, config: C):
         self._config = config
 
     @property
-    def config(self):
+    def config(self) -> C:
         return self._config
 
     @abstractmethod

elt-common/src/elt_common/sources/sqldatabase/__init__.py

@@ -66,7 +66,7 @@ class TableInfo(NamedTuple):
     destination_table_name: Optional[str] = None
 
 
-class SqlDatabaseExtract(BaseExtract):
+class SqlDatabaseExtract(BaseExtract[SqlDatabaseSourceConfig]):
     """Base class for defining SQL ingest Extract classes.
 
     Example usage, for an ingest script that reads from 3 tables::
@@ -90,8 +90,6 @@ def table_info(self):
 
     def __init__(self, config: SqlDatabaseSourceConfig):
         super().__init__(config)
-        self._chunk_size = config.chunk_size
-        self._row_limit = config.row_limit
 
         LOGGER.debug(
             f"Creating engine for {config.drivername} database at "
@@ -168,7 +166,7 @@ def _extract_table(
         watermark: Watermark | None = None,
         query_filter: Callable[[Select], Select] | None = None,
     ) -> Iterator[pa.Table]:
-        LOGGER.debug(f"Extracting table {name} in chunks of {self._chunk_size} rows.")
+        LOGGER.debug(f"Extracting table {name} in chunks of {self.config.chunk_size} rows.")
         table = sa.Table(
             name,
             self._metadata,
@@ -183,13 +181,13 @@ def _extract_table(
         if query_filter:
             query = query_filter(query)
 
-        query = query.limit(self._row_limit)
+        query = query.limit(self.config.row_limit)
 
         # If all the values in a column are null pyarrow won't know what type
         # the column should be, so we need to explicitly create a schema from
         # the table
         pa_schema = to_pyarrow_schema(table)
-        result = conn.execution_options(yield_per=self._chunk_size).execute(query)
+        result = conn.execution_options(yield_per=self.config.chunk_size).execute(query)
         for partition in result.mappings().partitions():
             table = pa.Table.from_pylist(partition, schema=pa_schema)
             yield table

elt-common/tests/unit_tests/test_extract.py

@@ -116,4 +116,4 @@ def test_create_extract_obj_sql_extract(monkeypatch):
 
     assert isinstance(extract_obj, BaseExtract)
     assert isinstance(extract_obj, SqlDatabaseExtract)
-    assert extract_obj._chunk_size == 100
+    assert extract_obj.config.chunk_size == 100

elt-pipelines/facility_ops/ingest/estates/electricity_sharepoint/electricity_sharepoint.py

@@ -45,16 +45,12 @@ def glob_patterns(self):
         return self.backfill_globs if self.backfill_globs else _default_backfill_globs
 
 
-class Extract(BaseExtract):
+class Extract(BaseExtract[Configuration]):
     config_cls = Configuration
 
     def __init__(self, cfg: Configuration):
         super().__init__(cfg)
         self._client = SPListClient(SITE_URL, cfg)
-        self._backfilling = cfg.backfill
-        self._glob_patterns = cfg.glob_patterns
-
-        LOGGER.debug(f"Searching for files matching: {self._glob_patterns}")
 
     def extract_resource_properties(self):
         yield (
@@ -69,8 +65,10 @@ def extract_resource_properties(self):
         )
 
     def _extract_electricity_usage(self, w: Watermark | None):
+        LOGGER.debug(f"Searching for files matching: {self.config.glob_patterns}")
+
         watermark_value: dt.datetime | None = None
-        if w and self._backfilling:
+        if w and self.config.backfill:
             LOGGER.debug("Ignoring watermark because this is a backfill")
         elif w:
             if not isinstance(w.value, dt.datetime):
@@ -82,7 +80,7 @@ def _extract_electricity_usage(self, w: Watermark | None):
                 LOGGER.debug(f"Only fetching files modified after {watermark_value}")
 
         files = []
-        for pattern in self._glob_patterns:
+        for pattern in self.config.glob_patterns:
             files.extend(
                 self._client.glob(
                     _root_path, pattern=pattern, modified_after=watermark_value