From 2331201e93ba46e4a3b3fa31b5f4553cf29b72ba Mon Sep 17 00:00:00 2001 From: Will Taylor Date: Thu, 13 Aug 2026 16:41:47 +0100 Subject: [PATCH] refactor(elt-common): Improve BaseExtract config typing This makes type checkers happy when extract classes which use custom configuration access the fields on their configuration --- elt-common/src/elt_common/extract.py | 12 ++++++------ .../src/elt_common/sources/sqldatabase/__init__.py | 10 ++++------ elt-common/tests/unit_tests/test_extract.py | 2 +- .../electricity_sharepoint/electricity_sharepoint.py | 12 +++++------- 4 files changed, 16 insertions(+), 20 deletions(-) diff --git a/elt-common/src/elt_common/extract.py b/elt-common/src/elt_common/extract.py index 03f15164..f1961146 100644 --- a/elt-common/src/elt_common/extract.py +++ b/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,10 +104,10 @@ 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. @@ -115,11 +115,11 @@ class BaseExtract(ABC): 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 diff --git a/elt-common/src/elt_common/sources/sqldatabase/__init__.py b/elt-common/src/elt_common/sources/sqldatabase/__init__.py index ef9d6da8..e075c2f7 100644 --- a/elt-common/src/elt_common/sources/sqldatabase/__init__.py +++ b/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 diff --git a/elt-common/tests/unit_tests/test_extract.py b/elt-common/tests/unit_tests/test_extract.py index df826f29..bfc2638d 100644 --- a/elt-common/tests/unit_tests/test_extract.py +++ b/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 diff --git a/elt-pipelines/facility_ops/ingest/estates/electricity_sharepoint/electricity_sharepoint.py b/elt-pipelines/facility_ops/ingest/estates/electricity_sharepoint/electricity_sharepoint.py index 18cdafaf..2f7237d2 100644 --- a/elt-pipelines/facility_ops/ingest/estates/electricity_sharepoint/electricity_sharepoint.py +++ b/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