Add ExternalDeploymentSensor

astronomer/providers/core/example_dags/example_astro.py

@@ -0,0 +1,63 @@
+import time
+from datetime import datetime
+
+from airflow import DAG
+from airflow.decorators import task
+from airflow.operators.trigger_dagrun import TriggerDagRunOperator
+
+from astronomer.providers.core.sensors.astro import ExternalDeploymentSensor
+
+with DAG(
+ dag_id="example_astro_task",
+ start_date=datetime(2022, 1, 1),
+ schedule=None,
+ catchup=False,
+ tags=["example", "async", "core"],
+):
+ ExternalDeploymentSensor(
+ task_id="test1",
+ external_dag_id="example_wait_to_test_example_astro_task",
+ )
+
+ ExternalDeploymentSensor(
+ task_id="test2",
+ external_dag_id="example_wait_to_test_example_astro_task",
+ external_task_id="wait_for_2_min",
+ )
+
+with DAG(
+ dag_id="example_wait_to_test_example_astro_task",
+ start_date=datetime(2022, 1, 1),
+ schedule=None,
+ catchup=False,
+ tags=["example", "async", "core"],
+):
+
+ @task
+ def wait_for_2_min() -> None:
+ """Wait for 2 min."""
+ time.sleep(120)
+
+ wait_for_2_min()
+
+
+with DAG(
+ dag_id="trigger_astro_test_and_example",
+ start_date=datetime(2022, 1, 1),
+ schedule=None,
+ catchup=False,
+ tags=["example", "async", "core"],
+):
+ run_wait_dag = TriggerDagRunOperator(
+ task_id="run_wait_dag",
+ trigger_dag_id="example_wait_to_test_example_astro_task",
+ wait_for_completion=False,
+ )
+
+ run_astro_dag = TriggerDagRunOperator(
+ task_id="run_astro_dag",
+ trigger_dag_id="example_astro_task",
+ wait_for_completion=False,
+ )
+
+ run_wait_dag >> run_astro_dag

