MM-56793: new push proxy dag (#1453)

* MM-56793: new push proxy dag * Fix failing test * Fix linter
mattermost · Feb 29, 2024 · e58ca90 · e58ca90
1 parent 59919e9
commit e58ca90
Showing 1 changed file with 83 additions and 0 deletions.
diff --git a/airflow/dags/mattermost_dags/extract/push_proxy_new.py b/airflow/dags/mattermost_dags/extract/push_proxy_new.py
@@ -0,0 +1,83 @@
+import pendulum
+from airflow.decorators import dag
+from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import (
+    KubernetesPodOperator,
+)
+
+from mattermost_dags.airflow_utils import MATTERMOST_DATAWAREHOUSE_IMAGE, pod_defaults, send_alert
+from mattermost_dags.kube_secrets import (
+    SNOWFLAKE_LOAD_USER,
+    SNOWFLAKE_LOAD_PASSWORD,
+    SNOWFLAKE_ACCOUNT,
+    SNOWFLAKE_LOAD_DATABASE,
+    SNOWFLAKE_LOAD_WAREHOUSE,
+)
+
+
+@dag(
+    schedule=None,
+    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
+    catchup=False,
+    on_failure_callback=send_alert,
+)
+def push_proxy_loader():
+    """
+    ### Push proxy DAG
+
+    Loads ALB logs to Snowflake.
+    """
+
+    extract_hpns_us_logs = KubernetesPodOperator(
+        **pod_defaults,
+        image=MATTERMOST_DATAWAREHOUSE_IMAGE,  # Uses latest build from master
+        task_id="task-push-proxy-us",
+        name=f"push-proxy-us",
+        secrets=[
+            SNOWFLAKE_LOAD_USER,
+            SNOWFLAKE_LOAD_PASSWORD,
+            SNOWFLAKE_ACCOUNT,
+            SNOWFLAKE_LOAD_DATABASE,
+            SNOWFLAKE_LOAD_WAREHOUSE,
+        ],
+        env_vars={},
+        arguments=[
+            "push_proxy PUSH_PROXY_LOGS_US_NEW LOGS_US_NEW "
+            " --prefix /AWSLogs/{{ var.value.push_proxy_aws_account_id }}/elasticloadbalancing/{{ var.value.push_proxy_aws_region }}"
+            " -s {{ var.value.push_proxy_target_schema }}"
+            " -d ${SNOWFLAKE_LOAD_DATABASE}"
+            " -w ${SNOWFLAKE_LOAD_WAREHOUSE}"
+            " -r ${SNOWFLAKE_LOAD_ROLE}"
+            " -u ${SNOWFLAKE_LOAD_USER}"
+            " -p ${SNOWFLAKE_LOAD_PASSWORD}"
+        ],
+    )
+
+    extract_tpns_logs = KubernetesPodOperator(
+        **pod_defaults,
+        image=MATTERMOST_DATAWAREHOUSE_IMAGE,  # Uses latest build from master
+        task_id="task-push-proxy-tests",
+        name=f"push-proxy-test",
+        secrets=[
+            SNOWFLAKE_LOAD_USER,
+            SNOWFLAKE_LOAD_PASSWORD,
+            SNOWFLAKE_ACCOUNT,
+            SNOWFLAKE_LOAD_DATABASE,
+            SNOWFLAKE_LOAD_WAREHOUSE,
+        ],
+        env_vars={},
+        arguments=[
+            "push_proxy PUSH_PROXY_LOGS_TEST_NEW LOGS_TEST_NEW "
+            " --prefix /AWSLogs/{{ var.value.push_proxy_aws_account_id }}/elasticloadbalancing/{{ var.value.push_proxy_aws_region }}"
+            " -s {{ var.value.push_proxy_target_schema }}"
+            " -d ${SNOWFLAKE_LOAD_DATABASE}"
+            " -w ${SNOWFLAKE_LOAD_WAREHOUSE}"
+            " -r ${SNOWFLAKE_LOAD_ROLE}"
+            " -u ${SNOWFLAKE_LOAD_USER}"
+            " -p ${SNOWFLAKE_LOAD_PASSWORD}"
+        ],
+    )
+
+    extract_hpns_us_logs >> extract_tpns_logs
+
+
+dag = push_proxy_loader()