Experimental BQ support to run dbt models with ExecutionMode.AIRFLOW_ASYNC

.github/workflows/test.yml

@@ -2,7 +2,7 @@ name: test
 
 on:
   push: # Run on pushes to the default branch
-    branches: [main]
+    branches: [main,poc-dbt-compile-task]
   pull_request_target: # Also run on pull requests originated from forks
     branches: [main]
 

cosmos/airflow/graph.py

@@ -81,6 +81,7 @@ def create_test_task_metadata(
     on_warning_callback: Callable[..., Any] | None = None,
     node: DbtNode | None = None,
     render_config: RenderConfig | None = None,
+    async_op_args: dict | None = None,
 ) -> TaskMetadata:
     """
     Create the metadata that will be used to instantiate the Airflow Task that will be used to run the Dbt test node.
@@ -125,6 +126,7 @@ def create_test_task_metadata(
         ),
         arguments=task_args,
         extra_context=extra_context,
+        async_op_args=async_op_args,
     )
 
 
@@ -288,6 +290,7 @@ def build_airflow_graph(
     render_config: RenderConfig,
     task_group: TaskGroup | None = None,
     on_warning_callback: Callable[..., Any] | None = None,  # argument specific to the DBT test command
+    async_op_args: dict | None = None,
 ) -> None:
     """
     Instantiate dbt `nodes` as Airflow tasks within the given `task_group` (optional) or `dag` (mandatory).
@@ -352,6 +355,7 @@ def build_airflow_graph(
             task_args=task_args,
             on_warning_callback=on_warning_callback,
             render_config=render_config,
+            async_op_args=async_op_args,
         )
         test_task = create_airflow_task(test_meta, dag, task_group=task_group)
         leaves_ids = calculate_leaves(tasks_ids=list(tasks_map.keys()), nodes=nodes)

cosmos/config.py

@@ -10,6 +10,7 @@
 from pathlib import Path
 from typing import Any, Callable, Iterator
 
+import yaml
 from airflow.version import version as airflow_version
 
 from cosmos.cache import create_cache_profile, get_cached_profile, is_profile_cache_enabled
@@ -286,6 +287,21 @@ def validate_profiles_yml(self) -> None:
         if self.profiles_yml_filepath and not Path(self.profiles_yml_filepath).exists():
             raise CosmosValueError(f"The file {self.profiles_yml_filepath} does not exist.")
 