astronomer/providers/core/hooks/astro.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import os
+from typing import Any
+from urllib.parse import quote
+
+import requests
+from aiohttp import ClientSession
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+
+
+class AstroHook(BaseHook):
+ """
+ Custom Apache Airflow Hook for interacting with Astro Cloud API.
+
+ :param astro_cloud_conn_id: The connection ID to retrieve Astro Cloud credentials.
+ """
+
+ conn_name_attr = "astro_cloud_conn_id"
+ default_conn_name = "astro_cloud_default"
+ conn_type = "Astro Cloud"
+ hook_name = "Astro Cloud"
+
+ def __init__(self, astro_cloud_conn_id: str = "astro_cloud_conn_id"):
+ super().__init__()
+ self.astro_cloud_conn_id = astro_cloud_conn_id
+
+ @classmethod
+ def get_ui_field_behaviour(cls) -> dict[str, Any]:
+ """
+ Returns UI field behavior customization for the Astro Cloud connection.
+
+ This method defines hidden fields, relabeling, and placeholders for UI display.
+ """
+ return {
+ "hidden_fields": ["login", "port", "schema", "extra"],
+ "relabeling": {
+ "password": "Astro Cloud API Token",
+ },
+ "placeholders": {
+ "host": "https://clmkpsyfc010391acjie00t1l.astronomer.run/d5lc9c9x",
+ "password": "Astro API JWT Token",
+ },
+ }
+
+ def get_conn(self) -> tuple[str, str]:
+ """Retrieves the Astro Cloud connection details."""
+ conn = BaseHook.get_connection(self.astro_cloud_conn_id)
+ base_url = conn.host or os.environ.get("AIRFLOW__WEBSERVER__BASE_URL")
+ if base_url is None:
+ raise AirflowException(f"Airflow host is missing in connection {self.astro_cloud_conn_id}")
+ token = conn.password
+ if token is None:
+ raise AirflowException(f"Astro API token is missing in connection {self.astro_cloud_conn_id}")
+ return base_url, token
+
+ @property
+ def _headers(self) -> dict[str, str]:
+ """Generates and returns headers for Astro Cloud API requests."""
+ _, token = self.get_conn()
+ headers = {"accept": "application/json", "Authorization": f"Bearer {token}"}
+ return headers
+
+ def get_dag_runs(self, external_dag_id: str) -> list[dict[str, str]]:
+ """
+ Retrieves information about running or queued DAG runs.
+
+ :param external_dag_id: External ID of the DAG.
+ """
+ base_url, _ = self.get_conn()
+ path = f"/api/v1/dags/{external_dag_id}/dagRuns"
+ params: dict[str, int | str | list[str]] = {
+ "limit": 1,
+ "state": ["running", "queued"],
+ "order_by": "-execution_date",
+ }
+ url = f"{base_url}{path}"
+ response = requests.get(url, headers=self._headers, params=params)
+ response.raise_for_status()
+ data: dict[str, list[dict[str, str]]] = response.json()
+ return data["dag_runs"]
+
+ def get_dag_run(self, external_dag_id: str, dag_run_id: str) -> dict[str, Any] | None:
+ """
+ Retrieves information about a specific DAG run.
+
+ :param external_dag_id: External ID of the DAG.
+ :param dag_run_id: ID of the DAG run.
+ """
+ base_url, _ = self.get_conn()
+ dag_run_id = quote(dag_run_id)
+ path = f"/api/v1/dags/{external_dag_id}/dagRuns/{dag_run_id}"
+ url = f"{base_url}{path}"
+ response = requests.get(url, headers=self._headers)
+ response.raise_for_status()
+ dr: dict[str, Any] = response.json()
+ return dr
+
+ async def get_a_dag_run(self, external_dag_id: str, dag_run_id: str) -> dict[str, Any] | None:
+ """
+ Retrieves information about a specific DAG run.
+
+ :param external_dag_id: External ID of the DAG.
+ :param dag_run_id: ID of the DAG run.
+ """
+ base_url, _ = self.get_conn()
+ dag_run_id = quote(dag_run_id)
+ path = f"/api/v1/dags/{external_dag_id}/dagRuns/{dag_run_id}"
+ url = f"{base_url}{path}"
+
+ async with ClientSession(headers=self._headers) as session:
+ async with session.get(url) as response:
+ response.raise_for_status()
+ dr: dict[str, Any] = await response.json()
+ return dr
+
+ def get_task_instance(
+ self, external_dag_id: str, dag_run_id: str, external_task_id: str
+ ) -> dict[str, Any] | None:
+ """
+ Retrieves information about a specific task instance within a DAG run.
+
+ :param external_dag_id: External ID of the DAG.
+ :param dag_run_id: ID of the DAG run.
+ :param external_task_id: External ID of the task.
+ """
+ base_url, _ = self.get_conn()
+ dag_run_id = quote(dag_run_id)
+ path = f"/api/v1/dags/{external_dag_id}/dagRuns/{dag_run_id}/taskInstances/{external_task_id}"
+ url = f"{base_url}{path}"
+ response = requests.get(url, headers=self._headers)
+ response.raise_for_status()
+ ti: dict[str, Any] = response.json()
+ return ti
+
+ async def get_a_task_instance(
+ self, external_dag_id: str, dag_run_id: str, external_task_id: str
+ ) -> dict[str, Any] | None:
+ """
+ Retrieves information about a specific task instance within a DAG run.
+
+ :param external_dag_id: External ID of the DAG.
+ :param dag_run_id: ID of the DAG run.
+ :param external_task_id: External ID of the task.
+ """
+ base_url, _ = self.get_conn()
+ dag_run_id = quote(dag_run_id)
+ path = f"/api/v1/dags/{external_dag_id}/dagRuns/{dag_run_id}/taskInstances/{external_task_id}"
+ url = f"{base_url}{path}"
+
+ async with ClientSession(headers=self._headers) as session:
+ async with session.get(url) as response:
+ response.raise_for_status()
+ ti: dict[str, Any] = await response.json()
+ return ti

