Add OrchestratorInfo class and update alert handling for class

Author: quickcoffee

elementary/monitor/alerts/alert.py

@@ -7,11 +7,23 @@
     ReportLinkData,
 )
 from elementary.utils.log import get_logger
+from elementary.utils.pydantic_shim import BaseModel
 from elementary.utils.time import DATETIME_WITH_TIMEZONE_FORMAT
 
 logger = get_logger(__name__)
 
 
+class OrchestratorInfo(BaseModel):
+    """Structured orchestrator metadata for alerts."""
+
+    job_id: Optional[str] = None
+    job_name: Optional[str] = None
+    run_id: Optional[str] = None
+    orchestrator: Optional[str] = None
+    job_url: Optional[str] = None
+    run_url: Optional[str] = None
+
+
 class AlertModel:
     def __init__(
         self,
@@ -98,7 +110,7 @@ def get_report_link(self) -> Optional[ReportLinkData]:
         raise NotImplementedError
 
     @property
-    def orchestrator_info(self) -> Optional[Dict[str, str]]:
+    def orchestrator_info(self) -> Optional[OrchestratorInfo]:
         """Returns structured orchestrator metadata if available."""
         if not any(
             [
@@ -111,18 +123,11 @@ def orchestrator_info(self) -> Optional[Dict[str, str]]:
         ):
             return None
 
-        info = {}
-        if self.job_id:
-            info["job_id"] = self.job_id
-        if self.job_name:
-            info["job_name"] = self.job_name
-        if self.job_run_id:
-            info["run_id"] = self.job_run_id
-        if self.orchestrator:
-            info["orchestrator"] = self.orchestrator
-        if self.job_url:
-            info["job_url"] = self.job_url
-        if self.job_run_url:
-            info["run_url"] = self.job_run_url
-
-        return info
+        return OrchestratorInfo(
+            job_id=self.job_id or None,
+            job_name=self.job_name or None,
+            run_id=self.job_run_id or None,
+            orchestrator=self.orchestrator or None,
+            job_url=self.job_url or None,
+            run_url=self.job_run_url or None,
+        )

elementary/monitor/alerts/alert_messages/builder.py

@@ -35,6 +35,7 @@
     TextStyle,
 )
 from elementary.messages.message_body import Color, MessageBlock, MessageBody
+from elementary.monitor.alerts.alert import OrchestratorInfo
 from elementary.monitor.alerts.alert_messages.alert_fields import AlertField
 from elementary.monitor.alerts.alerts_groups.alerts_group import AlertsGroup
 from elementary.monitor.alerts.alerts_groups.base_alerts_group import BaseAlertsGroup
@@ -111,7 +112,7 @@ def _get_run_alert_subtitle_block(
         suppression_interval: Optional[int] = None,
         env: Optional[str] = None,
         links: list[ReportLinkData] = [],
-        orchestrator_info: Optional[Dict[str, str]] = None,
+        orchestrator_info: Optional[OrchestratorInfo] = None,
     ) -> LinesBlock:
         summary = []
         summary.append((type.capitalize() + ":", name))
@@ -124,9 +125,9 @@ def _get_run_alert_subtitle_block(
         # Initialize subtitle lines with summary
         subtitle_lines = []
 
-        if orchestrator_info and orchestrator_info.get("job_name"):
-            orchestrator_name = orchestrator_info.get("orchestrator", "orchestrator")
-            job_info_text = f"{orchestrator_info['job_name']} (via {orchestrator_name})"
+        if orchestrator_info and orchestrator_info.job_name:
+            orchestrator_name = orchestrator_info.orchestrator or "orchestrator"
+            job_info_text = f"{orchestrator_info.job_name} (via {orchestrator_name})"
 
             # Create job info with inline orchestrator link
             orchestrator_link = create_orchestrator_link(orchestrator_info)
@@ -162,7 +163,7 @@ def _get_run_alert_subtitle_block(
             all_links.append((link.text, link.url, link.icon))
 
         # Add orchestrator link if available (only if not already added inline)
-        if orchestrator_info and not orchestrator_info.get("job_name"):
+        if orchestrator_info and not orchestrator_info.job_name:
             orchestrator_link = create_orchestrator_link(orchestrator_info)
             if orchestrator_link:
                 all_links.append(

elementary/monitor/data_monitoring/alerts/integrations/utils/orchestrator_link.py

@@ -1,6 +1,7 @@
-from typing import Dict, Optional
+from typing import Optional
 
 from elementary.messages.blocks import Icon
+from elementary.monitor.alerts.alert import OrchestratorInfo
 from elementary.utils.pydantic_shim import BaseModel
 
 
@@ -12,37 +13,37 @@ class OrchestratorLinkData(BaseModel):
 
 
 def create_orchestrator_link(
-    orchestrator_info: Dict[str, str]
+    orchestrator_info: OrchestratorInfo,
 ) -> Optional[OrchestratorLinkData]:
     """Create an orchestrator link from orchestrator info if URL is available."""
-    if not orchestrator_info or not orchestrator_info.get("run_url"):
+    if not orchestrator_info or not orchestrator_info.run_url:
         return None
 
-    orchestrator = orchestrator_info.get("orchestrator", "orchestrator")
+    orchestrator = orchestrator_info.orchestrator or "orchestrator"
 
     return OrchestratorLinkData(
-        url=orchestrator_info["run_url"],
+        url=orchestrator_info.run_url,
         text=f"View in {orchestrator}",
         orchestrator=orchestrator,
         icon=Icon.LINK,
     )
 
 
 def create_job_link(
-    orchestrator_info: Dict[str, str]
+    orchestrator_info: OrchestratorInfo,
 ) -> Optional[OrchestratorLinkData]:
     """Create a job-level orchestrator link if job URL is available."""
-    if not orchestrator_info or not orchestrator_info.get("job_url"):
+    if not orchestrator_info or not orchestrator_info.job_url:
         return None
 
-    orchestrator = orchestrator_info.get("orchestrator", "orchestrator")
-    job_name = orchestrator_info.get("job_name", "Job")
+    orchestrator = orchestrator_info.orchestrator or "orchestrator"
+    job_name = orchestrator_info.job_name or "Job"
 
     # Capitalize orchestrator name for display
     display_name = orchestrator.replace("_", " ").title()
 
     return OrchestratorLinkData(
-        url=orchestrator_info["job_url"],
+        url=orchestrator_info.job_url,
         text=f"{job_name} in {display_name}",
         orchestrator=orchestrator,
         icon=Icon.GEAR,

tests/unit/alerts/alert_messages/test_orchestrator_message_simple.py

@@ -144,14 +144,14 @@ def test_orchestrator_info_property_integration(self):
         orchestrator_info = alert.orchestrator_info
 
         assert orchestrator_info is not None
-        assert orchestrator_info["job_name"] == "integration_test"
-        assert orchestrator_info["run_id"] == "run_123"
-        assert orchestrator_info["orchestrator"] == "airflow"
+        assert orchestrator_info.job_name == "integration_test"
+        assert orchestrator_info.run_id == "run_123"
+        assert orchestrator_info.orchestrator == "airflow"
         assert (
-            orchestrator_info["job_url"]
+            orchestrator_info.job_url
             == "https://airflow.example.com/job/integration_test"
         )
-        assert orchestrator_info["run_url"] == "https://airflow.example.com/run/123"
+        assert orchestrator_info.run_url == "https://airflow.example.com/run/123"
 
         # Test message includes this data
         alert.get_report_link = Mock(return_value=None)

tests/unit/alerts/test_orchestrator_integration.py

@@ -1,6 +1,7 @@
 import pytest
 
 from elementary.messages.blocks import Icon
+from elementary.monitor.alerts.alert import OrchestratorInfo
 from elementary.monitor.alerts.model_alert import ModelAlertModel
 from elementary.monitor.alerts.source_freshness_alert import SourceFreshnessAlertModel
 from elementary.monitor.alerts.test_alert import TestAlertModel
@@ -15,12 +16,12 @@ class TestOrchestratorLinkCreation:
     """Test orchestrator link creation functionality."""
 
     def test_create_orchestrator_link_with_valid_data(self):
-        orchestrator_info = {
-            "job_name": "nightly_load",
-            "run_id": "12345",
-            "orchestrator": "airflow",
-            "run_url": "https://airflow.example.com/run/12345",
-        }
+        orchestrator_info = OrchestratorInfo(
+            job_name="nightly_load",
+            run_id="12345",
+            orchestrator="airflow",
+            run_url="https://airflow.example.com/run/12345",
+        )
 
         link = create_orchestrator_link(orchestrator_info)
 
@@ -30,25 +31,25 @@ def test_create_orchestrator_link_with_valid_data(self):
         assert link.icon == Icon.LINK
 
     def test_create_orchestrator_link_without_url_returns_none(self):
-        orchestrator_info = {
-            "job_name": "nightly_load",
-            "orchestrator": "airflow"
+        orchestrator_info = OrchestratorInfo(
+            job_name="nightly_load",
+            orchestrator="airflow",
             # No run_url
-        }
+        )
 
         link = create_orchestrator_link(orchestrator_info)
         assert link is None
 
     def test_create_orchestrator_link_with_empty_info_returns_none(self):
-        link = create_orchestrator_link({})
+        link = create_orchestrator_link(OrchestratorInfo())
         assert link is None
 
     def test_create_job_link_with_valid_data(self):
-        orchestrator_info = {
-            "job_name": "nightly_load",
-            "orchestrator": "airflow",
-            "job_url": "https://airflow.example.com/job/nightly_load",
-        }
+        orchestrator_info = OrchestratorInfo(
+            job_name="nightly_load",
+            orchestrator="airflow",
+            job_url="https://airflow.example.com/job/nightly_load",
+        )
 
         link = create_job_link(orchestrator_info)
 
@@ -59,11 +60,11 @@ def test_create_job_link_with_valid_data(self):
         assert link.icon == Icon.GEAR
 
     def test_create_job_link_without_url_returns_none(self):
-        orchestrator_info = {
-            "job_name": "nightly_load",
-            "orchestrator": "airflow"
+        orchestrator_info = OrchestratorInfo(
+            job_name="nightly_load",
+            orchestrator="airflow",
             # No job_url
-        }
+        )
 
         link = create_job_link(orchestrator_info)
         assert link is None
@@ -92,11 +93,11 @@ def test_test_alert_orchestrator_info_with_complete_data(self):
 
         info = alert.orchestrator_info
         assert info is not None
-        assert info["job_name"] == "nightly_load"
-        assert info["run_id"] == "12345"
-        assert info["orchestrator"] == "airflow"
-        assert info["job_url"] == "https://airflow.example.com/job/nightly_load"
-        assert info["run_url"] == "https://airflow.example.com/run/12345"
+        assert info.job_name == "nightly_load"
+        assert info.run_id == "12345"
+        assert info.orchestrator == "airflow"
+        assert info.job_url == "https://airflow.example.com/job/nightly_load"
+        assert info.run_url == "https://airflow.example.com/run/12345"
 
     def test_test_alert_orchestrator_info_with_minimal_data(self):
         alert = TestAlertModel(
@@ -115,8 +116,10 @@ def test_test_alert_orchestrator_info_with_minimal_data(self):
 
         info = alert.orchestrator_info
         assert info is not None
-        assert info["job_name"] == "test_job"
-        assert len(info) == 1
+        assert info.job_name == "test_job"
+        # Other fields should be None
+        assert info.run_id is None
+        assert info.orchestrator is None
 
     def test_test_alert_orchestrator_info_with_no_data_returns_none(self):
         alert = TestAlertModel(
@@ -152,10 +155,10 @@ def test_model_alert_orchestrator_info(self):
 
         info = alert.orchestrator_info
         assert info is not None
-        assert info["job_name"] == "nightly_build"
-        assert info["run_id"] == "67890"
-        assert info["orchestrator"] == "dbt_cloud"
-        assert info["run_url"] == "https://cloud.getdbt.com/run/67890"
+        assert info.job_name == "nightly_build"
+        assert info.run_id == "67890"
+        assert info.orchestrator == "dbt_cloud"
+        assert info.run_url == "https://cloud.getdbt.com/run/67890"
 
     def test_source_freshness_alert_orchestrator_info(self):
         alert = SourceFreshnessAlertModel(
@@ -174,9 +177,9 @@ def test_source_freshness_alert_orchestrator_info(self):
 
         info = alert.orchestrator_info
         assert info is not None
-        assert info["job_name"] == "freshness_check"
-        assert info["orchestrator"] == "airflow"
-        assert info["run_url"] == "https://airflow.example.com/run/111"
+        assert info.job_name == "freshness_check"
+        assert info.orchestrator == "airflow"
+        assert info.run_url == "https://airflow.example.com/run/111"
 
 
 class TestOrchestratorInfoEdgeCases:
@@ -201,7 +204,7 @@ def test_orchestrator_info_with_empty_strings(self):
         info = alert.orchestrator_info
         assert info is None
 
-    def test_orchestrator_info_filters_none_values(self):
+    def test_orchestrator_info_sets_none_for_empty_values(self):
         alert = TestAlertModel(
             id="test_id",
             test_unique_id="test_unique_id",
@@ -219,10 +222,9 @@ def test_orchestrator_info_filters_none_values(self):
 
         info = alert.orchestrator_info
         assert info is not None
-        assert "job_name" in info
-        assert "orchestrator" not in info
-        assert "run_url" not in info
-        assert len(info) == 1
+        assert info.job_name == "valid_job"
+        assert info.orchestrator is None
+        assert info.run_url is None
 
     def test_orchestrator_info_with_only_run_id(self):
         alert = TestAlertModel(
@@ -240,8 +242,9 @@ def test_orchestrator_info_with_only_run_id(self):
 
         info = alert.orchestrator_info
         assert info is not None
-        assert info["run_id"] == "12345"
-        assert len(info) == 1
+        assert info.run_id == "12345"
+        assert info.job_name is None
+        assert info.orchestrator is None
 
     def test_orchestrator_info_with_only_orchestrator(self):
         alert = TestAlertModel(
@@ -259,8 +262,9 @@ def test_orchestrator_info_with_only_orchestrator(self):
 
         info = alert.orchestrator_info
         assert info is not None
-        assert info["orchestrator"] == "dbt_cloud"
-        assert len(info) == 1
+        assert info.orchestrator == "dbt_cloud"
+        assert info.job_name is None
+        assert info.run_id is None
 
 
 class TestOrchestratorLinkDataModel: