fixup! chore(etl): separate past event extraction into different class

Author: sandeshit

alert_system/etl/Gdacs_flood/config.py

@@ -0,0 +1,12 @@
+from typing import TypedDict, Dict
+
+class ExtractionConfig(TypedDict):
+    event_collection_type: str
+    hazard_collection_type: str | None
+    impact_collection_type: str | None
+
+    filter_event: Dict | None
+    filter_hazard: Dict | None
+    filter_impact: Dict | None
+
+    people_exposed_threshold : int

alert_system/etl/Usgs_earthquake/config.py

@@ -1,14 +1,16 @@
 import logging
 from abc import ABC
-from datetime import datetime, timedelta, timezone
-from typing import Dict, Generator, List, Optional
+from datetime import timedelta
+from typing import Dict, Generator, List, Optional, Type
 
 import httpx
 from django.db import transaction
+from django.utils import timezone
 
 from alert_system.helpers import build_stac_search
 from alert_system.models import Connector, ExtractionItem, LoadItem
 
+from .config import ExtractionConfig
 from .loader import BaseLoaderClass
 from .transform import BaseTransformerClass
 
@@ -29,19 +31,28 @@ class BaseExtractionClass(ABC):
     """
 
     event_collection_type: str
-    hazard_collection_type: str | None
-    impact_collection_type: str | None
+    transformer_class: Type[BaseTransformerClass]
+    loader_class: Type[BaseLoaderClass]
+
+    hazard_collection_type: Optional[str] = None
+    impact_collection_type: Optional[str] = None
+
     filter_event: Optional[Dict] = None
     filter_hazard: Optional[Dict] = None
     filter_impact: Optional[Dict] = None
-    transformer_class: type[BaseTransformerClass]
-    loader_class: type[BaseLoaderClass]
+
+    config: ExtractionConfig
 
     def __init__(self, connector: Connector):
         self.connector = connector
         self.base_url = connector.source_url.rstrip("/")
+        self.load_config()
         self._validate_required_attributes()
 
+    def load_config(self):
+        for key, value in self.config.items():
+            setattr(self, key, value)
+
     def _validate_required_attributes(self):
         missing_attr = []
         if not getattr(self, "event_collection_type", None):
@@ -96,10 +107,10 @@ def get_datetime_filter(self) -> str:
             ISO 8601 datetime range string
         """
 
-        now = datetime.now(timezone.utc)
+        now = timezone.now()
         last_run = self.connector.last_success_run
 
-        start_time = last_run if last_run else (now - timedelta(days=20))  # NOTE: Arbitrary value for failure case.
+        start_time = last_run if last_run else (now - timedelta(days=30))  # NOTE: Arbitrary value for failure case.
         return f"{start_time.isoformat()}/{now.isoformat()}"
 
     def _save_stac_item(self, stac_id: str, defaults: Dict) -> Optional[ExtractionItem]:
@@ -200,7 +211,7 @@ def process_event_items(self, extraction_run_id: str, correlation_id: str | None
                 ),
             )
         except Exception as e:
-            logger.error(f"Failed to fetch events: {e}")
+            logger.warning(f"Failed to fetch events: {e}")
             raise
 
         for feature in event_items:
@@ -233,23 +244,20 @@ def process_event_items(self, extraction_run_id: str, correlation_id: str | None
                     logger.info(f"Successfully processed event {event_id}")
 
             except Exception as e:
-                logger.error(f"Failed to process event {event_id}: {e}", exc_info=True)
+                logger.warning(f"Failed to process event {event_id}: {e}", exc_info=True)
                 raise
 
     def run(self, extraction_run_id: str, correlation_id: str | None = None, is_past_event: bool = False) -> None:
         """Main entry point for running the connector."""
-        logger.info(f"Starting connector run for {self.connector}")
-
         try:
             self.process_event_items(extraction_run_id, correlation_id, is_past_event)
-            logger.info("Connector run completed successfully")
         except Exception as e:
-            logger.error(f"Connector run failed: {e}", exc_info=True)
+            logger.warning(f"Connector run failed: {e}", exc_info=True)
             raise
 
 
 class PastEventExtractionClass:
-    LOOKBACK_WEEKS = 16
+    LOOKBACK_WEEKS = 520
 
     def __init__(self, extractor: BaseExtractionClass):
         self.extractor = extractor
@@ -268,6 +276,13 @@ def _impact_filter(self, impact_metadata: list[dict]) -> str:
 
         return " OR ".join(f"({filter})" for filter in filters)
 
+    def _country_filter(self, country_codes) -> list[str]:
+        filters = []
+        if country_codes:
+            country_cql = " OR ".join(f"a_contains(monty:country_codes, '{code}')" for code in country_codes)
+            filters.append(f"({country_cql})")
+        return filters
+
     def _hazard_filter(self, unit: str, value: int) -> str:
         return f"monty:hazard_detail.severity_unit = '{unit}' AND " f"monty:hazard_detail.severity_value >= {value}"
 
@@ -280,21 +295,31 @@ def _collect_corr_ids(self, features, exclude: str) -> set[str]:
         return corr_ids
 
     def find_related_corr_ids(self, load_obj: LoadItem) -> set[str]:
-        start = datetime.now(timezone.utc) - timedelta(weeks=self.LOOKBACK_WEEKS)
-        end = datetime.now(timezone.utc)
+        start = timezone.now() - timedelta(weeks=self.LOOKBACK_WEEKS)
+        end = timezone.now()
 
         corr_ids = set()
 
         if self.extractor.impact_collection_type:
             impact_filter = self._impact_filter(load_obj.impact_metadata)
+            country_filters = self._country_filter(load_obj.country_codes)
+
+            additional_filters = []
+
+            if impact_filter:
+                additional_filters.append(impact_filter)
+
+            additional_filters.extend(country_filters)
+
             features = self.extractor.fetch_stac_data(
                 self.base_url,
                 build_stac_search(
                     collections=self.extractor.impact_collection_type,
-                    additional_filters=[impact_filter] if impact_filter else [],
+                    additional_filters=additional_filters,
                     datetime_range=f"{start.isoformat()}/{end.isoformat()}",
                 ),
             )
+
             corr_ids |= self._collect_corr_ids(features, load_obj.correlation_id)
 
         # NOTE: Returns too many correlation_ids.

alert_system/etl/base/config.py

@@ -41,13 +41,15 @@ def load(self, transformed_data: Dict, connector: Connector, run_id: str, is_pas
                 "total_people_exposed": transformed_data.get("people_exposed"),
                 "total_buildings_exposed": transformed_data.get("buildings_exposed"),
                 "impact_metadata": transformed_data.get("impact_metadata"),
+                "start_datetime": transformed_data.get("start_datetime"),
+                "end_datetime": transformed_data.get("end_datetime"),
                 "item_eligible": is_item_eligible,
                 "is_past_event": is_past_event,
                 "extraction_run_id": run_id,
             },
         )
 
         action = "Created" if created else "Updated"
