fix: More reliable electricity data extraction from M365 (#158)

### Summary

More robust extract of electricity data from sharepoint:

- Loading now stops if file content cannot be read to avoid
silently missing data.
- Number of threads capped to 8
- Retry on network errors when reading file content
- Content merged on filename
- Removed dlt lag widow that didn't function as expected

Refs #156

Author: martyngigg

.github/actions/run-pytest-with-uv/action.yml

@@ -11,6 +11,9 @@ inputs:
   pytest-file-or-dir:
     description: Path to a file or directory for pytest execution
     required: true
+  python-version:
+    description: The version of Python to install. Defaults to 3.13 as 3.14 was released but not all dependencies work.
+    default: "3.13"
   uv-cache-dependency-glob:
     description: Cache dependencies based on the supplied glob
   uv-version:
@@ -25,6 +28,7 @@ runs:
       with:
         cache-dependency-glob: ${{ inputs.uv-cache-dependency-glob }}
         version: ${{ inputs.uv-version }}
+        python-version: ${{ inputs.python-version }}
 
     - name: Bring up Docker Compose services
       if: inputs.compose-file-path != ''

elt-common/src/elt_common/dlt_sources/m365/__init__.py

@@ -1,30 +1,21 @@
 """Reads files from a SharePoint documents library"""
 
-from typing import Iterator, List, cast
+from typing import Iterator, List
 
 import dlt
-from dlt.common.storages.fsspec_filesystem import MTIME_DISPATCH, glob_files, FileItemDict
+from dlt.common.storages.fsspec_filesystem import MTIME_DISPATCH, glob_files
 from dlt.extract import decorators
 import dlt.common.logger as logger
 import pendulum
 
-from .helpers import M365CredentialsResource, M365DriveFS
+from .helpers import M365CredentialsResource, M365DriveFS, M365DriveItem
 from .settings import DEFAULT_CHUNK_SIZE
 
 # Add our M365DriveFS protocol(s) to the known modificaton time mappings
 for protocol in M365DriveFS.protocol:
     MTIME_DISPATCH[protocol] = MTIME_DISPATCH["file"]
 
 
-class M365DriveItem(FileItemDict):
-    """Specialises FileItemDict to add 'fetch_bytes' to bypass complicated file reading/caching in
-    'read_bytes' and just download the file content"""
-
-    def read_bytes(self) -> bytes:
-        drive_fs = cast(M365DriveFS, self.fsspec)
-        return drive_fs.fetch_all(self["file_url"])
-
-
 # This is designed to look similar to the dlt.filesystem resource where the resource returns DriveItem
 # objects that include the content as raw bytes. The bytes need to be parsed by an appropriate
 # transformer
@@ -50,8 +41,8 @@ def sharepoint(
     for file_model in glob_files(
         sp_library, bucket_url=M365DriveFS.protocol[0] + "://", file_glob=file_glob
     ):
-        log_msg = f"Found '{file_model['file_name']}' with modification date '{file_model['modification_date']}'"
-        if modified_after and file_model["modification_date"] <= modified_after:
+        log_msg = f"Found '{file_model['file_name']}' with modification date '{file_model['modification_date']}'"  # type: ignore
+        if modified_after and file_model["modification_date"] <= modified_after:  # type: ignore
             log_msg += ": skipped old item."
             continue
         else:

elt-common/src/elt_common/dlt_sources/m365/helpers.py

@@ -6,6 +6,7 @@
 from fsspec import AbstractFileSystem
 from dlt.common.configuration.specs import configspec
 from dlt.common.typing import TSecretStrValue
+from dlt.common.storages.fsspec_filesystem import FileItemDict
 from httpx import Response
 from httpx import HTTPStatusError, NetworkError, TimeoutException
 import tenacity
@@ -236,3 +237,13 @@ def _msgraph_get(self, url: str, **kwargs) -> Response:
     @tenacity.retry(**_RETRY_ARGS)
     def _msgraph_request(self, method: str, url: str, **kwargs) -> Response:
         return self.client.request(method, url, **kwargs)
+
+
+class M365DriveItem(FileItemDict):
+    """Specialises FileItemDict to add 'fetch_bytes' to bypass complicated file reading/caching in
+    'read_bytes' and just download the file content"""
+
+    @tenacity.retry(**_RETRY_ARGS)
+    def read_bytes(self) -> bytes:
+        drive_fs = cast(M365DriveFS, self.fsspec)
+        return drive_fs.fetch_all(self["file_url"])

infra/local/trino/trino-create-catalog.sql

@@ -9,7 +9,7 @@ with (
   "iceberg.rest-catalog.oauth2.credential" = 'localinfra:s3cr3t',
   "iceberg.rest-catalog.oauth2.scope" = 'lakekeeper offline_access',
   "fs.native-s3.enabled" = 'true',
-  "s3.endpoint" = 'http://traefik:9000',
+  "s3.endpoint" = 'http://minio:59000',
   "s3.region" = 'local-01',
   "s3.path-style-access" = 'true',
   "s3.aws-access-key" = 'adpuser',

warehouses/accelerator/extract_load/electricity_sharepoint/.dlt/config.toml

@@ -1,3 +1,5 @@
 [electricity_sharepoint.sources.m365]
+file_glob = "/General/RDM Data/**/*.*"
 files_per_page = 200
 skip_rows = 7
+max_workers = 6

warehouses/accelerator/extract_load/electricity_sharepoint/extract_and_load.py

@@ -33,15 +33,24 @@
 LOGGER = logging.getLogger(__name__)
 
 
+MAX_WORKERS_DEFAULT = min(8, (os.cpu_count() or 1) + 4)
 EXCEL_ENGINE = "calamine"
+ONE_DAY_SECS = 24 * 60 * 60
 PIPELINE_NAME = "electricity_sharepoint"
 RDM_TIMEZONE = "Europe/London"
 SITE_URL = "https://stfc365.sharepoint.com/sites/ISISSustainability"
-LAG_WINDOW_SECONDS = 24 * 60 * 60
-# It was observed that trying to load too many files concurrently from a sharepoint drive
-# randomly resulted in empty content. Lower the maximum number of threads that can be used
-# to extract the data
-MAX_WORKERS_DEFAULT = min(10, (os.cpu_count() or 1) + 4)
+
+
+def effective_max_workers(max_workers: int | None) -> int:
+    # For high-cpu-count machines the ThreadPoolExecutor default ends up being too high.
+    # It has been observed that too many concurrent threads cause issues with partial file reads from
+    # OneDrive. 8 threads seems to work correctly so clamp this as them maximum...
+
+    return (
+        min(MAX_WORKERS_DEFAULT, max_workers)
+        if max_workers is not None
+        else MAX_WORKERS_DEFAULT
+    )
 
 
 def to_utc(ts: pd.Series) -> pd.Series:
@@ -59,7 +68,7 @@ def read_power_consumption_csv(
     """
     df = pd.read_csv(file_content, skiprows=skip_rows)
     df["DateTime"] = to_utc(
-        pd.to_datetime(df["Date"] + " " + df["Time"], format="%d/%m/%y %H:%M:%S")
+        pd.to_datetime(df["Date"] + " " + df["Time"], format="%d/%m/%y %H:%M:%S")  # type: ignore
     )
     return df.drop(["Date", "Time"], axis=1)
 
@@ -92,34 +101,33 @@ def extract_content_and_read(
     :param items: An iterator of dicts describing the file content
     :param skip_rows: Number of rows in the csv/xlsx files to skip
     :param max_workers (optional): How many threads to use to process the files.
-                                   Defaults to a maximum of MAX_WORKERS_DEFAULT.
+                                   Defaults to a maximum defined by concurrent.futures.ThreadPoolExecutor
     """
 
     # The files are all independent. Process them in parallel and combine for a single yield
     def read_as_dataframe(file_obj: M365DriveItem) -> pd.DataFrame | None:
-        file_content = io.BytesIO(file_obj.read_bytes())
+        file_name = file_obj["file_name"]
+        file_bytes = file_obj.read_bytes()
+        LOGGER.debug(f"Filename '{file_name}' has size {len(file_bytes)} bytes.")
         try:
-            match pathlib.Path(file_obj["file_name"]).suffix:
+            match pathlib.Path(file_name).suffix:
                 case ".csv":
-                    df = read_power_consumption_csv(file_content, skip_rows)
+                    df = read_power_consumption_csv(io.BytesIO(file_bytes), skip_rows)
                 case ".xlsx":
-                    df = read_power_consumption_excel(file_content, skip_rows)
+                    df = read_power_consumption_excel(io.BytesIO(file_bytes), skip_rows)
                 case _:
-                    raise RuntimeError(
-                        f"Unsupported file extension in '{file_obj['file_name']}'"
-                    )
-        except ValueError as exc:
-            LOGGER.warning(
-                f"Error reading '{file_obj['file_name']}': {str(exc)}. Skipping"
-            )
-            df = None
+                    raise RuntimeError(f"Unsupported file extension in '{file_name}'")
+        except pd.errors.EmptyDataError as exc:
+            raise RuntimeError(
+                f"'{file_name} ({len(file_bytes)} bytes)': {str(exc)}"
+            ) from exc
 
+        df["file_name"] = file_obj["file_name"]
         return df
 
     df_batch = None
-    effective_max_workers = min(MAX_WORKERS_DEFAULT, max_workers or MAX_WORKERS_DEFAULT)
     with concurrent.futures.ThreadPoolExecutor(
-        max_workers=effective_max_workers
+        max_workers=effective_max_workers(max_workers)
     ) as executor:
         future_to_file_item = {
             executor.submit(read_as_dataframe, file_obj): file_obj for file_obj in items
@@ -132,24 +140,18 @@ def read_as_dataframe(file_obj: M365DriveItem) -> pd.DataFrame | None:
     yield df_batch
 
 
-# Occasionally pandas raises a EmptyDataError when parsing the CSV content indicating
-# there is no data. This is likely an issue with a file being listed but no content yet available.
-# We skip these files but want to make sure we grab them next time so we use the lag functionality
-# to look back LAG_WINDOW_SECONDS when the next load is run.
-@dlt.resource(merge_key="DateTime")
+@dlt.resource(merge_key="file_name")
 def rdm_data(
     datetime_cur=dlt.sources.incremental(
         "DateTime",
         initial_value=pendulum.DateTime.EPOCH,
-        lag=LAG_WINDOW_SECONDS,
-        last_value_func=max,
     ),
 ) -> Iterator[TDataItems]:
     files = sharepoint(
         site_url=SITE_URL,
-        file_glob="/General/RDM Data/**/*.*",
         extract_content=False,
-        modified_after=datetime_cur.start_value,
+        modified_after=datetime_cur.start_value
+        - pendulum.Duration(seconds=ONE_DAY_SECS),
     )
     reader = files | extract_content_and_read()
     yield from reader