airflow-priority 0.1.0__py3-none-any.whl

airflow_priority/__init__.py

@@ -0,0 +1,3 @@
+__version__ = "0.1.0"
+
+from .common import *

airflow_priority/common.py

@@ -0,0 +1,49 @@
+from types import MappingProxyType
+from typing import Literal, Optional, Tuple
+
+from airflow.models.dagrun import DagRun
+
+__all__ = (
+    "has_priority_tag",
+    "get_config_option",
+    "DagStatus",
+    "PriorityTags",
+    "AirflowPriorityConfigurationOptionNotFound",
+)
+
+
+DagStatus = Literal["running", "success", "failed"]
+
+PriorityTags = MappingProxyType(
+    {
+        "P1": 1,
+        "P2": 2,
+        "P3": 3,
+        "P4": 4,
+        "P5": 5,
+    }
+)
+
+
+class AirflowPriorityConfigurationOptionNotFound(RuntimeError): ...
+
+
+def get_config_option(section, key, default, required=True):
+    import airflow.configuration
+
+    config_option = airflow.configuration.conf.get(f"priority.{section}", key, default)
+    if not config_option and required:
+        raise AirflowPriorityConfigurationOptionNotFound(f"{section}.{key}")
+    return config_option
+
+
+def has_priority_tag(dag_run: DagRun) -> Optional[Tuple[str, int]]:
+    dag = dag_run.dag
+    dag_id = dag_run.dag_id
+    tags = dag.tags
+    if tags:
+        for tag in tags:
+            priority = PriorityTags.get(tag.upper())
+            if priority:
+                return dag_id, priority
+    return dag_id, None

airflow_priority/plugins/aws.py

@@ -0,0 +1,65 @@
+import os
+import sys
+from datetime import datetime
+
+from airflow.listeners import hookimpl
+from airflow.models.dagrun import DagRun
+from airflow.plugins_manager import AirflowPlugin
+from boto3 import client as Boto3Client
+from botocore.config import Config
+
+from airflow_priority import DagStatus, has_priority_tag
+
+config = Config(retries=dict(max_attempts=10, mode="adaptive"))
+client = Boto3Client("cloudwatch", config=config)
+
+
+def send_metric_cloudwatch(dag_id: str, priority: int, tag: DagStatus) -> None:
+    client.put_metric_data(
+        Namespace="Airflow/Custom",
+        MetricData=[
+            {
+                "MetricName": f"priority_{tag}",
+                "Dimensions": [
+                    {"Name": "environment", "Value": os.environ.get("AIRFLOW_ENV_NAME", "unknown-airflow-env")},
+                    {
+                        "Name": "dag",
+                        "Value": dag_id,
+                    },
+                    {
+                        "Name": "priority",
+                        "Value": str(priority),
+                    },
+                ],
+                "Timestamp": datetime.utcnow(),
+                "Value": int(priority),
+                "Unit": "Count",
+            }
+        ],
+    )
+
+
+@hookimpl
+def on_dag_run_running(dag_run: DagRun, msg: str):
+    dag_id, priority = has_priority_tag(dag_run=dag_run)
+    if priority:
+        send_metric_cloudwatch(dag_id, priority, "running")
+
+
+@hookimpl
+def on_dag_run_success(dag_run: DagRun, msg: str):
+    dag_id, priority = has_priority_tag(dag_run=dag_run)
+    if priority:
+        send_metric_cloudwatch(dag_id, priority, "success")
+
+
+@hookimpl
+def on_dag_run_failed(dag_run: DagRun, msg: str):
+    dag_id, priority = has_priority_tag(dag_run=dag_run)
+    if priority:
+        send_metric_cloudwatch(dag_id, priority, "failed")
+
+
+class AWSCloudWatchPriorityPlugin(AirflowPlugin):
+    name = "AWSCloudWatchPriorityPlugin"
+    listeners = [sys.modules[__name__]]

airflow_priority/plugins/datadog.py

