Feature/datacoves3.3 (#229)

* Updating daily_loan
* Update Airflow API example
* Fivetran refactor
* Change examples to use decorators
* Change airflow connection name
* Testing instance aware dags for airflow.
* Testing is_development_environment
* Move get_schedule to utils and add airflow aware dag
* Testing key pair airflow service connection
* updating versions and adding overrides
* Change airbyte connection id
* Add note for constants file
* update decorator with env
* update dataset outlet
* Update DAG to use >> dependency management
* Fix dlt
* Remove dbt out of airflow github action
* remove variable import

---------

Co-authored-by: Noel Gomez 
Co-authored-by: Mayra Pena

Author: mayrapena1324

.github/workflows/10_integrate_airflow_changes.yml

@@ -26,7 +26,6 @@ jobs:
       AIRFLOW__CORE__DAGS_FOLDER: /__w/${{ github.event.repository.name }}/${{ github.event.repository.name }}/orchestrate/dags
       AIRFLOW__CORE__DAGBAG_IMPORT_TIMEOUT: 300
       AIRFLOW__ARTIFACTS_PATH: /__w/${{ github.event.repository.name }}/${{ github.event.repository.name }}/orchestrate
-      DBT_PROFILES_DIR: /__w/${{ github.event.repository.name }}/${{ github.event.repository.name }}/automate/dbt
       DATACOVES__DBT_HOME: /__w/${{ github.event.repository.name }}/${{ github.event.repository.name }}/transform
       DATACOVES__REPO_PATH: /__w/${{ github.event.repository.name }}/${{ github.event.repository.name }}
       PYTHONPATH: /__w/${{ github.event.repository.name }}/${{ github.event.repository.name }}

.github/workflows/10_integrate_dbt_changes.yml

@@ -27,7 +27,7 @@ jobs:
     # environment: PR_ENV
 
     # most people should use this one
-    container: datacoves/ci-basic-dbt-snowflake:3.2
+    container: datacoves/ci-basic-dbt-snowflake:3.3
 
     defaults:
       run:

.github/workflows/20_drop_integration_db.yml

@@ -21,7 +21,7 @@ jobs:
     # Alternatively, You can define multiple ENV for different workflows.
     # https://github.com/<org>/<repo>/settings/environments
     # environment: PR_ENV
-    container: datacoves/ci-basic-dbt-snowflake:3.2
+    container: datacoves/ci-basic-dbt-snowflake:3.3
 
     defaults:
       run:

.github/workflows/30_deploy_changes_to_production.yml

@@ -26,7 +26,7 @@ jobs:
     # Alternatively, You can define multiple ENV for different workflows.
     # https://github.com/<org>/<repo>/settings/environments
     # environment: PR_ENV
-    container: datacoves/ci-basic-dbt-snowflake:3.2
+    container: datacoves/ci-basic-dbt-snowflake:3.3
 
     defaults:
       run:

.pre-commit-config.yaml

@@ -2,7 +2,7 @@ files: ^transform/models/
 
 repos:
   - repo: https://github.com/dbt-checkpoint/dbt-checkpoint
-    rev: v2.0.5
+    rev: v2.0.6
 
     hooks:
       - id: check-source-table-has-description
@@ -27,7 +27,7 @@ repos:
         additional_dependencies:
           [
             "sqlfluff-templater-dbt==3.1.1",
-            "dbt-core==1.8.7",
+            "dbt-core==1.8.8",
             "dbt-snowflake==1.8.4",
           ]
         args: [--config, transform/.sqlfluff]

orchestrate/dags/bad_variable_usage.py

@@ -1,14 +1,11 @@
 import datetime
-
-from airflow.decorators import dag, task_group
+from airflow.decorators import dag, task, task_group
 from airflow.models import Variable
-from operators.datacoves.bash import DatacovesBashOperator
 
-# This is Here to show what NOT to do. When done this way, Aiflow will
-# query for this variable on every parse (every 30 secs). This can be
-# bad if using an external secrets manager like AWS Secrets Manager.
-# Doing this will incur significant AWS charges
-# The proper way to get a value is to do this in a method with the @task decorator
+
+# ❌ BAD PRACTICE: Fetching a variable at the top level
+# This will cause Airflow to query for this variable on EVERY DAG PARSE (every 30 seconds),
+# which can be costly when using an external secrets manager (e.g., AWS Secrets Manager).
 bad_used_variable = Variable.get("bad_used_variable", "default_value")
 
 @dag(
@@ -19,20 +16,24 @@
         "email_on_failure": True,
         "retries": 3
     },
-    description="Sample DAG for dbt build",
+    description="Sample DAG demonstrating bad variable usage",
     schedule="0 0 1 */12 *",
     tags=["extract_and_load"],
     catchup=False,
 )
 def bad_variable_usage():
