Infra/claude mcp oauth

orchestrate/dags/daily_loan_run__dc_test_cluster.py

@@ -0,0 +1,100 @@
+"""
+## Sample DAG showing end-to-end ELT (DC test cluster)
+This DAG shows how to load data with 3 tools, then run dbt, then other tasks.
+
+Same pipeline as `daily_loan_run`, kept separate because the Airbyte
+connection id differs between Datacoves clusters.
+"""
+
+from airflow.decorators import dag, task, task_group
+from orchestrate.utils import datacoves_utils
+
+from fivetran_provider_async.operators import FivetranOperator
+from fivetran_provider_async.sensors import FivetranSensor
+from airflow.providers.airbyte.operators.airbyte import AirbyteTriggerSyncOperator
+
+@dag(
+    doc_md = __doc__,
+    catchup = False,
+
+    default_args = datacoves_utils.set_default_args(
+        owner = "Noel Gomez",
+        owner_email = "noel@example.com"
+    ),
+
+    description = "Sample DAG to synchronize the Airflow database (DC test cluster)",
+    schedule = datacoves_utils.set_schedule("0 0 1 */12 *"),
+    tags=["extract_and_load", "transform", "marketing_automation", "update_catalog"],
+)
+def daily_loan_run__dc_test_cluster():
+
+    @task_group(
+        group_id="extract_and_load_airbyte",
+        tooltip="Airbyte Extract and Load"
+    )
+    def extract_and_load_airbyte():
+        # Extact and load
+        sync_airbyte = AirbyteTriggerSyncOperator(
+            task_id="country_populations_datacoves_snowflake",
+            connection_id="9cccf1c4-978e-4b0a-9677-43c0111dcda9",
+            airbyte_conn_id="airbyte_connection",
+        )
+
+
+    @task_group(
+        group_id="extract_and_load_fivetran",
+        tooltip="Fivetran Extract and Load"
+    )
+    def extract_and_load_fivetran():
+        trigger_fivetran =  FivetranOperator(
+            task_id="datacoves_snowflake_google_analytics_4_trigger",
+            fivetran_conn_id="fivetran_connection",
+            connector_id="speak_menial",
+            wait_for_completion=False,
+        )
+        sensor_fivetran = FivetranSensor(
+            task_id="datacoves_snowflake_google_analytics_4_sensor",
+            fivetran_conn_id="fivetran_connection",
+            connector_id="speak_menial",
+            poke_interval=60,
+        )
+        trigger_fivetran >> sensor_fivetran
+
+
+    @task_group(
+        group_id="extract_and_load_dlt",
+        tooltip="dlt Extract and Load"
+    )
+    def extract_and_load_dlt():
+        @task.datacoves_bash(
+            env = datacoves_utils.set_dlt_env_vars({"destinations": ["main_load_keypair"]}),
+            append_env=True
+        )
+        def load_loans_data():
+
+            return "cd load/dlt && ./loans_data.py"
+
+        load_loans_data()
+
+
+    # Transform Data
+    @task.datacoves_dbt(
+        connection_id="main_key_pair"
+    )
+    def transform():
+        return "dbt build -s 'tag:daily_run_airbyte+ tag:daily_run_fivetran+'"
+
+
+    # Post transformation tasks
+    @task.datacoves_bash
+    def marketing_automation():
+        return "echo 'send data to marketing tool'"
+
+    @task.datacoves_bash
+    def update_catalog():
+        return "echo 'refresh data catalog'"
+
+    extract_and_load = [extract_and_load_airbyte(), extract_and_load_fivetran(), extract_and_load_dlt()]
+    extract_and_load >> transform() >> [marketing_automation(), update_catalog()]
+
+daily_loan_run__dc_test_cluster()

orchestrate/dags/notifications_examples/yaml_slack_dag.py

@@ -1,18 +1,20 @@
 import datetime
+
 from airflow.decorators import dag, task
 from airflow.providers.slack.notifications.slack import send_slack_notification
 
+
 @dag(
+    description="Sample DAG with Slack notification, custom image, and resource requests",
+    schedule="0 0 1 */12 *",
+    tags=["transform", "slack_notification"],
     default_args={
         "start_date": datetime.datetime(2024, 1, 1, 0, 0),
         "owner": "Noel Gomez",
         "email": "gomezn@example.com",
         "email_on_failure": True,
         "retries": 3,
     },
-    description="Sample DAG with Slack notification, custom image, and resource requests",
-    schedule="0 0 1 */12 *",
-    tags=["transform", "slack_notification"],
     catchup=False,
     on_success_callback=send_slack_notification(
         text="The DAG {{ dag.dag_id }} succeeded", channel="#general"
@@ -22,12 +24,13 @@
     ),
 )
 def yaml_slack_dag():