-        logger.info(f"{action} Event for correlation_id={correlation_id}")
+        logger.info(f"{action} Event for {correlation_id=}")
 
         return load_obj

alert_system/etl/base/extraction.py

@@ -1,5 +1,6 @@
 import logging
 from abc import ABC, abstractmethod
+from datetime import datetime
 from typing import List, Optional, TypedDict
 
 from alert_system.models import ExtractionItem
@@ -23,6 +24,8 @@ class EventType(TypedDict):
         title: str
         description: str
         country: str
+        start_datetime: datetime
+        end_datetime: datetime
 
     def __init__(
         self, event_obj: ExtractionItem, hazard_obj: Optional[ExtractionItem] = None, impact_obj: List[ExtractionItem] = []

alert_system/etl/base/loader.py

@@ -0,0 +1,14 @@
+# NOTE: Store Config files here. Might need to refactor if source supports filtering with hazards.
+from alert_system.etl.base.config import ExtractionConfig
+
+gdacs_flood_config: ExtractionConfig = {
+    "event_collection_type": "gdacs-events",
+    "hazard_collection_type": "gdacs-hazards",
+    "impact_collection_type": "gdacs-impacts",
+    "filter_event": {
+        "hazard_codes": ["FL", "MH0600", "nat-hyd-flo-flo"]
+    },
+    "filter_hazard": None,
+    "filter_impact": None,
+    "people_exposed_threshold": 500,
+}

alert_system/etl/base/transform.py

@@ -10,9 +10,6 @@
 
 
 class GdacsFloodExtraction(BaseExtractionClass):
-    event_collection_type = gdacs_flood_config.event_collection_type
-    hazard_collection_type = getattr(gdacs_flood_config, "hazard_collection_type", None)
-    impact_collection_type = getattr(gdacs_flood_config, "impact_collection_type", None)
-    filter_event = getattr(gdacs_flood_config, "filter_event", None)
+    config = gdacs_flood_config
     transformer_class = GdacsTransformer
     loader_class = GdacsLoader

alert_system/etl/gdacs_flood/config.py

@@ -4,10 +4,9 @@
 
 
 class GdacsLoader(BaseLoaderClass):
-    people_exposed_threshold = gdacs_flood_config.people_exposed_threshold
 
     # NOTE: Add additional changes to the filter here. This is example only.
     def filter_eligible_items(self, load_obj):
-        if load_obj.get("people_exposed") > GdacsLoader.people_exposed_threshold:
+        if load_obj.get("people_exposed") > gdacs_flood_config["people_exposed_threshold"]:
             return True
         return False

…ert_system/etl/Gdacs_flood/extraction.py …ert_system/etl/gdacs_flood/extraction.pyalert_system/etl/Gdacs_flood/extraction.py renamed to alert_system/etl/gdacs_flood/extraction.py alert_system/etl/Gdacs_flood/extraction.py renamed to alert_system/etl/gdacs_flood/extraction.py

@@ -76,4 +76,6 @@ def process_event(self, event_item) -> BaseTransformerClass.EventType:
             "title": properties.get("title", ""),
             "description": properties.get("description", ""),
             "country": properties.get("monty:country_codes", ""),
+            "start_datetime": properties.get("start_datetime"),
+            "end_datetime": properties.get("end_datetime"),
         }