+    
     @task_group(group_id="extract_and_load_dlt", tooltip="dlt Extract and Load")
     def extract_and_load_dlt():
-        load_us_population = DatacovesBashOperator(
-            task_id="load_us_population",
-            bash_command="cd load/dlt && ./loans_data.py",
-        )
+        """Task group for DLT extract and load process"""
+
+        @task.datacoves_bash(env={"BAD_VAR": bad_used_variable})  # ✅ Passing the bad variable to the task
+        def load_us_population():
+            return "cd load/dlt && ./loans_data.py"
 
-    tg_extract_and_load_dlt = extract_and_load_dlt()
+        load_us_population()
 
+    extract_and_load_dlt()
 
+# Invoke DAG
 dag = bad_variable_usage()

orchestrate/dags/daily_loan_run.py

@@ -1,12 +1,6 @@
 import datetime
-
-from airflow.decorators import dag, task_group
-from airflow.providers.airbyte.operators.airbyte import \
-    AirbyteTriggerSyncOperator
-from fivetran_provider_async.operators import FivetranOperator
-from fivetran_provider_async.sensors import FivetranSensor
-from operators.datacoves.bash import DatacovesBashOperator
-from operators.datacoves.dbt import DatacovesDbtOperator
+from airflow.decorators import dag, task, task_group
+from datahub_airflow_plugin.entities import Dataset
 
 @dag(
     default_args={"start_date": datetime.datetime(2024, 1, 1, 0, 0), "retries": 3},
@@ -16,79 +10,98 @@
     catchup=False,
 )
 def daily_loan_run():
+
     @task_group(group_id="extract_and_load_airbyte", tooltip="Airbyte Extract and Load")
     def extract_and_load_airbyte():
-        country_populations_datacoves_snowflake = AirbyteTriggerSyncOperator(
-            task_id="country_populations_datacoves_snowflake",
-            connection_id="ac02ea96-58a1-4061-be67-78900bb5aaf6",
-            airbyte_conn_id="airbyte_connection",
-        )
+        @task
+        def sync_airbyte():
+            from airflow.providers.airbyte.operators.airbyte import AirbyteTriggerSyncOperator
+            return AirbyteTriggerSyncOperator(
+                task_id="country_populations_datacoves_snowflake",
+                connection_id="b293aaea-6557-4506-8cfb-6b621ec4c6ef",
+                airbyte_conn_id="airbyte_connection",
+            ).execute({})
+
+        sync_airbyte()
 
     tg_extract_and_load_airbyte = extract_and_load_airbyte()
 
-    @task_group(
-        group_id="extract_and_load_fivetran", tooltip="Fivetran Extract and Load"
-    )
+    @task_group(group_id="extract_and_load_fivetran", tooltip="Fivetran Extract and Load")
     def extract_and_load_fivetran():