@@ -0,0 +1,79 @@
+import sys
+from datetime import datetime
+from functools import lru_cache
+
+from airflow.listeners import hookimpl
+from airflow.models.dagrun import DagRun
+from airflow.plugins_manager import AirflowPlugin
+from datadog_api_client import ApiClient, Configuration
+from datadog_api_client.v2.api.metrics_api import MetricsApi
+from datadog_api_client.v2.model.metric_intake_type import MetricIntakeType
+from datadog_api_client.v2.model.metric_payload import MetricPayload
+from datadog_api_client.v2.model.metric_point import MetricPoint
+from datadog_api_client.v2.model.metric_resource import MetricResource
+from datadog_api_client.v2.model.metric_series import MetricSeries
+
+from airflow_priority import DagStatus, get_config_option, has_priority_tag
+
+
+@lru_cache
+def get_configuration():
+    return Configuration(
+        api_key={"apiKeyAuth": get_config_option("datadog", "api_key", "")},
+    )
+
+
+def send_metric_datadog(dag_id: str, priority: int, tag: DagStatus) -> None:
+    with ApiClient(get_configuration()) as api_client:
+        api_instance = MetricsApi(api_client)
+
+        body = MetricPayload(
+            series=[
+                MetricSeries(
+                    metric=f"airflow.custom.priority.p{priority}.{tag}",
+                    type=MetricIntakeType.UNSPECIFIED,
+                    points=[
+                        MetricPoint(
+                            timestamp=int(datetime.now().timestamp()),
+                            value=priority,
+                        ),
+                    ],
+                    resources=[
+                        MetricResource(
+                            name=dag_id,
+                            type="dag",
+                        ),
+                    ],
+                    tags=["airflow", "priority", f"dag_{dag_id}"],
+                ),
+            ],
+        )
+
+        resp = api_instance.submit_metrics(body=body)
+        assert resp["errors"] == []
+
+
+@hookimpl
+def on_dag_run_running(dag_run: DagRun, msg: str):
+    dag_id, priority = has_priority_tag(dag_run=dag_run)
+    if priority:
+        send_metric_datadog(dag_id, priority, "running")
+
+
+@hookimpl
+def on_dag_run_success(dag_run: DagRun, msg: str):
+    dag_id, priority = has_priority_tag(dag_run=dag_run)
+    if priority:
+        send_metric_datadog(dag_id, priority, "success")
+
+
+@hookimpl
+def on_dag_run_failed(dag_run: DagRun, msg: str):
+    dag_id, priority = has_priority_tag(dag_run=dag_run)
+    if priority:
+        send_metric_datadog(dag_id, priority, "failed")
+
+
+class DatadogPriorityPlugin(AirflowPlugin):
+    name = "DatadogPriorityPlugin"
+    listeners = [sys.modules[__name__]]

airflow_priority/plugins/newrelic.py

@@ -0,0 +1,55 @@
+import sys
+from functools import lru_cache
+
+from airflow.listeners import hookimpl
+from airflow.models.dagrun import DagRun
+from airflow.plugins_manager import AirflowPlugin
+from newrelic_telemetry_sdk import GaugeMetric, MetricClient
+
+from airflow_priority import DagStatus, get_config_option, has_priority_tag
+
+__all__ = (
+    "send_metric_newrelic",
+    "on_dag_run_running",
+    "on_dag_run_success",
+    "on_dag_run_failed",
+    "NewRelicPriorityPlugin",
+)
+
+
+@lru_cache
+def get_client():
+    return MetricClient(get_config_option("newrelic", "api_key", "", True))
+
+
+def send_metric_newrelic(dag_id: str, priority: int, tag: DagStatus) -> None:
+    priority = GaugeMetric(
+        f"airflow.custom.priority.p{priority}.{tag}", 1, tags={"application": "airflow", "dag": dag_id, "priority": priority, "tag": tag}
+    )
+    get_client().send_batch([priority])
+
+
+@hookimpl
+def on_dag_run_running(dag_run: DagRun, msg: str):
+    dag_id, priority = has_priority_tag(dag_run=dag_run)
+    if priority:
+        send_metric_newrelic(dag_id, priority, "running")
+
+
+@hookimpl
+def on_dag_run_success(dag_run: DagRun, msg: str):
+    dag_id, priority = has_priority_tag(dag_run=dag_run)
+    if priority:
+        send_metric_newrelic(dag_id, priority, "success")
+
+
+@hookimpl
+def on_dag_run_failed(dag_run: DagRun, msg: str):
+    dag_id, priority = has_priority_tag(dag_run=dag_run)
+    if priority:
+        send_metric_newrelic(dag_id, priority, "failed")
+
+
+class NewRelicPriorityPlugin(AirflowPlugin):
+    name = "NewRelicPriorityPlugin"
+    listeners = [sys.modules[__name__]]

airflow_priority/plugins/slack.py