astronomer/providers/core/sensors/astro.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import datetime
+from typing import Any, cast
+
+from airflow.exceptions import AirflowException, AirflowSkipException
+from airflow.sensors.base import BaseSensorOperator, PokeReturnValue
+
+from astronomer.providers.core.hooks.astro import AstroHook
+from astronomer.providers.core.triggers.astro import AstroDeploymentTrigger
+from astronomer.providers.utils.typing_compat import Context
+
+
+class ExternalDeploymentSensor(BaseSensorOperator):
+ """
+ Custom Apache Airflow sensor for monitoring external deployments using Astro Cloud.
+
+ :param external_dag_id: External ID of the DAG being monitored.
+ :param astro_cloud_conn_id: The connection ID to retrieve Astro Cloud credentials.
+ Defaults to "astro_cloud_default".
+ :param external_task_id: External ID of the task being monitored. If None, monitors the entire DAG.
+ :param kwargs: Additional keyword arguments passed to the BaseSensorOperator constructor.
+ """
+
+ def __init__(
+ self,
+ external_dag_id: str,
+ astro_cloud_conn_id: str = "astro_cloud_default",
+ external_task_id: str | None = None,
+ **kwargs: Any,
+ ):
+ super().__init__(**kwargs)
+ self.astro_cloud_conn_id = astro_cloud_conn_id
+ self.external_task_id = external_task_id
+ self.external_dag_id = external_dag_id
+ self._dag_run_id: str = ""
+
+ def poke(self, context: Context) -> bool | PokeReturnValue:
+ """
+ Check the status of a DAG/task in another deployment.
+
+ Queries Airflow's REST API for the status of the specified DAG or task instance.
+ Returns True if successful, False otherwise.
+
+ :param context: The task execution context.
+ """
+ hook = AstroHook(self.astro_cloud_conn_id)
+ dag_runs: list[dict[str, Any]] = hook.get_dag_runs(self.external_dag_id)
+ if not dag_runs:
+ self.log.info("No DAG runs found for DAG %s", self.external_dag_id)
+ return True
+ self._dag_run_id = cast(str, dag_runs[0]["dag_run_id"])
+ if self.external_task_id is not None:
+ task_instance = hook.get_task_instance(
+ self.external_dag_id, self._dag_run_id, self.external_task_id
+ )
+ task_state = task_instance.get("state") if task_instance else None
+ if task_state == "success":
+ return True
+ else:
+ state = dag_runs[0].get("state")
+ if state == "success":
+ return True
+ return False
+
+ def execute(self, context: Context) -> Any:
+ """
+ Executes the sensor.
+
+ If the external deployment is not successful, it defers the execution using an AstroDeploymentTrigger.
+
+ :param context: The task execution context.
+ """
+ if not self.poke(context):
+ self.defer(
+ timeout=datetime.timedelta(seconds=self.timeout),
+ trigger=AstroDeploymentTrigger(
+ astro_cloud_conn_id=self.astro_cloud_conn_id,
+ external_task_id=self.external_task_id,
+ external_dag_id=self.external_dag_id,
+ poke_interval=self.poke_interval,
+ dag_run_id=self._dag_run_id,
+ ),
+ method_name="execute_complete",
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, str]) -> None:
+ """
+ Handles the completion event from the deferred execution.
+
+ Raises AirflowSkipException if the upstream job failed and `soft_fail` is True.
+ Otherwise, raises AirflowException.
+
+ :param context: The task execution context.
+ :param event: The event dictionary received from the deferred execution.
+ """
+ if event.get("status") == "failed":
+ if self.soft_fail:
+ raise AirflowSkipException("Upstream job failed. Skipping the task.")
+ raise AirflowException("Upstream job failed.")