-        datacoves_snowflake_google_analytics_4_trigger = FivetranOperator(
-            task_id="datacoves_snowflake_google_analytics_4_trigger",
-            fivetran_conn_id="fivetran_connection",
-            connector_id="speak_menial",
-            wait_for_completion=False,
-        )
-        datacoves_snowflake_google_analytics_4_sensor = FivetranSensor(
-            task_id="datacoves_snowflake_google_analytics_4_sensor",
-            fivetran_conn_id="fivetran_connection",
-            connector_id="speak_menial",
-            poke_interval=60,
-        )
-        (
-            datacoves_snowflake_google_analytics_4_trigger
-            >> datacoves_snowflake_google_analytics_4_sensor
-        )
+
+        @task
+        def trigger_fivetran():
+            from fivetran_provider_async.operators import FivetranOperator
+            return FivetranOperator(
+                task_id="datacoves_snowflake_google_analytics_4_trigger",
+                fivetran_conn_id="fivetran_connection",
+                connector_id="speak_menial",
+                wait_for_completion=False,
+            ).execute({})
+
+        @task
+        def sensor_fivetran():
+            from fivetran_provider_async.sensors import FivetranSensor
+            return FivetranSensor(
+                task_id="datacoves_snowflake_google_analytics_4_sensor",
+                fivetran_conn_id="fivetran_connection",
+                connector_id="speak_menial",
+                poke_interval=60,
+            ).poke({})
+
+        trigger = trigger_fivetran()
+        sensor = sensor_fivetran()
+
+        trigger >> sensor
+        return sensor  # Return last task in the group
 
     tg_extract_and_load_fivetran = extract_and_load_fivetran()
 
     @task_group(group_id="extract_and_load_dlt", tooltip="dlt Extract and Load")
     def extract_and_load_dlt():
