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

Author: sandeshit

alert_system/etl/Gdacs_flood/config.py

@@ -4,7 +4,7 @@ def __init__(self):
         self.event_collection_type = "gdacs-events"
         self.hazard_collection_type = "gdacs-hazards"
         self.impact_collection_type = "gdacs-impacts"
-        self.people_exposed_threshold = 5
+        self.people_exposed_threshold = 500
         self.filter_event = {"hazard_codes": ["FL", "MH0600", "nat-hyd-flo-flo"]}
 
 

alert_system/etl/Gdacs_flood/transform.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),

alert_system/etl/Usgs_earthquake/transform.py

@@ -28,6 +28,7 @@ def compute_buildings_exposed(self, metadata_list) -> int:
                 return data["value"]
         return 0
 
+    # NOTE: To be changed.
     def process_impact(self, impact_items) -> BaseTransformerClass.ImpactType:
         metadata = []
         values_metadata = {}

alert_system/etl/base/extraction.py

@@ -99,7 +99,7 @@ def get_datetime_filter(self) -> str:
         now = datetime.now(timezone.utc)
         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=20))  # 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]:
@@ -236,74 +236,110 @@ def process_event_items(self, extraction_run_id: str, correlation_id: str | None
                 logger.error(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."""
+        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
 
-        for detail in impact_metadata:
-            category = detail.get("category")
-            type_ = detail.get("type")
-            value = detail.get("value")
 
-            if category and type_ and value is not None:
+class PastEventExtractionClass:
+    LOOKBACK_WEEKS = 16
+
+    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 _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 = datetime.now(timezone.utc) - timedelta(weeks=self.LOOKBACK_WEEKS)
+        end = datetime.now(timezone.utc)
+
+        corr_ids = set()
+
+        if self.extractor.impact_collection_type:
+            impact_filter = self._impact_filter(load_obj.impact_metadata)
+            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 [],
+                    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/mappings.py

@@ -1,24 +1,18 @@
 from dataclasses import dataclass
 
 from alert_system.etl.base.extraction import BaseExtractionClass
-from alert_system.etl.base.loader import BaseLoaderClass
-from alert_system.etl.base.transform import BaseTransformerClass
 from alert_system.etl.Gdacs_flood.extraction import GdacsFloodExtraction
-from alert_system.etl.Gdacs_flood.loader import GdacsLoader
-from alert_system.etl.Gdacs_flood.transform import GdacsTransformer
+from alert_system.etl.Usgs_earthquake.extraction import USGSEarthquakeExtraction
 from alert_system.models import Connector
 
 
 # NOTE: Store all the mappings here.
 @dataclass
 class ConnectorClasses:
     extractor: type[BaseExtractionClass]
-    transfomer: type[BaseTransformerClass]
-    loader: type[BaseLoaderClass]
 
 
 CONNECTOR_REGISTRY = {
-    Connector.ConnectorType.GDACS_FLOOD: ConnectorClasses(
-        extractor=GdacsFloodExtraction, transfomer=GdacsTransformer, loader=GdacsLoader
-    ),
+    Connector.ConnectorType.GDACS_FLOOD: ConnectorClasses(extractor=GdacsFloodExtraction),
+    Connector.ConnectorType.USGS_EARTHQUAKE: ConnectorClasses(extractor=USGSEarthquakeExtraction),
 }

alert_system/tasks.py

@@ -5,6 +5,7 @@
 from celery.exceptions import MaxRetriesExceededError
 from django.db import transaction
 
+from alert_system.etl.base.extraction import PastEventExtractionClass
 from api.models import Event
 
 from .helpers import get_connector_processor, set_connector_status
@@ -75,6 +76,7 @@ def fetch_past_events_from_monty(self, extraction_run_id):
         connector_id = first_item.connector.id
 
         processor, _ = get_connector_processor(connector_id)
+        past_event_extraction_service = PastEventExtractionClass(processor)
 
         # Process each eligible item
         processed = 0
@@ -83,7 +85,7 @@ def fetch_past_events_from_monty(self, extraction_run_id):
         for load_obj in eligible_items.iterator():
             try:
                 with transaction.atomic():
-                    processor.fetch_past_events(load_obj)
+                    past_event_extraction_service.extract_past_events(load_obj=load_obj)
                     processed += 1
             except Exception as e:
                 failed += 1