+    def get_profile_type(self) -> str:
+        if isinstance(self.profile_mapping, BaseProfileMapping):
+            return str(self.profile_mapping.dbt_profile_type)
+
+        profile_path = self._get_profile_path()
+
+        with open(profile_path) as file:
+            profiles = yaml.safe_load(file)
+
+            profile = profiles[self.profile_name]
+            target_type = profile["outputs"][self.target_name]["type"]
+            return str(target_type)
+
+        return "undefined"
+
     def _get_profile_path(self, use_mock_values: bool = False) -> Path:
         """
         Handle the profile caching mechanism.

cosmos/converter.py

@@ -207,6 +207,7 @@ def __init__(
         task_group: TaskGroup | None = None,
         operator_args: dict[str, Any] | None = None,
         on_warning_callback: Callable[..., Any] | None = None,
+        async_op_args: dict[str, Any] | None = None,
         *args: Any,
         **kwargs: Any,
     ) -> None:
@@ -256,6 +257,7 @@ def __init__(
             cache_identifier=cache_identifier,
             dbt_vars=dbt_vars,
             airflow_metadata=cache._get_airflow_metadata(dag, task_group),
+            async_op_args=async_op_args,
         )
         self.dbt_graph.load(method=render_config.load_method, execution_mode=execution_config.execution_mode)
 
@@ -301,6 +303,7 @@ def __init__(
             dbt_project_name=project_config.project_name,
             on_warning_callback=on_warning_callback,
             render_config=render_config,
+            async_op_args=async_op_args,
         )
 
         current_time = time.perf_counter()

cosmos/core/airflow.py

@@ -31,6 +31,9 @@ def get_airflow_task(task: Task, dag: DAG, task_group: TaskGroup | None = None)
     if task.owner != "":
         task_kwargs["owner"] = task.owner
 
+    if task.async_op_args:
+        task_kwargs["async_op_args"] = task.async_op_args
+
     airflow_task = Operator(
         task_id=task.id,
         dag=dag,

cosmos/core/graph/entities.py

@@ -61,3 +61,4 @@ class Task(CosmosEntity):
     operator_class: str = "airflow.operators.empty.EmptyOperator"
     arguments: Dict[str, Any] = field(default_factory=dict)
     extra_context: Dict[str, Any] = field(default_factory=dict)
+    async_op_args: Dict[str, Any] = field(default_factory=dict)

cosmos/dbt/graph.py

@@ -217,13 +217,15 @@ def __init__(
         dbt_vars: dict[str, str] | None = None,
         airflow_metadata: dict[str, str] | None = None,
         operator_args: dict[str, Any] | None = None,
+        async_op_args: dict[str, Any] | None = None,
     ):
         self.project = project
         self.render_config = render_config
         self.profile_config = profile_config
         self.execution_config = execution_config
         self.cache_dir = cache_dir
         self.airflow_metadata = airflow_metadata or {}
+        self.async_op_args = async_op_args
         if cache_identifier:
             self.dbt_ls_cache_key = cache.create_cache_key(cache_identifier)
         else:
@@ -467,7 +469,6 @@ def should_use_dbt_ls_cache(self) -> bool:
 
     def load_via_dbt_ls_cache(self) -> bool:
         """(Try to) load dbt ls cache from an Airflow Variable"""
-
         logger.info(f"Trying to parse the dbt project using dbt ls cache {self.dbt_ls_cache_key}...")
         if self.should_use_dbt_ls_cache():
             project_path = self.project_path

cosmos/operators/airflow_async.py

@@ -1,18 +1,26 @@
+from __future__ import annotations
+
+from typing import Any
+
+from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
+from airflow.providers.google.cloud.operators.bigquery import BigQueryInsertJobOperator
+from airflow.utils.context import Context
+
+from cosmos import settings
+from cosmos.exceptions import CosmosValueError
 from cosmos.operators.local import (
     DbtBuildLocalOperator,
     DbtCompileLocalOperator,
-    DbtDocsAzureStorageLocalOperator,
-    DbtDocsGCSLocalOperator,
-    DbtDocsLocalOperator,
-    DbtDocsS3LocalOperator,
     DbtLSLocalOperator,
-    DbtRunLocalOperator,
     DbtRunOperationLocalOperator,
     DbtSeedLocalOperator,
     DbtSnapshotLocalOperator,
     DbtSourceLocalOperator,
     DbtTestLocalOperator,
 )
+from cosmos.settings import remote_target_path, remote_target_path_conn_id
+
+_SUPPORTED_DATABASES = ["bigquery"]
 
 
 class DbtBuildAirflowAsyncOperator(DbtBuildLocalOperator):
@@ -35,8 +43,81 @@ class DbtSourceAirflowAsyncOperator(DbtSourceLocalOperator):
     pass
 
 
-class DbtRunAirflowAsyncOperator(DbtRunLocalOperator):
-    pass
+class DbtRunAirflowAsyncOperator(BigQueryInsertJobOperator):  # type: ignore
+    def __init__(self, *args, full_refresh: bool = False, **kwargs):  # type: ignore
+        # dbt task param
+        self.profile_config = kwargs.get("profile_config")
+        self.project_dir = kwargs.get("project_dir")
+        self.file_path = kwargs.get("extra_context", {}).get("dbt_node_config", {}).get("file_path")
+        self.profile_type: str = self.profile_config.get_profile_type()  # type: ignore
+        self.full_refresh = full_refresh
+
+        # airflow task param
+        self.async_op_args = kwargs.pop("async_op_args", {})
+        self.configuration: dict[str, object] = {}
+        self.job_id = self.async_op_args.get("job_id", "")
+        self.impersonation_chain = self.async_op_args.get("impersonation_chain", "")
+        self.gcp_project = self.async_op_args.get("project_id", "astronomer-dag-authoring")
+        self.gcp_conn_id = self.profile_config.profile_mapping.conn_id  # type: ignore
+        self.dataset = self.async_op_args.get("dataset", "my_dataset")
+        self.location = self.async_op_args.get("location", "US")
+        self.async_op_args["deferrable"] = True
+        self.reattach_states: set[str] = self.async_op_args.get("reattach_states") or set()
+
+        super().__init__(*args, configuration=self.configuration, task_id=kwargs.get("task_id"), **self.async_op_args)
+
+        if self.profile_type not in _SUPPORTED_DATABASES:
+            raise CosmosValueError(f"Async run are only supported: {_SUPPORTED_DATABASES}")
+
+    def get_remote_sql(self) -> str:
+        if not settings.AIRFLOW_IO_AVAILABLE:
+            raise CosmosValueError(f"Cosmos async support is only available starting in Airflow 2.8 or later.")
+        from airflow.io.path import ObjectStoragePath
+
+        if not self.file_path or not self.project_dir:
+            raise CosmosValueError("file_path and project_dir are required to be set on the task for async execution")
+        project_dir_parent = str(self.project_dir.parent)
+        relative_file_path = str(self.file_path).replace(project_dir_parent, "").lstrip("/")
+        remote_model_path = f"{str(remote_target_path).rstrip('/')}/{self.dag_id}/{relative_file_path}"
+
+        print("remote_model_path: ", remote_model_path)
+        object_storage_path = ObjectStoragePath(remote_model_path, conn_id=remote_target_path_conn_id)
+        with object_storage_path.open() as fp:  # type: ignore
+            return fp.read()  # type: ignore
+
+    def drop_table_sql(self) -> None:
+        model_name = self.task_id.split(".")[0]
+        sql = f"DROP TABLE IF EXISTS {self.gcp_project}.{self.dataset}.{model_name};"
+        hook = BigQueryHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        self.configuration = {
+            "query": {
+                "query": sql,
+                "useLegacySql": False,
+            }
+        }
+        hook.insert_job(configuration=self.configuration, location=self.location, project_id=self.gcp_project)
+
+    def execute(self, context: Context) -> Any | None:
+        if not self.full_refresh:
+            raise CosmosValueError("The async execution only supported for full_refresh")
+        else:
+            self.drop_table_sql()
+
+            sql = self.get_remote_sql()
+            model_name = self.task_id.split(".")[0]
+            # prefix explicit create command to create table
+            sql = f"CREATE TABLE {self.gcp_project}.{self.dataset}.{model_name} AS  {sql}"
+
+            self.configuration = {
+                "query": {
+                    "query": sql,
+                    "useLegacySql": False,
+                }
+            }
+            return super().execute(context)
 
 
 class DbtTestAirflowAsyncOperator(DbtTestLocalOperator):
@@ -47,21 +128,5 @@ class DbtRunOperationAirflowAsyncOperator(DbtRunOperationLocalOperator):
     pass
 
 
-class DbtDocsAirflowAsyncOperator(DbtDocsLocalOperator):
-    pass
-
-
-class DbtDocsS3AirflowAsyncOperator(DbtDocsS3LocalOperator):
-    pass
-
-
-class DbtDocsAzureStorageAirflowAsyncOperator(DbtDocsAzureStorageLocalOperator):
-    pass
-
-
-class DbtDocsGCSAirflowAsyncOperator(DbtDocsGCSLocalOperator):
-    pass
-
-
 class DbtCompileAirflowAsyncOperator(DbtCompileLocalOperator):
     pass

cosmos/operators/base.py

@@ -110,6 +110,7 @@ def __init__(
         dbt_cmd_global_flags: list[str] | None = None,
         cache_dir: Path | None = None,
         extra_context: dict[str, Any] | None = None,
+        # configuration: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         self.project_dir = project_dir
@@ -140,6 +141,7 @@ def __init__(
         self.cache_dir = cache_dir
         self.extra_context = extra_context or {}
         kwargs.pop("full_refresh", None)  # usage of this param should be implemented in child classes
+        # kwargs["configuration"] = {}
         super().__init__(**kwargs)
 
     def get_env(self, context: Context) -> dict[str, str | bytes | os.PathLike[Any]]:

cosmos/operators/local.py

@@ -32,7 +32,7 @@
 from cosmos.dataset import get_dataset_alias_name
 from cosmos.dbt.project import get_partial_parse_path, has_non_empty_dependencies_file
 from cosmos.exceptions import AirflowCompatibilityError, CosmosValueError
-from cosmos.settings import AIRFLOW_IO_AVAILABLE, remote_target_path, remote_target_path_conn_id
+from cosmos.settings import remote_target_path, remote_target_path_conn_id
 
 try:
     from airflow.datasets import Dataset
@@ -143,6 +143,7 @@ def __init__(
         should_store_compiled_sql: bool = True,
         should_upload_compiled_sql: bool = False,
         append_env: bool = True,
+        async_op_args: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         self.task_id = task_id
@@ -157,6 +158,7 @@ def __init__(
         self.invoke_dbt: Callable[..., FullOutputSubprocessResult | dbtRunnerResult]
         self.handle_exception: Callable[..., None]
         self._dbt_runner: dbtRunner | None = None
+        self.async_op_args = async_op_args
         if self.invocation_mode:
             self._set_invocation_methods()
 
@@ -294,7 +296,7 @@ def _configure_remote_target_path() -> tuple[Path, str] | tuple[None, None]:
         if remote_conn_id is None:
             return None, None
 
-        if not AIRFLOW_IO_AVAILABLE:
+        if not settings.AIRFLOW_IO_AVAILABLE:
             raise CosmosValueError(
                 f"You're trying to specify remote target path {target_path_str}, but the required "
                 f"Object Storage feature is unavailable in Airflow version {airflow_version}. Please upgrade to "
@@ -340,6 +342,7 @@ def upload_compiled_sql(self, tmp_project_dir: str, context: Context) -> None:
             return
 
         dest_target_dir, dest_conn_id = self._configure_remote_target_path()
+
         if not dest_target_dir:
             raise CosmosValueError(
                 "You're trying to upload compiled SQL files, but the remote target path is not configured. "
@@ -349,6 +352,7 @@ def upload_compiled_sql(self, tmp_project_dir: str, context: Context) -> None:
 
         source_compiled_dir = Path(tmp_project_dir) / "target" / "compiled"
         files = [str(file) for file in source_compiled_dir.rglob("*") if file.is_file()]
+
         for file_path in files:
             dest_file_path = self._construct_dest_file_path(dest_target_dir, file_path, source_compiled_dir, context)
             dest_object_storage_path = ObjectStoragePath(dest_file_path, conn_id=dest_conn_id)

cosmos/settings.py

@@ -34,7 +34,6 @@
 # This will be merged with the `cache_dir` config parameter in upcoming releases.
 remote_cache_dir = conf.get("cosmos", "remote_cache_dir", fallback=None)
 remote_cache_dir_conn_id = conf.get("cosmos", "remote_cache_dir_conn_id", fallback=None)
-
 remote_target_path = conf.get("cosmos", "remote_target_path", fallback=None)
 remote_target_path_conn_id = conf.get("cosmos", "remote_target_path_conn_id", fallback=None)
 

docs/getting_started/execution-modes.rst

@@ -243,6 +243,7 @@ Each task will create a new Cloud Run Job execution, giving full isolation. The
         },
     )
 
+
 Airflow Async (experimental)
 ----------------------------