From 3e38af9ba63a55fbf6d02657097ebf3136e93510 Mon Sep 17 00:00:00 2001 From: Tatiana Al-Chueyr Date: Thu, 20 Jun 2024 15:45:47 +0100 Subject: [PATCH] Change example DAGs to use `example_conn` as opposed to `airflow_db` (#1054) Using the Airflow metadata database connection as an example connection is misleading. The mismatch in the environment variable value used in the Cosmos integration tests, particularly with sqlite as the Airflow metadata database, is an issue that can hide other underlining problems. This PR decouples the test connection used by Cosmos example DAGs from the Airflow metadata Database connection. Since this change affects the Github action configuration, it will only work for the branch-triggered GH action runs, such as: https://github.com/astronomer/astronomer-cosmos/actions/runs/9596066209 Because this is a breaking change to the CI script itself, all the tests `pull_request_target` are expected to fail during the PR - and will pass once this is merged to `main`. This improvement was originally part of #1014 --------- Co-authored-by: Pankaj Koti --- .github/workflows/test.yml | 18 +++++++++--------- dev/dags/basic_cosmos_dag.py | 4 ++-- dev/dags/basic_cosmos_task_group.py | 4 ++-- dev/dags/cosmos_manifest_example.py | 4 ++-- dev/dags/cosmos_profile_mapping.py | 2 +- dev/dags/cosmos_seed_dag.py | 2 +- dev/dags/dbt_docs.py | 2 +- dev/dags/example_model_version.py | 4 ++-- dev/dags/example_virtualenv.py | 4 ++-- dev/dags/performance_dag.py | 4 ++-- tests/dbt/test_graph.py | 4 ++-- tests/operators/test_local.py | 2 +- 12 files changed, 27 insertions(+), 27 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index dcac01496..379b96b65 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -140,7 +140,7 @@ jobs: hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-integration env: AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ - AIRFLOW_CONN_AIRFLOW_DB: postgres://postgres:postgres@0.0.0.0:5432/postgres + AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:postgres@0.0.0.0:5432/postgres DATABRICKS_HOST: mock DATABRICKS_WAREHOUSE_ID: mock DATABRICKS_TOKEN: mock @@ -210,7 +210,7 @@ jobs: DATABRICKS_UNIQUE_ID="${{github.run_id}}" hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-integration-expensive env: AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ - AIRFLOW_CONN_AIRFLOW_DB: postgres://postgres:postgres@0.0.0.0:5432/postgres + AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:postgres@0.0.0.0:5432/postgres PYTHONPATH: /home/runner/work/astronomer-cosmos/astronomer-cosmos/:$PYTHONPATH AIRFLOW_CONN_DATABRICKS_DEFAULT: ${{ secrets.AIRFLOW_CONN_DATABRICKS_DEFAULT }} DATABRICKS_CLUSTER_ID: ${{ secrets.DATABRICKS_CLUSTER_ID }} @@ -231,7 +231,7 @@ jobs: env: AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ - AIRFLOW_CONN_AIRFLOW_DB: postgres://postgres:postgres@0.0.0.0:5432/postgres + AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:postgres@0.0.0.0:5432/postgres PYTHONPATH: /home/runner/work/astronomer-cosmos/astronomer-cosmos/:$PYTHONPATH AIRFLOW_CONN_DATABRICKS_DEFAULT: ${{ secrets.AIRFLOW_CONN_DATABRICKS_DEFAULT }} DATABRICKS_CLUSTER_ID: ${{ secrets.DATABRICKS_CLUSTER_ID }} @@ -272,7 +272,7 @@ jobs: hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-integration-sqlite env: AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ - AIRFLOW_CONN_AIRFLOW_DB: postgres://postgres:postgres@0.0.0.0:5432/postgres + AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:postgres@0.0.0.0:5432/postgres AIRFLOW__CORE__DAGBAG_IMPORT_TIMEOUT: 90.0 PYTHONPATH: /home/runner/work/astronomer-cosmos/astronomer-cosmos/:$PYTHONPATH COSMOS_CONN_POSTGRES_PASSWORD: ${{ secrets.COSMOS_CONN_POSTGRES_PASSWORD }} @@ -295,7 +295,7 @@ jobs: env: AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ - AIRFLOW_CONN_AIRFLOW_DB: postgres://postgres:postgres@0.0.0.0:5432/postgres + AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:postgres@0.0.0.0:5432/postgres PYTHONPATH: /home/runner/work/astronomer-cosmos/astronomer-cosmos/:$PYTHONPATH Run-Integration-Tests-DBT-1-5-4: @@ -345,7 +345,7 @@ jobs: hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-integration-dbt-1-5-4 env: AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ - AIRFLOW_CONN_AIRFLOW_DB: postgres://postgres:postgres@0.0.0.0:5432/postgres + AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:postgres@0.0.0.0:5432/postgres AIRFLOW__CORE__DAGBAG_IMPORT_TIMEOUT: 90.0 PYTHONPATH: /home/runner/work/astronomer-cosmos/astronomer-cosmos/:$PYTHONPATH COSMOS_CONN_POSTGRES_PASSWORD: ${{ secrets.COSMOS_CONN_POSTGRES_PASSWORD }} @@ -368,7 +368,7 @@ jobs: env: AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ - AIRFLOW_CONN_AIRFLOW_DB: postgres://postgres:postgres@0.0.0.0:5432/postgres + AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:postgres@0.0.0.0:5432/postgres PYTHONPATH: /home/runner/work/astronomer-cosmos/astronomer-cosmos/:$PYTHONPATH Run-Performance-Tests: @@ -424,7 +424,7 @@ jobs: cat /tmp/performance_results.txt > $GITHUB_STEP_SUMMARY env: AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ - AIRFLOW_CONN_AIRFLOW_DB: postgres://postgres:postgres@0.0.0.0:5432/postgres + AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:postgres@0.0.0.0:5432/postgres AIRFLOW__CORE__DAGBAG_IMPORT_TIMEOUT: 90.0 PYTHONPATH: /home/runner/work/astronomer-cosmos/astronomer-cosmos/:$PYTHONPATH COSMOS_CONN_POSTGRES_PASSWORD: ${{ secrets.COSMOS_CONN_POSTGRES_PASSWORD }} @@ -437,7 +437,7 @@ jobs: MODEL_COUNT: ${{ matrix.num-models }} env: AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/ - AIRFLOW_CONN_AIRFLOW_DB: postgres://postgres:postgres@0.0.0.0:5432/postgres + AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:postgres@0.0.0.0:5432/postgres PYTHONPATH: /home/runner/work/astronomer-cosmos/astronomer-cosmos/:$PYTHONPATH Code-Coverage: diff --git a/dev/dags/basic_cosmos_dag.py b/dev/dags/basic_cosmos_dag.py index c961638b3..f71f351f0 100644 --- a/dev/dags/basic_cosmos_dag.py +++ b/dev/dags/basic_cosmos_dag.py @@ -1,5 +1,5 @@ """ -An example DAG that uses Cosmos to render a dbt project. +An example DAG that uses Cosmos to render a dbt project into an Airflow DAG. """ import os @@ -16,7 +16,7 @@ profile_name="default", target_name="dev", profile_mapping=PostgresUserPasswordProfileMapping( - conn_id="airflow_db", + conn_id="example_conn", profile_args={"schema": "public"}, disable_event_tracking=True, ), diff --git a/dev/dags/basic_cosmos_task_group.py b/dev/dags/basic_cosmos_task_group.py index f230e87d4..4221e3019 100644 --- a/dev/dags/basic_cosmos_task_group.py +++ b/dev/dags/basic_cosmos_task_group.py @@ -1,5 +1,5 @@ """ -An example DAG that uses Cosmos to render a dbt project as a TaskGroup. +An example DAG that uses Cosmos to render a dbt project as an Airflow TaskGroup. """ import os @@ -20,7 +20,7 @@ profile_name="default", target_name="dev", profile_mapping=PostgresUserPasswordProfileMapping( - conn_id="airflow_db", + conn_id="example_conn", profile_args={"schema": "public"}, ), ) diff --git a/dev/dags/cosmos_manifest_example.py b/dev/dags/cosmos_manifest_example.py index e8721aefd..225b38ddc 100644 --- a/dev/dags/cosmos_manifest_example.py +++ b/dev/dags/cosmos_manifest_example.py @@ -1,5 +1,5 @@ """ -An example DAG that uses Cosmos to render a dbt project. +An example DAG that uses Cosmos to render a dbt project into Airflow using a dbt manifest file. """ import os @@ -16,7 +16,7 @@ profile_name="default", target_name="dev", profile_mapping=PostgresUserPasswordProfileMapping( - conn_id="airflow_db", + conn_id="example_conn", profile_args={"schema": "public"}, dbt_config_vars=DbtProfileConfigVars(send_anonymous_usage_stats=True), ), diff --git a/dev/dags/cosmos_profile_mapping.py b/dev/dags/cosmos_profile_mapping.py index 3c9a503ba..36f344446 100644 --- a/dev/dags/cosmos_profile_mapping.py +++ b/dev/dags/cosmos_profile_mapping.py @@ -41,7 +41,7 @@ def cosmos_profile_mapping() -> None: profile_name="default", target_name="dev", profile_mapping=get_automatic_profile_mapping( - conn_id="airflow_db", + conn_id="example_conn", profile_args={"schema": "public"}, ), ), diff --git a/dev/dags/cosmos_seed_dag.py b/dev/dags/cosmos_seed_dag.py index b682ecc36..d0b4f4103 100644 --- a/dev/dags/cosmos_seed_dag.py +++ b/dev/dags/cosmos_seed_dag.py @@ -29,7 +29,7 @@ profile_name="default", target_name="dev", profile_mapping=PostgresUserPasswordProfileMapping( - conn_id="airflow_db", + conn_id="example_conn", profile_args={"schema": "public"}, ), ) diff --git a/dev/dags/dbt_docs.py b/dev/dags/dbt_docs.py index a7703d4e4..e26e10d53 100644 --- a/dev/dags/dbt_docs.py +++ b/dev/dags/dbt_docs.py @@ -35,7 +35,7 @@ profile_name="default", target_name="dev", profile_mapping=PostgresUserPasswordProfileMapping( - conn_id="airflow_db", + conn_id="example_conn", profile_args={"schema": "public"}, ), ) diff --git a/dev/dags/example_model_version.py b/dev/dags/example_model_version.py index 909c7494a..df9205f3e 100644 --- a/dev/dags/example_model_version.py +++ b/dev/dags/example_model_version.py @@ -1,5 +1,5 @@ """ -An example DAG that uses Cosmos to render a dbt project. +An example DAG that uses Cosmos to render a dbt project as an Airflow DAG. """ import os @@ -16,7 +16,7 @@ profile_name="default", target_name="dev", profile_mapping=PostgresUserPasswordProfileMapping( - conn_id="airflow_db", + conn_id="example_conn", profile_args={"schema": "public"}, ), ) diff --git a/dev/dags/example_virtualenv.py b/dev/dags/example_virtualenv.py index 6275b1048..cd38cba9e 100644 --- a/dev/dags/example_virtualenv.py +++ b/dev/dags/example_virtualenv.py @@ -1,5 +1,5 @@ """ -An example DAG that uses Cosmos to render a dbt project. +An example DAG that uses Cosmos to render a dbt project as an Airflow DAG. """ import os @@ -17,7 +17,7 @@ profile_name="default", target_name="dev", profile_mapping=PostgresUserPasswordProfileMapping( - conn_id="airflow_db", + conn_id="example_conn", profile_args={"schema": "public"}, ), ) diff --git a/dev/dags/performance_dag.py b/dev/dags/performance_dag.py index 1c8d639e0..3ed23c94c 100644 --- a/dev/dags/performance_dag.py +++ b/dev/dags/performance_dag.py @@ -1,5 +1,5 @@ """ -An airflow DAG that uses Cosmos to render a dbt project for performance testing. +An Airflow DAG that uses Cosmos to render a dbt project for performance testing. """ import os @@ -17,7 +17,7 @@ profile_name="default", target_name="dev", profile_mapping=PostgresUserPasswordProfileMapping( - conn_id="airflow_db", + conn_id="example_conn", profile_args={"schema": "public"}, ), ) diff --git a/tests/dbt/test_graph.py b/tests/dbt/test_graph.py index 5f778b763..0166dd89f 100644 --- a/tests/dbt/test_graph.py +++ b/tests/dbt/test_graph.py @@ -51,7 +51,7 @@ def postgres_profile_config() -> ProfileConfig: profile_name="default", target_name="default", profile_mapping=PostgresUserPasswordProfileMapping( - conn_id="airflow_db", + conn_id="example_conn", profile_args={"schema": "public"}, ), ) @@ -265,7 +265,7 @@ def test_load_automatic_without_manifest_with_profile_mapping(mock_load_via_dbt_ profile_name="test", target_name="test", profile_mapping=PostgresUserPasswordProfileMapping( - conn_id="airflow_db", + conn_id="example_conn", profile_args={"schema": "public"}, ), ) diff --git a/tests/operators/test_local.py b/tests/operators/test_local.py index f90237082..ed46caf88 100644 --- a/tests/operators/test_local.py +++ b/tests/operators/test_local.py @@ -56,7 +56,7 @@ profile_name="default", target_name="dev", profile_mapping=PostgresUserPasswordProfileMapping( - conn_id="airflow_db", + conn_id="example_conn", profile_args={"schema": "public"}, ), )