@@ -0,0 +1,57 @@
+import sys
+from functools import lru_cache
+
+from airflow.listeners import hookimpl
+from airflow.models.dagrun import DagRun
+from airflow.plugins_manager import AirflowPlugin
+from slack_sdk import WebClient
+
+from airflow_priority import DagStatus, get_config_option, has_priority_tag
+
+__all__ = ("get_client", "get_channel_id", "send_metric_slack", "on_dag_run_failed", "SlackPriorityPlugin")
+
+
+@lru_cache
+def get_client():
+    return WebClient(token=get_config_option("slack", "token", ""))
+
+
+@lru_cache
+def get_channel_id():
+    channel_name = get_config_option("slack", "channel", "")
+    conversations = get_client().conversations_list(types=["public_channel", "private_channel"])
+    if conversations.data["ok"]:
+        for channel in conversations.data["channels"]:
+            if channel["name"] == channel_name:
+                return channel["id"]
+    raise Exception("TODO")
+
+
+def send_metric_slack(dag_id: str, priority: int, tag: DagStatus) -> None:
+    get_client().chat_postMessage(channel=get_channel_id(), text=f'A P{priority} DAG "{dag_id}" has {tag}!')
+
+
+# @hookimpl
+# def on_dag_run_running(dag_run: DagRun, msg: str):
+#     dag_id, priority = has_priority_tag(dag_run=dag_run)
+#     if priority:
+#         send_metric_slack(dag_id, priority, "running")
+
+
+# @hookimpl
+# def on_dag_run_success(dag_run: DagRun, msg: str):
+#     dag_id, priority = has_priority_tag(dag_run=dag_run)
+#     if priority:
+#         send_metric_slack(dag_id, priority, "succeeded")
+
+
+@hookimpl
+def on_dag_run_failed(dag_run: DagRun, msg: str):
+    dag_id, priority = has_priority_tag(dag_run=dag_run)
+    if priority:
+        send_metric_slack(dag_id, priority, "failed")
+
+
+class SlackPriorityPlugin(AirflowPlugin):
+    name = "SlackPriorityPlugin"
+    listeners = [sys.modules[__name__]]

airflow_priority/tests/airflow.cfg.jinja

