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

Author: sandeshit

alert_system/etl/Gdacs_flood/config.py

@@ -0,0 +1,13 @@
+from typing import Dict, TypedDict
+
+
+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=10))  # 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,77 +244,127 @@ 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 _construct_filter_for_past_events(self, impact_metadata: list[dict]) -> str:
-        filters = []
+    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."""
+        try:
+            self.process_event_items(extraction_run_id, correlation_id, is_past_event)
+        except Exception as e:
+            logger.warning(f"Connector run failed: {e}", exc_info=True)
+            raise
+
 
-        for detail in impact_metadata:
-            category = detail.get("category")
-            type_ = detail.get("type")
-            value = detail.get("value")
+class PastEventExtractionClass:
+    LOOKBACK_WEEKS = 520
 
-            if category and type_ and value is not None:
+    def __init__(self, extractor: BaseExtractionClass):
+        self.extractor = extractor
+        self.base_url = extractor.base_url
+
+    def _impact_filter(self, impact_metadata: list[dict]) -> str:
+        filters = []
+
+        for data in impact_metadata or []:
+            if data.get("category") and data.get("type") and data.get("value") is not None:
                 filters.append(
-                    f"monty:impact_detail.category = '{category}' AND "
-                    f"monty:impact_detail.type = '{type_}' AND "
-                    f"monty:impact_detail.value >= {value}"
+                    f"monty:impact_detail.category = '{data['category']}' AND "
+                    f"monty:impact_detail.type = '{data['type']}' AND "
+                    f"monty:impact_detail.value >= {data['value']}"
                 )
 
-        return " OR ".join(f"({f})" for f in filters) if filters else ""
+        return " OR ".join(f"({filter})" for filter in filters)
 
-    def fetch_past_events(self, load_obj):
-        if not self.impact_collection_type:
-            logger.warning(f"Impact does not exist for event {load_obj}")
-            return
-        start_time = datetime.now(timezone.utc) - timedelta(weeks=16)  # NOTE: Arbitrary value for lookback.
-        filters = [self._construct_filter_for_past_events(load_obj.impact_metadata)]
-        impact_data = self.fetch_stac_data(
-            self.base_url,
-            build_stac_search(
-                collections=self.impact_collection_type,
-                additional_filters=filters,
-                datetime_range=f"{start_time.isoformat()}/{datetime.now(timezone.utc).isoformat()}",
-            ),
-        )
-        load_obj_corr_id = load_obj.correlation_id
-        related_ids = []
-        logger.info(f"Fetching past event for event={load_obj.id}")
+    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}"
+
+    def _collect_corr_ids(self, features, exclude: str) -> set[str]:
         corr_ids = set()
-        for feature in impact_data:
-            corr_id = self._get_correlation_id(feature)
-            if corr_id and corr_id != load_obj_corr_id:
+        for feature in features or []:
+            corr_id = self.extractor._get_correlation_id(feature)
+            if corr_id and corr_id != exclude:
                 corr_ids.add(corr_id)
+        return corr_ids
+
+    def find_related_corr_ids(self, load_obj: LoadItem) -> set[str]:
+        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=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.
+        # if self.extractor.hazard_collection_type:
+        #     hazard_filter = self._hazard_filter(
+        #         load_obj.severity_unit,
+        #         load_obj.severity_value,
+        #     )
+        #     features = self.extractor.fetch_stac_data(
+        #         self.base_url,
+        #         build_stac_search(
+        #             collections=self.extractor.hazard_collection_type,
+        #             additional_filters=[hazard_filter],
+        #             datetime_range=f"{start.isoformat()}/{end.isoformat()}",
+        #         ),
+        #     )
+        #     corr_ids |= self._collect_corr_ids(features, load_obj.correlation_id)
+
+        return corr_ids
+
+    def extract_past_events(self, load_obj: LoadItem) -> None:
+        corr_ids = self.find_related_corr_ids(load_obj)
 
         if not corr_ids:
             return
 
         existing_items = LoadItem.objects.filter(correlation_id__in=corr_ids)
-        existing_map = {item.correlation_id: item for item in existing_items}
+        existing_map = {i.correlation_id: i for i in existing_items}
+
+        related_ids = []
 
         for corr_id in corr_ids:
             item = existing_map.get(corr_id)
+
+            if not item:
+                self.extractor.run(
+                    extraction_run_id=load_obj.extraction_run_id,
+                    correlation_id=corr_id,
+                    is_past_event=True,
+                )
+                item = LoadItem.objects.filter(correlation_id=corr_id).first()
+
             if item:
                 related_ids.append(item.id)
                 item.related_montandon_events.add(load_obj.id)
-            else:
-                self.run(extraction_run_id=load_obj.extraction_run_id, correlation_id=corr_id, is_past_event=True)
-                new_item = LoadItem.objects.filter(correlation_id=corr_id).first()
-                if new_item:
-                    related_ids.append(new_item.id)
-                    new_item.related_montandon_events.add(load_obj.id)
 
         if related_ids:
             load_obj.related_montandon_events.set(related_ids)
-
-    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)
-            raise

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,12 @@
+# 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

@@ -31,25 +31,22 @@ def compute_buildings_exposed(self, metadata_list) -> int:
     # NOTE: This logic will change with changes in montandon.
     def process_impact(self, impact_items) -> BaseTransformerClass.ImpactType:
         metadata = []
-        largest_values_metadata = {}
         for item in impact_items:
             properties = item.resp_data.get("properties", {})
             impact_detail = properties.get("monty:impact_detail", {})
             category = impact_detail.get("category")
             type_ = impact_detail.get("type")
             value = impact_detail.get("value")
-            if category and type_:
-                key = (category, type_)
-
-                if key not in largest_values_metadata or value > largest_values_metadata[key]["value"]:
-                    largest_values_metadata[key] = {
+            if category == "people" and type_ == "affected_total":
+                metadata = [
+                    {
                         "category": category,
                         "type": type_,
                         "value": value,
                         "unit": impact_detail.get("unit", ""),
                         "estimate_type": impact_detail.get("estimate_type", ""),
                     }
-        metadata.extend(largest_values_metadata.values())
+                ]
         return {
             "people_exposed": self.compute_people_exposed(metadata),
             "buildings_exposed": self.compute_buildings_exposed(metadata),
@@ -79,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"),
         }