update python dag

astronomer · Jul 24, 2023 · 6f64246 · 6f64246
1 parent dac3179
commit 6f64246
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git 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
@@ -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),