Support both virtualenv use-cases in two task groups in the example DAG

astronomer · Oct 23, 2023 · 02dbd9a · 02dbd9a
1 parent 89913c2
commit 02dbd9a
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 23 deletions.
diff --git a/dev/dags/example_virtualenv.py b/dev/dags/example_virtualenv.py
@@ -5,7 +5,11 @@
 from datetime import datetime
 from pathlib import Path
 
-from cosmos import DbtDag, ExecutionMode, ExecutionConfig, ProjectConfig, ProfileConfig
+from airflow.decorators import dag
+from airflow.configuration import get_airflow_home
+from airflow.operators.empty import EmptyOperator
+
+from cosmos import DbtTaskGroup, ExecutionMode, ExecutionConfig, ProjectConfig, ProfileConfig
 from cosmos.profiles import PostgresUserPasswordProfileMapping
 
 DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt"
@@ -21,29 +25,57 @@
     ),
 )
 
-# [START virtualenv_example]
-example_virtualenv = DbtDag(
-    # dbt/cosmos-specific parameters
-    project_config=ProjectConfig(
-        DBT_ROOT_PATH / "jaffle_shop",
-    ),
-    profile_config=profile_config,
-    execution_config=ExecutionConfig(
-        execution_mode=ExecutionMode.VIRTUALENV,
-        # We can enable this flag if we want Airflow to create one virtualenv
-        # and reuse that within the whole DAG.
-        # virtualenv_dir=f"{get_airflow_home()}/persistent-venv",
-    ),
-    operator_args={
-        "py_system_site_packages": False,
-        "py_requirements": ["dbt-postgres==1.6.0b1"],
-        "install_deps": True,
-    },
-    # normal dag parameters
+@dag(
     schedule_interval="@daily",
     start_date=datetime(2023, 1, 1),
     catchup=False,
-    dag_id="example_virtualenv",
-    default_args={"retries": 2},
 )
+def example_virtualenv() -> None:
+    start_task = EmptyOperator(task_id='start-venv-examples')
+    end_task = EmptyOperator(task_id='end-venv-examples')
+
+    tmp_venv_task_group = DbtTaskGroup(
+        group_id='tmp-venv-group',
+        # dbt/cosmos-specific parameters
+        project_config=ProjectConfig(
+            DBT_ROOT_PATH / "jaffle_shop",
+        ),
+        profile_config=profile_config,
+        execution_config=ExecutionConfig(
+            execution_mode=ExecutionMode.VIRTUALENV,
+            # We can enable this flag if we want Airflow to create one virtualenv
+            # and reuse that within the whole DAG.
+            # virtualenv_dir=f"{get_airflow_home()}/persistent-venv",
+        ),
+        operator_args={
+            "py_system_site_packages": False,
+            "py_requirements": ["dbt-postgres==1.6.0b1"],
+            "install_deps": True,
+        },
+    )
+
+    cached_venv_task_group = DbtTaskGroup(
+        group_id='cached-venv-group',
+        # dbt/cosmos-specific parameters
+        project_config=ProjectConfig(
+            DBT_ROOT_PATH / "jaffle_shop",
+        ),
+        profile_config=profile_config,
+        execution_config=ExecutionConfig(
+            execution_mode=ExecutionMode.VIRTUALENV,
+            # We can enable this flag if we want Airflow to create one virtualenv
+            # and reuse that within the whole DAG.
+            virtualenv_dir=Path(f"{get_airflow_home()}/persistent-venv"),
+        ),
+        operator_args={
+            "py_system_site_packages": False,
+            "py_requirements": ["dbt-postgres==1.6.0b1"],
+            "install_deps": True,
+        },
+    )
+
+    start_task >> [tmp_venv_task_group, cached_venv_task_group] >> end_task
+
+example_virtualenv()
 # [END virtualenv_example]
+
diff --git a/tests/test_converter.py b/tests/test_converter.py
@@ -81,7 +81,8 @@ def test_converter_creates_dag_with_seed(mock_load_dbt_graph, execution_mode, op
 @patch("cosmos.converter.DbtGraph.load")
 def test_converter_raises_warning(mock_load_dbt_graph, execution_mode, virtualenv_dir, operator_args, caplog):
     """
-    This test will raise a warning if we are trying to pass ExecutionMode != `VirtualEnv` andm still pass a defined `virtualenv_dir`
+    This test will raise a warning if we are trying to pass ExecutionMode != `VirtualEnv` 
+    and still pass a defined `virtualenv_dir`
     """
     project_config = ProjectConfig(dbt_project_path=SAMPLE_DBT_PROJECT)
     execution_config = ExecutionConfig(execution_mode=execution_mode, virtualenv_dir=virtualenv_dir)