Merge pull request #108 from gulfofmaine/initial_s3_pipeline

Initial S3 timeseries tests

Author: abkfenris

.github/workflows/pipeline_s3_timeseries.yml

@@ -25,7 +25,9 @@ jobs:
             image_name: buoy_retriever_s3_timeseries
             image_tag: ${{ needs.shortsha.outputs.shortsha }}
             # push_image: ${{ github.event_name == 'push' || github.event_name == 'workflow_dispatch' }}
-            # test_command: "pixi run pytest"
+            test_command: "pixi run pytest --cov=. --cov-report=xml --cov-report=term-missing --junitxml=junit.xml -o junit_family=legacy"
+            coverage_flags: s3_timeseries
+            upload_tests: true
 
     # deploy:
     #     needs: [shortsha, build_test_push]

.gitignore

@@ -50,6 +50,7 @@ coverage.xml
 .hypothesis/
 .pytest_cache/
 cover/
+junit.xml
 
 # Translations
 *.mo

Makefile

@@ -51,3 +51,16 @@ user:
 # Open a Django shell inside the backend container
 shell:
 	docker compose exec backend pixi run python manage.py shell
+
+test-common:
+	cd common; uv run pytest --cov=.
+
+test-hohonu:
+	docker build -f pipeline/hohonu/Dockerfile -t buoy_retriever-hohonu .
+	docker run buoy_retriever-hohonu pixi run pytest --cov=.
+
+test-s3-timeseries:
+	docker build -f pipeline/s3_timeseries/Dockerfile -t buoy_retriever-s3_timeseries .
+	docker run buoy_retriever-s3_timeseries pixi run pytest --cov=.
+
+test-all: test-common test-s3-timeseries test-hohonu

README.md

@@ -91,6 +91,10 @@ Spotlight is also accessible to Github Copilot and other agents via a MCP server
 
 Currently there is some testing for common utilities and Hohonu pipelines.
 
-For common, cd into `common/` then `uv run pytest`. Add `--cov=.` to see coverage (along with other pytest-cov options).
+`make test-all` to try to test everything.
 
-For Hohonu, `docker compose exec hohonu pixi run pytest`.
+- Common - `make test-common` or cd into `common/` then `uv run pytest`. Add `--cov=.` to see coverage (along with other pytest-cov options).
+- Pipelines
+  - S3 Timeseries - `make test-s3-timeseries` for more isolated tests, or `docker compose exec/run s3_timeseries pixi run pytest` which will mount volumes (in case snapshots need to be updated).
+    - To run tests with real AWS data, include the `--aws` to `pytest`. This requires AWS credentials that can access the buckets.
+  - Hohonu - `make test-hohonu` for a more isolated test, or  `docker compose exec/run hohonu pixi run pytest` will mount volumes (in case snapshots need to be updated).

common/test_utils/__init__.py

@@ -1,3 +1,4 @@
 from .assets import get_asset_by_name
+from .sensors import get_sensor_by_name
 
-__all__ = ["get_asset_by_name"]
+__all__ = ["get_asset_by_name", "get_sensor_by_name"]

common/test_utils/sensors.py

@@ -0,0 +1,12 @@
+import dagster as dg
+
+
+def get_sensor_by_name(defs, name: str) -> dg.SensorDefinition:
+    """Extract a specific sensor from Dagster Definitions by name.
+
+    Uses the function name, not the full sensor key/prefix.
+    """
+    for sensor in defs.sensors:
+        if sensor.name == name:
+            return sensor
+    raise KeyError(f"Sensor with name {name} not found")

pipeline/hohonu/tests/test_hohonu_pipeline.py

@@ -15,6 +15,7 @@
 def vcr_config():
     return {
         "filter_headers": [("Authorization", "FAKE")],
+        "ignore_hosts": ["spotlight"],
     }
 
 

pipeline/s3_timeseries/Dockerfile

@@ -13,7 +13,7 @@ ENV PYTHONWARNINGS once
 ENV PIXI_FROZEN true
 
 WORKDIR ${HOME}
-COPY ${SOURCE}/pixi.toml ${SOURCE}/pixi.lock ${HOME}/
+COPY ${SOURCE}/pyproject.toml ${SOURCE}/pixi.lock ${HOME}/
 
 RUN --mount=type=cache,id=br_s3_timeseries,target=/home/ioos/.cache/rattler/cache,uid=1000,gid=1000 \
     pixi install -e default