@@ -0,0 +1,363 @@
+{% raw %}
+[core]
+dags_folder = ~/airflow/dags
+hostname_callable = airflow.utils.net.getfqdn
+might_contain_dag_callable = airflow.utils.file.might_contain_dag_via_default_heuristic
+default_timezone = utc
+executor = SequentialExecutor
+auth_manager = airflow.providers.fab.auth_manager.fab_auth_manager.FabAuthManager
+parallelism = 32
+max_active_tasks_per_dag = 16
+dags_are_paused_at_creation = True
+max_active_runs_per_dag = 16
+max_consecutive_failed_dag_runs_per_dag = 0
+load_examples = True
+plugins_folder = ~/airflow/plugins
+execute_tasks_new_python_interpreter = False
+fernet_key = 
+donot_pickle = True
+dagbag_import_timeout = 30.0
+dagbag_import_error_tracebacks = True
+dagbag_import_error_traceback_depth = 2
+dag_file_processor_timeout = 50
+task_runner = StandardTaskRunner
+default_impersonation = 
+security = 
+unit_test_mode = False
+enable_xcom_pickling = False
+allowed_deserialization_classes = airflow.*
+allowed_deserialization_classes_regexp = 
+killed_task_cleanup_time = 60
+dag_run_conf_overrides_params = True
+dag_discovery_safe_mode = True
+dag_ignore_file_syntax = regexp
+default_task_retries = 0
+default_task_retry_delay = 300
+max_task_retry_delay = 86400
+default_task_weight_rule = downstream
+default_task_execution_timeout = 
+min_serialized_dag_update_interval = 30
+compress_serialized_dags = False
+min_serialized_dag_fetch_interval = 10
+max_num_rendered_ti_fields_per_task = 30
+check_slas = True
+xcom_backend = airflow.models.xcom.BaseXCom
+lazy_load_plugins = True
+lazy_discover_providers = True
+hide_sensitive_var_conn_fields = True
+sensitive_var_conn_names = 
+default_pool_task_slot_count = 128
+max_map_length = 1024
+daemon_umask = 0o077
+strict_dataset_uri_validation = False
+database_access_isolation = False
+test_connection = Disabled
+max_templated_field_length = 4096
+
+[database]
+alembic_ini_file_path = alembic.ini
+sql_alchemy_conn = sqlite:////~/airflow/airflow.db
+sql_engine_encoding = utf-8
+sql_alchemy_pool_enabled = True
+sql_alchemy_pool_size = 5
+sql_alchemy_max_overflow = 10
+sql_alchemy_pool_recycle = 1800
+sql_alchemy_pool_pre_ping = True
+sql_alchemy_schema = 
+load_default_connections = True
+max_db_retries = 3
+check_migrations = True
+
+[logging]
+base_log_folder = ~/airflow/logs
+remote_logging = False
+remote_log_conn_id = 
+delete_local_logs = False
+google_key_path = 
+remote_base_log_folder = 
+remote_task_handler_kwargs = 
+encrypt_s3_logs = False
+logging_level = INFO
+celery_logging_level = 
+fab_logging_level = WARNING
+logging_config_class = 
+colored_console_log = True
+colored_log_format = [%%(blue)s%%(asctime)s%%(reset)s] {%%(blue)s%%(filename)s:%%(reset)s%%(lineno)d} %%(log_color)s%%(levelname)s%%(reset)s - %%(log_color)s%%(message)s%%(reset)s
+colored_formatter_class = airflow.utils.log.colored_log.CustomTTYColoredFormatter
+log_format = [%%(asctime)s] {%%(filename)s:%%(lineno)d} %%(levelname)s - %%(message)s
+simple_log_format = %%(asctime)s %%(levelname)s - %%(message)s
+dag_processor_log_target = file
+dag_processor_log_format = [%%(asctime)s] [SOURCE:DAG_PROCESSOR] {%%(filename)s:%%(lineno)d} %%(levelname)s - %%(message)s
+log_formatter_class = airflow.utils.log.timezone_aware.TimezoneAware
+secret_mask_adapter = 
+task_log_prefix_template = 
+log_filename_template = dag_id={{ ti.dag_id }}/run_id={{ ti.run_id }}/task_id={{ ti.task_id }}/{%% if ti.map_index >= 0 %%}map_index={{ ti.map_index }}/{%% endif %%}attempt={{ try_number }}.log
+log_processor_filename_template = {{ filename }}.log
+dag_processor_manager_log_location = ~/airflow/logs/dag_processor_manager/dag_processor_manager.log
+dag_processor_manager_log_stdout = False
+task_log_reader = task
+extra_logger_names = 
+worker_log_server_port = 8793
+trigger_log_server_port = 8794
+file_task_handler_new_folder_permissions = 0o775
+file_task_handler_new_file_permissions = 0o664
+celery_stdout_stderr_separation = False
+enable_task_context_logger = True
+
+[metrics]
+metrics_use_pattern_match = False
+metrics_allow_list = 
+metrics_block_list = 
+statsd_on = False
+statsd_host = localhost
+statsd_port = 8125
+statsd_prefix = airflow
+stat_name_handler = 
+statsd_datadog_enabled = False
+statsd_datadog_tags = 
+statsd_datadog_metrics_tags = True
+statsd_disabled_tags = job_id,run_id
+statsd_influxdb_enabled = False
+otel_on = False
+otel_host = localhost
+otel_port = 8889
+otel_prefix = airflow
+otel_interval_milliseconds = 60000
+otel_debugging_on = False
+otel_ssl_active = False
+
+[secrets]
+backend = 
+backend_kwargs = 
+use_cache = False
+cache_ttl_seconds = 900
+
+[cli]
+api_client = airflow.api.client.local_client
+endpoint_url = http://localhost:8080
+
+[debug]
+fail_fast = False
+
+[api]
+enable_experimental_api = False
+auth_backends = airflow.api.auth.backend.session
+maximum_page_limit = 100
+fallback_page_limit = 100
+google_oauth2_audience = 
+google_key_path = 
+access_control_allow_headers = 
+access_control_allow_methods = 
+access_control_allow_origins = 
+enable_xcom_deserialize_support = False
+
+[lineage]
+backend = 
+
+[operators]
+default_owner = airflow
+default_deferrable = false
+default_cpus = 1
+default_ram = 512
+default_disk = 512
+default_gpus = 0
+default_queue = default
+allow_illegal_arguments = False
+
+[webserver]
+access_denied_message = Access is Denied
+config_file = ~/airflow/webserver_config.py
+base_url = http://localhost:8080
+default_ui_timezone = UTC
+web_server_host = 0.0.0.0
+web_server_port = 8080
+web_server_ssl_cert = 
+web_server_ssl_key = 
+session_backend = database
+web_server_master_timeout = 120
+web_server_worker_timeout = 120
+worker_refresh_batch_size = 1
+worker_refresh_interval = 6000
+reload_on_plugin_change = False
+secret_key = ANcdrL+FG0dj9+7o4K6LGA==
+workers = 4
+worker_class = sync
+access_logfile = -
+error_logfile = -
+access_logformat = 
+expose_config = False
+expose_hostname = False
+expose_stacktrace = False
+dag_default_view = grid
+dag_orientation = LR
+grid_view_sorting_order = topological
+log_fetch_timeout_sec = 5
+log_fetch_delay_sec = 2
+log_auto_tailing_offset = 30
+log_animation_speed = 1000
+hide_paused_dags_by_default = False
+page_size = 100
+navbar_color = #fff
+navbar_text_color = #51504f
+navbar_hover_color = #eee
+navbar_text_hover_color = #51504f
+navbar_logo_text_color = #51504f
+default_dag_run_display_number = 25
+enable_proxy_fix = False
+proxy_fix_x_for = 1
+proxy_fix_x_proto = 1
+proxy_fix_x_host = 1
+proxy_fix_x_port = 1
+proxy_fix_x_prefix = 1
+cookie_secure = False
+cookie_samesite = Lax
+default_wrap = False
+x_frame_enabled = True
+show_recent_stats_for_completed_runs = True
+session_lifetime_minutes = 43200
+instance_name_has_markup = False
+auto_refresh_interval = 3
+warn_deployment_exposure = True
+enable_swagger_ui = True
+run_internal_api = False
+caching_hash_method = md5
+show_trigger_form_if_no_params = False
+num_recent_configurations_for_trigger = 5
+allow_raw_html_descriptions = False
+allowed_payload_size = 1.0
+require_confirmation_dag_change = False
+
+[email]
+email_backend = airflow.utils.email.send_email_smtp
+email_conn_id = smtp_default
+default_email_on_retry = True
+default_email_on_failure = True
+ssl_context = default
+
+[smtp]
+smtp_host = localhost
+smtp_starttls = True
+smtp_ssl = False
+smtp_port = 25
+smtp_mail_from = airflow@example.com
+smtp_timeout = 30
+smtp_retry_limit = 5
+
+[sentry]
+sentry_on = false
+sentry_dsn = 
+
+[scheduler]
+job_heartbeat_sec = 5
+scheduler_heartbeat_sec = 5
+local_task_job_heartbeat_sec = 0
+num_runs = -1
+scheduler_idle_sleep_time = 1
+min_file_process_interval = 30
+parsing_cleanup_interval = 60
+stale_dag_threshold = 50
+dag_dir_list_interval = 300
+print_stats_interval = 30
+pool_metrics_interval = 5.0
+scheduler_health_check_threshold = 30
+enable_health_check = False
+scheduler_health_check_server_host = 0.0.0.0
+scheduler_health_check_server_port = 8974
+orphaned_tasks_check_interval = 300.0
+child_process_log_directory = ~/airflow/logs/scheduler
+scheduler_zombie_task_threshold = 300
+zombie_detection_interval = 10.0
+catchup_by_default = True
+ignore_first_depends_on_past_by_default = True
+max_tis_per_query = 16
+use_row_level_locking = True
+max_dagruns_to_create_per_loop = 10
+max_dagruns_per_loop_to_schedule = 20
+schedule_after_task_execution = True
+parsing_pre_import_modules = True
+parsing_processes = 2
+file_parsing_sort_mode = modified_time
+standalone_dag_processor = False
+max_callbacks_per_loop = 20
+dag_stale_not_seen_duration = 600
+use_job_schedule = True
+allow_trigger_in_future = False
+trigger_timeout_check_interval = 15
+task_queued_timeout = 600.0
+task_queued_timeout_check_interval = 120.0
+allowed_run_id_pattern = ^[A-Za-z0-9_.~:+-]+$
+create_cron_data_intervals = True
+
+[triggerer]
+default_capacity = 1000
+job_heartbeat_sec = 5
+triggerer_health_check_threshold = 30
+
+[kerberos]
+ccache = /tmp/airflow_krb5_ccache
+principal = airflow
+reinit_frequency = 3600
+kinit_path = kinit
+keytab = airflow.keytab
+forwardable = True
+include_ip = True
+
+[sensors]
+default_timeout = 604800
+
+[celery_kubernetes_executor]
+kubernetes_queue = kubernetes
+
+[celery]
+celery_app_name = airflow.providers.celery.executors.celery_executor
+worker_concurrency = 16
+worker_prefetch_multiplier = 1
+worker_enable_remote_control = true
+broker_url = redis://redis:6379/0
+result_backend_sqlalchemy_engine_options = 
+flower_host = 0.0.0.0
+flower_url_prefix = 
+flower_port = 5555
+flower_basic_auth = 
+sync_parallelism = 0
+celery_config_options = airflow.providers.celery.executors.default_celery.DEFAULT_CELERY_CONFIG
+ssl_active = False
+ssl_key = 
+ssl_cert = 
+ssl_cacert = 
+pool = prefork
+operation_timeout = 1.0
+task_acks_late = True
+task_track_started = True
+task_publish_max_retries = 3
+worker_precheck = False
+
+[celery_broker_transport_options]
+
+[common.io]
+xcom_objectstorage_path = 
+xcom_objectstorage_threshold = -1
+xcom_objectstorage_compression = 
+
+[fab]
+auth_rate_limited = True
+auth_rate_limit = 5 per 40 second
+update_fab_perms = True
+
+[imap]
+[smtp_provider]
+[atlas]
+test = "blerg"
+[airflow.priority.slack]
+test="blerg"
+{% endraw %}
+
+[priority.slack]
+token = {{ SLACK_TOKEN }}
+channel = {{ SLACK_CHANNEL }}
+
+[priority.datadog]
+api_key = {{ DATADOG_API_KEY }}
+
+[priority.newrelic]
+api_key = {{ NEWRELIC_API_KEY }}