-
-    @task.datacoves_dbt(connection_id="main_key_pair")
+    @task.datacoves_dbt(
+        connection_id="main_key_pair",
+    )
     def transform():
         return "dbt run -s personal_loans"
 
-    transform()
+    transform = transform()
+
 
-# Invoke DAG
 dag = yaml_slack_dag()

orchestrate/dags/notifications_examples/yaml_teams_dag.py

@@ -1,18 +1,20 @@
 import datetime
+
 from airflow.decorators import dag, task
 from notifiers.datacoves.ms_teams import MSTeamsNotifier
 
+
 @dag(
+    description="Sample DAG with MS Teams notification",
+    schedule="0 0 1 */12 *",
+    tags=["transform", "ms_teams_notification"],
     default_args={
         "start_date": datetime.datetime(2024, 1, 1, 0, 0),
         "owner": "Noel Gomez",
         "email": "gomezn@example.com",
         "email_on_failure": True,
         "retries": 3,
     },
-    description="Sample DAG with MS Teams notification",
-    schedule="0 0 1 */12 *",
-    tags=["transform", "ms_teams_notification"],
     catchup=False,
     on_success_callback=MSTeamsNotifier(
         connection_id="DATACOVES_MS_TEAMS", theme_color="0000FF"
@@ -22,12 +24,13 @@
     ),
 )
 def yaml_teams_dag():
-
-    @task.datacoves_dbt(connection_id="main_key_pair")
+    @task.datacoves_dbt(
+        connection_id="main_key_pair",
+    )
     def transform():
         return "dbt run -s personal_loans"
 
-    transform()
+    transform = transform()
+
 
-# Invoke DAG
 dag = yaml_teams_dag()

orchestrate/dags/other_examples/dbt_templated_overrides_dag.py

@@ -0,0 +1,50 @@
+"""
+## Sample DAG showing templated overrides on the dbt decorator
+This DAG shows how to parameterize the `@task.datacoves_dbt` `overrides` kwarg
+at trigger time using Jinja templates and DAG run `params`, instead of having
+to hardcode the schema/role/warehouse or fall back to an Airflow Variable.
+"""
+
+from airflow.decorators import dag, task
+from airflow.models.param import Param
+from orchestrate.utils import datacoves_utils
+
+@dag(
+    doc_md = __doc__,
+    catchup = False,
+
+    default_args = datacoves_utils.set_default_args(
+        owner = "Ian",
+        owner_email = "ian@example.com"
+    ),
+
+    description="Sample DAG showing templated overrides on the dbt decorator",
+    schedule = datacoves_utils.set_schedule(None),
+
+    params = {
+        "dbt_command": Param(
+            "dbt build --select models/l1/schema_name+",
+            type="string",
+            description="dbt command to run",
+        ),
+        "schema_override": Param(
+            "MANUAL_REFRESH",
+            type="string",
+            description="Schema to build into for this run",
+        ),
+    },
+
+    tags=["transform", "parameters"],
+)
+def dbt_templated_overrides():
+
+    @task.datacoves_dbt(
+        connection_id="main_key_pair",
+        overrides={"schema": "{{ params.schema_override }}"},
+    )
+    def run_dbt(dbt_command):
+        return dbt_command
+
+    run_dbt("{{ params.dbt_command }}")
+
+dbt_templated_overrides()

orchestrate/dags/other_examples/retry_dbt_failures.py

@@ -12,7 +12,7 @@
 
     default_args=datacoves_utils.set_default_args(
         owner = "Noel Gomez",
-        owner_email = "noel@example.com"
+        owner_email = "gomezn@example.com"
     ),
 
     schedule = datacoves_utils.set_schedule("0 0 1 */12 *"),

orchestrate/dags/other_examples/test_oom.py

