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="wait_to_test_example_astro_dag",
+ 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_external_task_async_waits_for_me",
+ 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,112 @@
+from __future__ import annotations
+
+import os
+from typing import Any
+from urllib.parse import quote
+
+import requests
+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": "JWT API 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"]}
+ 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
+
+ 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

astronomer/providers/core/sensors/astro.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import datetime
+
+# import time
+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 wait_for_dag_start(self, second_to_wait: int = 120, sleep: int = 5) -> None:
+ # """TODO"""
+ # hook = AstroHook(self.astro_cloud_conn_id)
+ # end_time = datetime.datetime.now() + datetime.timedelta(seconds=second_to_wait)
+ # while end_time >= datetime.datetime.now():
+ # try:
+ # dag_runs = hook.get_dag_runs(self.external_dag_id)
+ # if dag_runs is not None:
+ # return
+ # except Exception:
+ # time.sleep(sleep)
+
+ 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 dag_runs is None or len(dag_runs) == 0:
+ 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.")
+ else:
+ raise AirflowException("Upstream job failed.")

astronomer/providers/core/triggers/astro.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import asyncio
+from typing import Any, AsyncIterator
+
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+from astronomer.providers.core.hooks.astro import AstroHook
+
+
+class AstroDeploymentTrigger(BaseTrigger):
+ """
+ Custom Apache Airflow trigger for monitoring the completion status of an external deployment using Astro Cloud.
+
+ :param external_dag_id: External ID of the DAG being monitored.
+ :param dag_run_id: ID of the DAG run being monitored.
+ :param external_task_id: External ID of the task being monitored. If None, monitors the entire DAG.
+ :param astro_cloud_conn_id: The connection ID to retrieve Astro Cloud credentials. Defaults to "astro_cloud_default".
+ :param poke_interval: Time in seconds to wait between consecutive checks for completion status.
+ :param kwargs: Additional keyword arguments passed to the BaseTrigger constructor.
+ """
+
+ def __init__(
+ self,
+ external_dag_id: str,
+ dag_run_id: str,
+ external_task_id: str | None = None,
+ astro_cloud_conn_id: str = "astro_cloud_default",
+ poke_interval: float = 5.0,
+ **kwargs: Any,
+ ):
+ super().__init__(**kwargs)
+ self.external_dag_id = external_dag_id
+ self.dag_run_id = dag_run_id
+ self.external_task_id = external_task_id
+ self.astro_cloud_conn_id = astro_cloud_conn_id
+ self.poke_interval = poke_interval
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize the trigger for storage in the database."""
+ return (
+ "astronomer.providers.core.triggers.astro.AstroDeploymentTrigger",
+ {
+ "external_dag_id": self.external_dag_id,
+ "external_task_id": self.external_task_id,
+ "dag_run_id": self.dag_run_id,
+ "astro_cloud_conn_id": self.astro_cloud_conn_id,
+ "poke_interval": self.poke_interval,
+ },
+ )
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ """
+ Asynchronously runs the trigger and yields completion status events.
+
+ Checks the status of the external deployment using Astro Cloud at regular intervals.
+ Yields TriggerEvent with the status "done" if successful, "failed" if failed.
+ """
+ hook = AstroHook(self.astro_cloud_conn_id)
+ while True:
+ 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
+ )
+ state = task_instance.get("state") if task_instance else None
+ if state in ("success", "skipped"):
+ yield TriggerEvent({"status": "done"})
+ elif state in ("failed", "upstream_failed"):
+ yield TriggerEvent({"status": "failed"})
+ else:
+ dag_run = hook.get_dag_run(self.external_dag_id, self.dag_run_id)
+ state = dag_run.get("state") if dag_run else None
+ if state == "success":
+ yield TriggerEvent({"status": "done"})
+ elif state == "failed":
+ yield TriggerEvent({"status": "failed"})
+ self.log.info("Job status is %s sleeping for %s seconds.", state, self.poke_interval)
+ await asyncio.sleep(self.poke_interval)

astronomer/providers/package.py

@@ -10,6 +10,11 @@ def get_provider_info() -> Dict[str, Any]:
  "description": "Apache Airflow Providers containing Deferrable Operators & Sensors from Astronomer",
  "versions": "1.18.4",
  # Optional.
- "hook-class-names": [],
+ "connection-types": [
+ {
+ "hook-class-name": "astronomer.providers.core.hooks.astro.AstroHook",
+ "connection-type": "Astro Cloud",
+ }
+ ],
  "extra-links": [],
  }