-        load_us_population = DatacovesBashOperator(
-            task_id="load_loads_data",
-            bash_command="""
-                cd load/dlt \
-                && ./loans_data.py
-            """,
+        @task.datacoves_bash(
+            outlets=[Dataset("snowflake", "raw.loans_data.loans_data")],
             env={
                 "UV_CACHE_DIR": "/tmp/uv_cache",
-                "EXTRACT__NEXT_ITEM_MODE":"fifo",
-                "EXTRACT__MAX_PARALLEL_ITEMS":"1",
-                "EXTRACT__WORKERS":"1",
-                "NORMALIZE__WORKERS":"1",
-                "LOAD__WORKERS":"1",
+                "EXTRACT__NEXT_ITEM_MODE": "fifo",
+                "EXTRACT__MAX_PARALLEL_ITEMS": "1",
+                "EXTRACT__WORKERS": "1",
+                "NORMALIZE__WORKERS": "1",
+                "LOAD__WORKERS": "1",
             },
-            append_env=True,
+            append_env=True
         )
+        def load_loans_data():
+            return "cd load/dlt && ./loans_data.py"
+        load_loans_data()
 
     tg_extract_and_load_dlt = extract_and_load_dlt()
-    transform = DatacovesDbtOperator(
-        task_id="transform",
-        bash_command="dbt build -s 'tag:daily_run_airbyte+ tag:daily_run_fivetran+'",
-    )
-    transform.set_upstream(
-        [
-            tg_extract_and_load_airbyte,
-            tg_extract_and_load_dlt,
-            tg_extract_and_load_fivetran,
+
+    @task.datacoves_dbt(
+        connection_id="main",
+        inlets=[
+            Dataset("snowflake", "raw.loans_data.loans_data"),
+            Dataset("snowflake", "raw.google_analytics_4.engagement_events_report")
         ]
     )
-    marketing_automation = DatacovesBashOperator(
-        task_id="marketing_automation",
-        bash_command="echo 'send data to marketing tool'",
-    )
-    marketing_automation.set_upstream([transform])
-    update_catalog = DatacovesBashOperator(
-        task_id="update_catalog", bash_command="echo 'refresh data catalog'"
-    )
-    update_catalog.set_upstream([transform])
+    def transform():
+        return "dbt build -s 'tag:daily_run_airbyte+ tag:daily_run_fivetran+ -t prd'"
+
+    @task.datacoves_bash
+    def marketing_automation():
+        return "echo 'send data to marketing tool'"
+
+    @task.datacoves_bash
+    def update_catalog():
+        return "echo 'refresh data catalog'"
+
+
+    transform_task = transform()
+    marketing_automation_task = marketing_automation()
+    update_catalog_task = update_catalog()
 
+    [tg_extract_and_load_airbyte, tg_extract_and_load_dlt, tg_extract_and_load_fivetran] >> transform_task
+    transform_task >> [marketing_automation_task, update_catalog_task]
 
+# Invoke DAG
 dag = daily_loan_run()

orchestrate/dags/data_aware/consumer.py

@@ -2,7 +2,7 @@
 
 from airflow.decorators import dag, task
 from airflow.datasets import Dataset
-from operators.datacoves.bash import DatacovesBashOperator
+
 
 MY_SOURCE = Dataset("upstream_data")
 
@@ -20,10 +20,10 @@
 )
 def data_aware_consumer_dag():
     @task
-    def run_dbt():
+    def run_consumer():
         print("I'm the consumer")
 
-    run_dbt()
+    run_consumer()
 
 
 dag = data_aware_consumer_dag()

orchestrate/dags/data_aware/producer.py

@@ -2,7 +2,7 @@
 
 from airflow.decorators import dag, task
 from airflow.datasets import Dataset
-from operators.datacoves.bash import DatacovesBashOperator
+
 
 # A dataset can be anything, it will be a poiner in the Airflow db.
 # If you need to access url like s3://my_bucket/my_file.txt then you can set

orchestrate/dags/ng_test.py

@@ -0,0 +1,65 @@
+import datetime
+from airflow.decorators import dag, task
+from airflow.models import Variable
+from datahub_airflow_plugin.entities import Dataset
+
+
+# ❌ BAD PRACTICE: Fetching a variable at the top level
+# This will cause Airflow to query for this variable on EVERY DAG PARSE (every 30 seconds),
+# which can be costly when using an external secrets manager (e.g., AWS Secrets Manager).
+bad_used_variable = Variable.get("bad_used_variable", "default_value")
+
+@dag(
+    default_args={
+        "start_date": datetime.datetime(2024, 1, 1, 0, 0),
+        "owner": "Noel Gomez",
+        "email": "[email protected]",
+        "email_on_failure": True,
+        "retries": 3
+    },
+    description="Sample DAG demonstrating bad variable usage",
+    schedule="0 0 1 */12 *",
+    tags=["extract_and_load","transform"],
+    catchup=False,
+)
+def ng_test():
+
+    # @task.datacoves_dbt(connection_id="main")
+    # def show_env_value():
+    #     return """
+    #         echo dbt_home: && echo $DATACOVES__DBT_HOME &&
+    #         echo repo_path: && echo $DATACOVES__REPO_PATH &&
+    #         echo cwd: && pwd
+    #     """
+    # show_env_value()
+
+    @task.datacoves_bash(
+        outlets=[Dataset("snowflake",  "raw.us_population.us_population")],
+        env={
+            "UV_CACHE_DIR": "/tmp/uv_cache",
+            "EXTRACT__NEXT_ITEM_MODE":"fifo",
+            "EXTRACT__MAX_PARALLEL_ITEMS":"1",
+            "EXTRACT__WORKERS":"1",
+            "NORMALIZE__WORKERS":"1",
+            "LOAD__WORKERS":"1",
+        },
+        append_env=True
+    )
+    def load_us_population():
+        return "cd load/dlt/ && ./us_population.py"
+
+
+    @task.datacoves_dbt(
+        connection_id="main",
+        inlets=[
+            Dataset("snowflake",  "raw.us_population.us_population"),
+            Dataset("snowflake",  "raw.google_analytics_4.engagement_events_report")
+            ]
+    )
+    def run_dbt():
+        return "dbt build -s 'tag:daily_run_airbyte+ tag:daily_run_fivetran+ -t prd'"
+
+    load_us_population() >> run_dbt()
+
+# Invoke DAG
+dag = ng_test()