from importlib import import_module
from logging import getLogger
from typing import TYPE_CHECKING, cast
if TYPE_CHECKING:
from airflow.models.dagrun import DagRun
from .common import (
AirflowPriorityConfigurationOptionNotFound,
BackendSpecificDagContext,
DagContextIdentifier,
DagStatus,
SendMetricFunction,
get_config_option,
has_priority_tag,
)
_log = getLogger(__name__)
[docs]
class Tracker:
[docs]
def __init__(self):
self.dagruns: dict[DagContextIdentifier, dict[DagStatus, BackendSpecificDagContext]] = {}
self.backends: dict[str, SendMetricFunction] = {}
self.thresholds: dict[str, int] = {}
def register(
self,
backend: str,
necessary_configs: list[tuple[str, str]],
):
if backend in self.backends:
_log.warning(f"Backend '{backend}' is already registered. Skipping registration.")
return
_log.info(f"Registering backend '{backend}' with necessary configs: {necessary_configs}")
# Ensure all configs are set
for config in necessary_configs:
try:
get_config_option(*config)
except AirflowPriorityConfigurationOptionNotFound:
_log.warning(
f"Configuration option '{config[0]}.{config[1]}' is not set for backend '{backend}'. Metrics will not be sent for this backend."
)
return
# Import the backend
try:
module = import_module(f"airflow_priority.plugins.{backend}")
except ImportError:
_log.exception(f"Failed to import plugin for backend '{backend}'. Metrics will not be sent for this backend")
return
# register
self.backends[backend] = module.send_metric
# set threshold if available
self.thresholds[backend] = int(get_config_option(backend, "threshold", default=get_config_option("threshold", default=5)))
def _setup_context(self, backend: str, dag_instance_id: int):
key = (backend, dag_instance_id)
if key not in self.dagruns:
self.dagruns[key] = {}
def running(self, dag_run: "DagRun"):
dag_id, priority = has_priority_tag(dag_run=dag_run)
if priority:
_log.info(f"DAG Running: {dag_id} / P{priority}")
for backend, sender in self.backends.items():
_log.info(f"Backend: {backend}")
if priority > self.thresholds.get(backend, 5):
_log.info(f"Skipping running metric for backend {backend} with priority {priority} as it exceeds the threshold.")
continue
# Setup the context if needed
self._setup_context(backend, cast(int, dag_run.id))
try:
sender(dag_id, priority, "running", self.dagruns[(backend, cast(int, dag_run.id))])
except Exception: # noqa: BLE001
_log.exception(f"Failed to send running metric for DAG {dag_id} with priority {priority} on backend {backend}")
def success(self, dag_run: "DagRun"):
dag_id, priority = has_priority_tag(dag_run=dag_run)
if priority:
_log.info(f"DAG Success: {dag_id} / P{priority}")
for backend, sender in self.backends.items():
_log.info(f"Backend: {backend}")
if priority > self.thresholds.get(backend, 5):
_log.info(f"Skipping success metric for backend {backend} with priority {priority} as it exceeds the threshold.")
continue
# Setup the context if needed
self._setup_context(backend, cast(int, dag_run.id))
try:
sender(dag_id, priority, "success", self.dagruns[(backend, cast(int, dag_run.id))])
except Exception: # noqa: BLE001
_log.exception(f"Failed to send success metric for DAG {dag_id} with priority {priority} on backend {backend}")
def failed(self, dag_run: "DagRun"):
dag_id, priority = has_priority_tag(dag_run=dag_run)
if priority:
_log.info(f"DAG Failed: {dag_id} / P{priority}")
for backend, sender in self.backends.items():
_log.info(f"Backend: {backend}")
if priority > self.thresholds.get(backend, 5):
_log.info(f"Skipping failed metric for backend {backend} with priority {priority} as it exceeds the threshold.")
continue
# Setup the context if needed
self._setup_context(backend, cast(int, dag_run.id))
try:
sender(dag_id, priority, "failed", self.dagruns[(backend, cast(int, dag_run.id))])
except Exception: # noqa: BLE001
_log.exception(f"Failed to send failed metric for DAG {dag_id} with priority {priority} on backend {backend}")
# Single instance of Tracker
tracker_inst = Tracker()
tracker_inst.register(
backend="aws",
necessary_configs=[
("aws", "region"),
],
)
tracker_inst.register(
backend="datadog",
necessary_configs=[
("datadog", "api_key"),
],
)
tracker_inst.register(
backend="discord",
necessary_configs=[
("discord", "token"),
("discord", "channel"),
],
)
tracker_inst.register(
backend="logfire",
necessary_configs=[
("logfire", "token"),
],
)
tracker_inst.register(
backend="newrelic",
necessary_configs=[
("newrelic", "api_key"),
],
)
tracker_inst.register(
backend="opsgenie",
necessary_configs=[
("opsgenie", "api_key"),
],
)
tracker_inst.register(
backend="pagerduty",
necessary_configs=[
("pagerduty", "routing_key"),
],
)
tracker_inst.register(
backend="slack",
necessary_configs=[
("slack", "token"),
("slack", "channel"),
],
)
tracker_inst.register(
backend="symphony",
necessary_configs=[
("symphony", "room_name"),
("symphony", "message_create_url"),
("symphony", "cert_file"),
("symphony", "key_file"),
("symphony", "session_auth"),
("symphony", "key_auth"),
("symphony", "room_search_url"),
],
)