airflow_priority/tests/conftest.py

@@ -0,0 +1,44 @@
+import os
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import pytest
+from jinja2 import DictLoader, Environment
+
+
+@pytest.fixture(autouse=True)
+def airflow_config():
+    config_template = (Path(__file__).parent / "airflow.cfg.jinja").read_text()
+    j2 = Environment(loader=DictLoader({"airflow.cfg": config_template}), trim_blocks=True)
+    with TemporaryDirectory() as td:
+        tmpl = j2.get_template("airflow.cfg").render(
+            SLACK_TOKEN=os.environ.get("SLACK_TOKEN", ""),
+            SLACK_CHANNEL=os.environ.get("SLACK_CHANNEL", ""),
+            DATADOG_HOST=os.environ.get("DATADOG_HOST", ""),
+            DATADOG_API_KEY=os.environ.get("DATADOG_API_KEY", ""),
+            NEWRELIC_API_KEY=os.environ.get("NEWRELIC_API_KEY", ""),
+        )
+        (Path(td) / "airflow.cfg").write_text(tmpl)
+        os.environ["AIRFLOW_HOME"] = str(Path(td))
+        os.environ["AIRFLOW_CONFIG"] = str((Path(td) / "airflow.cfg"))
+        import airflow.configuration
+
+        airflow.configuration.AIRFLOW_HOME = os.environ["AIRFLOW_HOME"]
+        airflow.configuration.AIRFLOW_CONFIG = os.environ["AIRFLOW_CONFIG"]
+        airflow_config_parser = airflow.configuration.AirflowConfigParser()
+        airflow.configuration.load_standard_airflow_configuration(airflow_config_parser)
+        airflow_config_parser.validate()
+        airflow.configuration.conf = airflow_config_parser
+        yield
+
+
+@pytest.fixture(scope="function", autouse=True)
+def dag_run():
+    from unittest.mock import MagicMock
+
+    from airflow.models.dagrun import DagRun
+
+    dag_run = DagRun("UNIT TEST")
+    dag_run.dag = MagicMock()
+    dag_run.dag.tags = ["P1"]
+    return dag_run