@@ -28,6 +28,8 @@ COPY ${SOURCE}/*.py \
     ${SOURCE}/*.yaml \
     ${HOME}/
 
+COPY ${SOURCE}/tests/ ${HOME}/tests/
+
 # COPY ${SOURCE}/platforms \
 #     ${HOME}/platforms
 

pipeline/s3_timeseries/pipeline.py

@@ -1,11 +1,14 @@
-from datetime import date
+from datetime import date, datetime
 from textwrap import dedent
 from typing import Annotated
 
+import boto3
 import dagster as dg
 import pandas as pd
 import sentry_sdk
 import xarray as xr
+from dagster_aws.s3.sensor import get_objects
+from parse import parse
 from pydantic import BaseModel, Field
 
 from common import assets, config, io
@@ -88,7 +91,7 @@ def monthly_partition_path(self):
         )
 
 
-def defs_for_dataset(dataset: S3TimeseriesDataset) -> dg.Definitions:
+def defs_for_dataset(dataset: S3TimeseriesDataset) -> dg.Definitions:  # noqa: C901
     """Definitions for a single S3 Timeseries dataset."""
     common_asset_kwargs = {
         "key_prefix": ["s3_timeseries", dataset.safe_slug],
@@ -116,7 +119,13 @@ def daily_df(context: dg.AssetExecutionContext, s3fs: S3FSResource) -> pd.DataFr
         """Download daily dataframe from S3."""
         partition_date_string = context.asset_partition_key_for_output()
         partition_date = date.fromisoformat(partition_date_string)
-        day_glob = f"{dataset.config.s3_source.bucket}{dataset.config.s3_source.prefix}{dataset.config.file_pattern.day_pattern.format(partition_date=partition_date)}"
+        day_glob = (
+            dataset.config.s3_source.bucket
+            + dataset.config.s3_source.prefix
+            + dataset.config.file_pattern.day_pattern.format(
+                partition_date=partition_date,
+            )
+        )
 
         context.log.info(
             f"Reading daily data for {partition_date_string} from S3 with glob: {day_glob}",
@@ -137,6 +146,16 @@ def daily_df(context: dg.AssetExecutionContext, s3fs: S3FSResource) -> pd.DataFr
                 daily_dfs.append(df)
         df = pd.concat(daily_dfs)
 
+        # some flexibility in datetime/time column name
+        for col_name in ["datetime", "time"]:
+            try:
+                df[col_name] = pd.to_datetime(df[col_name])
+                df = df.sort_values(col_name)
+            except KeyError:
+                pass
+
+        df = df.reset_index(drop=True)
+
         return df
 
     @dg.asset(
@@ -206,7 +225,103 @@ def monthly_ds(
 
         return ds
 
-    return dg.Definitions(assets=[daily_df, monthly_ds])
+    daily_job = dg.define_asset_job(
+        f"update_{dataset.safe_slug}_daily",
+        selection=[daily_df],
+    )
+
+    @dg.sensor(
+        job=daily_job,
+        name=dataset.safe_slug + "_s3_sensor",
+        minimum_interval_seconds=5 * 60,
+    )
+    def s3_sensor(context: dg.SensorEvaluationContext, s3_credentials: S3Credentials):
+        """Sensor to detect new files for a day in S3."""
+        with sentry_sdk.start_transaction(
+            op=f"{dataset.safe_slug}_s3_sensor",
+            name=f"S3 Sensor for {dataset.safe_slug}",
+        ):
+            since_time = context.cursor or None
+            if since_time:
+                since_time = datetime.fromisoformat(since_time)
+
+            client = boto3.client(
+                "s3",
+                aws_access_key_id=s3_credentials.access_key_id,
+                aws_secret_access_key=s3_credentials.secret_access_key,
+            )
+            file_start = dataset.config.file_pattern.day_pattern.split("{")[0]
+
+            new_s3_keys = get_objects(
+                bucket=dataset.config.s3_source.bucket,
+                prefix=file_start,
+                since_last_modified=since_time,
+                client=client,
+            )
+            if not new_s3_keys:
+                return dg.SkipReason("No new files found in S3.")
+
+            existing_partitions = set()
+            known_partitions = set(daily_partitions.get_partition_keys())
+            for run in context.instance.get_runs(
+                filters=dg.RunsFilter(
+                    job_name=daily_job.name,
+                    statuses=[
+                        dg.DagsterRunStatus.QUEUED,
+                        dg.DagsterRunStatus.STARTING,
+                        dg.DagsterRunStatus.STARTED,
+                        dg.DagsterRunStatus.NOT_STARTED,
+                    ],
+                ),
+            ):
+                try:
+                    existing_partitions.add(run.tags["dagster/partition"])
+                except KeyError:
+                    pass
+
+            for key in new_s3_keys:
+                object_key = key.get("Key")
+
+                object_name = object_key.removeprefix(dataset.config.s3_source.prefix)
+                name_pattern = dataset.config.file_pattern.day_pattern.replace(
+                    "*",
+                    "{}",
+                )
+                result = parse(name_pattern, object_name)
+                dt = result.named["partition_date"].strftime("%Y-%m-%d")
+
+                if dt not in existing_partitions:
+                    last_modified = key.get("LastModified")
+
+                    if since_time is None or last_modified > since_time:
+                        existing_partitions.add(dt)
+                        run_key = f"{dt}_{last_modified.isoformat()}"
+
+                        if dt in known_partitions:
+                            yield dg.RunRequest(
+                                run_key=run_key,
+                                partition_key=dt,
+                            )
+                        else:
+                            context.log.info(
+                                f"Skipping partition {dt} as it is not a known partition",
+                            )
+
+            latest_key_dt = max(key.get("LastModified") for key in new_s3_keys)
+            context.update_cursor(latest_key_dt.isoformat())
+
+    dataset_assets = [daily_df, monthly_ds]
+
+    return dg.Definitions(
+        assets=dataset_assets,
+        sensors=[
+            s3_sensor,
+            dg.AutomationConditionSensorDefinition(
+                dataset.safe_slug + "_automation_sensor",
+                target=dataset_assets,
+            ),
+        ],
+    )
 
 
 @dg.definitions