From 6f64246eed9e5653ebab5a2fcd4502b175e5657c Mon Sep 17 00:00:00 2001
From: Julian LaNeve <lanevejulian@gmail.com>
Date: Sun, 23 Jul 2023 23:31:10 -0400
Subject: [PATCH] update python dag

---
 dev/Dockerfile                           |  2 +-
 dev/dags/example_cosmos_python_models.py | 23 +++++++++++++----------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/dev/Dockerfile b/dev/Dockerfile
index 64dece432..90c49ed6c 100644
--- a/dev/Dockerfile
+++ b/dev/Dockerfile
@@ -7,7 +7,7 @@ COPY ./README.rst  ${AIRFLOW_HOME}/astronomer_cosmos/
 COPY ./cosmos/  ${AIRFLOW_HOME}/astronomer_cosmos/cosmos/
 
 # install the package in editable mode
-RUN pip install -e "${AIRFLOW_HOME}/astronomer_cosmos"[dbt-postgres]
+RUN pip install -e "${AIRFLOW_HOME}/astronomer_cosmos"[dbt-postgres,dbt-databricks]
 
 # make sure astro user owns the package
 RUN chown -R astro:astro ${AIRFLOW_HOME}/astronomer_cosmos
diff --git a/dev/dags/example_cosmos_python_models.py b/dev/dags/example_cosmos_python_models.py
index 2876e38c2..df4d61e71 100644
--- a/dev/dags/example_cosmos_python_models.py
+++ b/dev/dags/example_cosmos_python_models.py
@@ -17,7 +17,8 @@
 from datetime import datetime
 from pathlib import Path
 
-from cosmos import DbtDag
+from cosmos import DbtDag, ProjectConfig, ProfileConfig, ExecutionConfig
+from cosmos.profiles import DatabricksTokenProfileMapping
 
 DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt"
 DBT_ROOT_PATH = Path(os.getenv("DBT_ROOT_PATH", DEFAULT_DBT_ROOT_PATH))
@@ -26,15 +27,17 @@
 # [START example_cosmos_python_models]
 example_cosmos_python_models = DbtDag(
     # dbt/cosmos-specific parameters
-    dbt_root_path=DBT_ROOT_PATH,
-    dbt_project_name="jaffle_shop_python",
-    conn_id="databricks_default",
-    profile_args={
-        "schema": SCHEMA,
-    },
-    operator_args={"append_env": True},
-    profile_name_override="airflow",
-    target_name_override="dev_target",
+    project_config=ProjectConfig(
+        DBT_ROOT_PATH / "jaffle_shop_python",
+    ),
+    profile_config=ProfileConfig(
+        profile_name="default",
+        target_name="dev",
+        profile_mapping=DatabricksTokenProfileMapping(conn_id="databricks_default", profile_args={"schema": SCHEMA}),
+    ),
+    execution_config=ExecutionConfig(
+        append_env=True,
+    ),
     # normal dag parameters
     schedule_interval="@daily",
     start_date=datetime(2023, 1, 1),