airflow_priority/tests/test_common.py

@@ -0,0 +1,4 @@
+def test_has_priority_tag(dag_run):
+    from airflow_priority import has_priority_tag
+
+    assert has_priority_tag(dag_run) == ("UNIT TEST", 1)

airflow_priority/tests/test_datadog.py

@@ -0,0 +1,31 @@
+from unittest.mock import patch
+
+
+def test_datadog_send(airflow_config, dag_run):
+    from airflow_priority.plugins.datadog import send_metric_datadog
+
+    send_metric_datadog("UNIT TEST", 1, "testing")
+
+
+def test_datadog_priority_failed(airflow_config, dag_run):
+    from airflow_priority.plugins.datadog import on_dag_run_failed
+
+    with patch("airflow_priority.plugins.datadog.send_metric_datadog") as p1:
+        on_dag_run_failed(dag_run, "test")
+    assert p1.call_count == 1
+
+
+def test_datadog_priority_running(airflow_config, dag_run):
+    from airflow_priority.plugins.datadog import on_dag_run_running
+
+    with patch("airflow_priority.plugins.datadog.send_metric_datadog") as p1:
+        on_dag_run_running(dag_run, "test")
+    assert p1.call_count == 1
+
+
+def test_datadog_priority_success(airflow_config, dag_run):
+    from airflow_priority.plugins.datadog import on_dag_run_success
+
+    with patch("airflow_priority.plugins.datadog.send_metric_datadog") as p1:
+        on_dag_run_success(dag_run, "test")
+    assert p1.call_count == 1

airflow_priority/tests/test_newrelic.py

@@ -0,0 +1,31 @@
+from unittest.mock import patch
+
+
+def test_newrelic_send(airflow_config, dag_run):
+    from airflow_priority.plugins.newrelic import send_metric_newrelic
+
+    send_metric_newrelic("UNIT TEST", 1, "testing")
+
+
+def test_newrelic_priority_failed(airflow_config, dag_run):
+    from airflow_priority.plugins.newrelic import on_dag_run_failed
+
+    with patch("airflow_priority.plugins.newrelic.send_metric_newrelic") as p1:
+        on_dag_run_failed(dag_run, "test")
+    assert p1.call_count == 1
+
+
+def test_newrelic_priority_running(airflow_config, dag_run):
+    from airflow_priority.plugins.newrelic import on_dag_run_running
+
+    with patch("airflow_priority.plugins.newrelic.send_metric_newrelic") as p1:
+        on_dag_run_running(dag_run, "test")
+    assert p1.call_count == 1
+
+
+def test_newrelic_priority_success(airflow_config, dag_run):
+    from airflow_priority.plugins.newrelic import on_dag_run_success
+
+    with patch("airflow_priority.plugins.newrelic.send_metric_newrelic") as p1:
+        on_dag_run_success(dag_run, "test")
+    assert p1.call_count == 1

airflow_priority/tests/test_slack.py

@@ -0,0 +1,15 @@
+from unittest.mock import patch
+
+
+def test_slack_send(airflow_config, dag_run):
+    from airflow_priority.plugins.slack import send_metric_slack
+
+    send_metric_slack("UNIT TEST", 1, "BEEN TESTED")
+
+
+def test_slack_priority_failed(airflow_config, dag_run):
+    from airflow_priority.plugins.slack import on_dag_run_failed
+
+    with patch("airflow_priority.plugins.slack.send_metric_slack") as p1:
+        on_dag_run_failed(dag_run, "test")
+    assert p1.call_count == 1