@@ -0,0 +1,62 @@
+"""
+## OOM Test DAG
+Deliberately allocates ~5GB of memory in a single task, then holds it for
+2 minutes before releasing. Used to test that OOM detection/alerting pops up a
+message when a task/pod runs out of memory.
+
+Trigger manually - not scheduled.
+
+Note: whether this actually gets OOM-killed depends on the memory
+limit configured for the worker pod. If the pod's limit is below the
+`target_gb` below, the task will be killed while allocating. If the limit
+is higher, the task will succeed after holding the memory for `hold_seconds`.
+"""
+from pendulum import datetime
+
+from airflow.decorators import dag, task
+
+GB = 1024 ** 3
+
+default_args = {
+    "start_date": datetime(2024, 1, 1),
+    "owner": "Fernando Mercado",
+    "email_on_failure": False,
+    "retries": 0,
+}
+
+
+@dag(
+    doc_md=__doc__,
+    catchup=False,
+    schedule=None,
+    default_args=default_args,
+    tags=["test", "oom"],
+    dag_id="test_oom",
+)
+def test_oom():
+
+    @task
+    def consume_memory(target_gb: int = 5, hold_seconds: int = 120):
+        import time
+
+        chunk_size = 256 * 1024 * 1024  # 256MB per chunk
+        target_bytes = target_gb * GB
+
+        chunks = []
+        allocated = 0
+
+        while allocated < target_bytes:
+            # bytearray() zero-fills on creation, which forces the pages
+            # to become resident rather than just reserved
+            chunks.append(bytearray(chunk_size))
+            allocated += chunk_size
+            print(f"Allocated {allocated / GB:.2f} GB")
+
+        print(f"Holding {allocated / GB:.2f} GB for {hold_seconds}s")
+        time.sleep(hold_seconds)
+        print("Releasing memory")
+
+    consume_memory()
+
+
+test_oom()

orchestrate/dags_yml_definitions/daily_loan_run.yml

@@ -1,8 +1,9 @@
 doc_md: |
-  """
   ## Sample DAG showing end-to-end ELT
   This DAG shows how to load data with 3 tools, then run dbt, then other tasks
-  """
+
+imports:
+  - "from orchestrate.utils import datacoves_utils"
 description: "Loan Run"
 schedule: "0 0 1 */12 *"
 tags:
@@ -35,15 +36,17 @@ nodes:
     tooltip: "dlt Extract and Load"
 
     tasks:
-      load_us_population:
+      load_loans_data:
         task_decorator: datacoves_bash
-        bash_command: "./load/dlt/load_data.py"
+        bash_command: "cd load/dlt && ./loans_data.py"
+        env: !py 'datacoves_utils.set_dlt_env_vars({"destinations": ["main_load_keypair"]})'
+        append_env: true
 
   transform:
     type: task
     task_decorator: datacoves_dbt
-    connection_id: main
-    bash_command: "dbt build -s 'tag:daily_run_airbyte+ tag:daily_run_fivetran+ -t prd'"
+    connection_id: main_key_pair
+    bash_command: "dbt build -s 'tag:daily_run_airbyte+ tag:daily_run_fivetran+'"
     dependencies:
       [
         "extract_and_load_airbyte",
@@ -54,6 +57,7 @@ nodes:
   marketing_automation:
     task_decorator: datacoves_bash
     type: task
+
     bash_command: "echo 'send data to marketing tool'"
     dependencies: ["transform"]
 

orchestrate/dags_yml_definitions/dbt_dag.yml

@@ -25,4 +25,10 @@ nodes:
     type: task
     task_decorator: datacoves_dbt
     connection_id: main_key_pair
+    # Task-level override: fail fast instead of inheriting the retries
+    # from default_args.
+    retries: 0
+    # execution_timeout wants a timedelta, and !py is only resolved for
+    # DAG-level keys -- at node level dbt-coves emits it as a quoted
+    # string, which Airflow rejects. Left out until that is fixed.
     bash_command: "dbt debug"

orchestrate/dags_yml_definitions/notifications_examples/yaml_slack_dag.yml

@@ -29,7 +29,7 @@ notifications:
 nodes:
   transform:
     task_decorator: datacoves_dbt
-    connection_id: main
+    connection_id: main_key_pair
     type: task
     bash_command: "dbt run -s personal_loans"
 

orchestrate/dags_yml_definitions/notifications_examples/yaml_teams_dag.yml

@@ -31,7 +31,7 @@ notifications:
 nodes:
   transform:
     task_decorator: datacoves_dbt
-    connection_id: main
+    connection_id: main_key_pair
     type: task
     bash_command: "dbt run -s personal_loans"