airflow_priority-0.1.0.dist-info/METADATA

@@ -0,0 +1,154 @@
+Metadata-Version: 2.3
+Name: airflow-priority
+Version: 0.1.0
+Summary: Priority Tags for Airflow Dags
+Project-URL: Repository, https://github.com/airflow-laminar/airflow-priority
+Project-URL: Homepage, https://github.com/airflow-laminar/airflow-priority
+Author-email: the airflow-priority authors <3105306+timkpaine@users.noreply.github.com>
+License: Apache-2.0
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Programming Language :: Python :: Implementation :: PyPy
+Requires-Python: >=3.9
+Requires-Dist: apache-airflow<3,>=2
+Provides-Extra: aws
+Requires-Dist: boto3; extra == 'aws'
+Requires-Dist: botocore; extra == 'aws'
+Provides-Extra: datadog
+Requires-Dist: datadog-api-client; extra == 'datadog'
+Provides-Extra: develop
+Requires-Dist: boto3; extra == 'develop'
+Requires-Dist: botocore; extra == 'develop'
+Requires-Dist: build; extra == 'develop'
+Requires-Dist: bump-my-version; extra == 'develop'
+Requires-Dist: check-manifest; extra == 'develop'
+Requires-Dist: datadog-api-client; extra == 'develop'
+Requires-Dist: hatchling; extra == 'develop'
+Requires-Dist: newrelic-telemetry-sdk; extra == 'develop'
+Requires-Dist: pytest; extra == 'develop'
+Requires-Dist: pytest-cov; extra == 'develop'
+Requires-Dist: ruff; extra == 'develop'
+Requires-Dist: slack-sdk; extra == 'develop'
+Requires-Dist: twine; extra == 'develop'
+Requires-Dist: wheel; extra == 'develop'
+Provides-Extra: newrelic
+Requires-Dist: newrelic-telemetry-sdk; extra == 'newrelic'
+Provides-Extra: slack
+Requires-Dist: slack-sdk; extra == 'slack'
+Description-Content-Type: text/markdown
+
+# airflow-priority
+
+Priority Tags for Airflow Dags
+
+[![Build Status](https://github.com/airflow-laminar/airflow-priority/actions/workflows/build.yml/badge.svg?branch=main&event=push)](https://github.com/airflow-laminar/airflow-priority/actions/workflows/build.yml)
+[![codecov](https://codecov.io/gh/airflow-laminar/airflow-priority/branch/main/graph/badge.svg)](https://codecov.io/gh/airflow-laminar/airflow-priority)
+[![License](https://img.shields.io/github/license/airflow-laminar/airflow-priority)](https://github.com/airflow-laminar/airflow-priority)
+[![PyPI](https://img.shields.io/pypi/v/airflow-priority.svg)](https://pypi.python.org/pypi/airflow-priority)
+
+## Overview
+
+This repo provides [Airflow Plugins](https://airflow.apache.org/docs/apache-airflow/stable/authoring-and-scheduling/plugins.html) for priority-driven DAG failure alerting. In layman's terms, one need only add a [tag](https://airflow.apache.org/docs/apache-airflow/stable/howto/add-dag-tags.html) to their DAG in `P1, P2, P3, P4, P5`, and that dag will send a notification to:
+
+- [New Relic](https://newrelic.com)
+- [Datadog](https://www.datadoghq.com)
+- [Slack](http://slack.com)
+
+Where `P1` corresponds to highest priority, and `P5` corresponds to lowest.
+
+## Installation
+
+You can install from pip:
+
+```bash
+pip install airflow-priority
+```
+
+Or via conda:
+
+```bash
+conda install airflow-priority -c conda-forge
+```
+
+## Integrations
+| Integration | Metric / Tag |
+| :---------- | :----------- |
+| [New Relic](https://newrelic.com) | `airflow.custom.priority.p{1,2,3,4,5}.{failed,succeeded,running}` |
+| [Datadog](https://www.datadoghq.com) | `airflow.custom.priority.p{1,2,3,4,5}.{failed,succeeded,running}` |
+| [Slack](http://slack.com) | `N/A` | 
+
+
+### New Relic
+
+Create a new New Relic API Key [following their guide](https://docs.newrelic.com/docs/apis/intro-apis/new-relic-api-keys/). Note that the type should have `INGEST - LICENSE`.
+
+Copy this api key into your `airflow.cfg` like so:
+
+```
+[priority.newrelic]
+api_key = the api key
+```
+
+Under `Query Your Data` in the New Relic UI, you can create a query for the new custom metric:
+
+```
+SELECT sum(`airflow.custom.priority.p1.failed`) FROM Metric FACET dag
+```
+
+With this, you can now [create a custom alert](https://docs.newrelic.com/docs/alerts/create-alert/examples/define-custom-metrics-alert-condition/). For fast alerting, we recommend the following parameters:
+
+```raw
+Window duration - 30 seconds
+Sliding window aggregation - Disabled
+Slide by interval - Not set
+Streaming method - Event timer
+Timer - 5 seconds
+
+Fill data gaps with - None
+Evaluation delay - Not set
+
+Thresholds: Critical: Query result is above or equals 1 at least once in 1 minute
+```
+
+### Slack
+
+Configure a new slack application following the [Slack Quickstart](https://api.slack.com/quickstart).
+
+Ensure your application has the following scopes for public and private channel access:
+
+- `channels:read`
+- `groups:read`
+- `chat:write`
+
+Enable and install your Slack application into your workspace, and add it as an integration in whatever channel you want it to post.
+
+Copy your Slack application's Oauth Token (starting with `xoxb-`) and your desired channel into your `airflow.cfg` like so:
+
+```
+[priority.slack]
+token = xoxb-...
+channel = channel-name
+```
+
+### Datadog
+
+Create a new Datadog api key [following their guide](https://docs.datadoghq.com/account_management/api-app-keys/#add-an-api-key-or-client-token).
+
+Copy this api key into your `airflow.cfg` like so:
+
+```
+[priority.datadog]
+api_key = the api key
+```
+
+Ensure your dags are configured with tags and run some, it can often be convenient to have an intentionally failing `P1` dag to test the integration. With this, you can now [create custom monitors](https://docs.datadoghq.com/getting_started/monitors/) for the tags.
+
+
+## License

airflow_priority-0.1.0.dist-info/RECORD

@@ -0,0 +1,18 @@
+airflow_priority/__init__.py,sha256=JRAlq9ECAesJA0mEaoYsdF-a81-p96Ne327m2Bz5cPo,45
+airflow_priority/common.py,sha256=kOUXLY6wBp2iyKafjiYDWqcq5OxpkJu-t4eaEi4y5ro,1175
+airflow_priority/plugins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+airflow_priority/plugins/aws.py,sha256=waws2w3Zr7Sd-KPdHh0euqDcQ6ALZbhSkg1zhO4m-38,1960
+airflow_priority/plugins/datadog.py,sha256=3Twjn-cqsUcfGIxcO6Tcrg2-dRyI8hxTirMiC27B4O8,2620
+airflow_priority/plugins/newrelic.py,sha256=u7K0jPsTXeGzepsvxJoyhy8wRoRkQ-xvw0OllS6oTG4,1567
+airflow_priority/plugins/slack.py,sha256=dqMHKheBWn63tgvsOYtdlIpARnzzKcQPCGm4h88t8K4,1784
+airflow_priority/tests/airflow.cfg.jinja,sha256=0NFCaKzB4y-SfdcU01Okzsn5yu9Wolc_rtiEFwmsMQk,9985
+airflow_priority/tests/conftest.py,sha256=mUl_LWdzbI4Tfvjg66aqZCBsv9BoQbygGGeNC1mLcLE,1683
+airflow_priority/tests/test_common.py,sha256=1ctwps0n4VmZJFm6eqiYcHyyd-jnf3sXaTnur_zwAz4,144
+airflow_priority/tests/test_datadog.py,sha256=ly-NbB3yAmS5KHN53e5odInzaHgGdqOpocZ-md4CV_8,1049
+airflow_priority/tests/test_newrelic.py,sha256=heEc_MV_cYT0dumg3ChYDa5MPNZuOGUGUoxre1AD8oo,1065
+airflow_priority/tests/test_slack.py,sha256=tZvYVMQ-aYOuZbYyPa4A9DkaYJvgccsZaHHVb8NUe9s,471
+airflow_priority-0.1.0.dist-info/METADATA,sha256=cC3_1Od7GcspN2HdzPx4K__SDvfQS7mOmpoonFKM-5k,5775
+airflow_priority-0.1.0.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+airflow_priority-0.1.0.dist-info/entry_points.txt,sha256=rANAxUPYmf0fI8I2YZSsDyeEZeAfknp6wz7P75CV3uI,352
+airflow_priority-0.1.0.dist-info/licenses/LICENSE,sha256=SmmTABBMCdZmZ8-nO_YaPUSw6Hvd8SD2kptj4jEO7SE,11358
+airflow_priority-0.1.0.dist-info/RECORD,,

airflow_priority-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.25.0
+Root-Is-Purelib: true
+Tag: py3-none-any

airflow_priority-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,5 @@
+[airflow.plugins]
+airflow_priority_aws_cloudwatch = airflow_priority.plugins.aws:AWSCloudWatchPriorityPlugin
+airflow_priority_datadog = airflow_priority.plugins.datadog:DatadogPriorityPlugin
+airflow_priority_newrelic = airflow_priority.plugins.newrelic:NewRelicPriorityPlugin
+airflow_priority_slack = airflow_priority.plugins.slack:SlackPriorityPlugin

airflow_priority-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2024 the airflow-priority authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.