Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Monkeypatch BiqQuery adapter to retrive SQL for async execution #1474

Draft
wants to merge 21 commits into
base: main
Choose a base branch
from
Draft
Changes from 1 commit
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Merge branch 'main' into monkeypatch-bq-adapter
pankajkoti authored Jan 27, 2025

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
commit 7a85d271ae4134c1c199fc367ce7d449eade0872
66 changes: 12 additions & 54 deletions cosmos/operators/airflow_async.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,9 @@
from __future__ import annotations

import inspect
from typing import Any, Sequence

from airflow.providers.google.cloud.operators.bigquery import BigQueryInsertJobOperator
from airflow.utils.context import Context

from cosmos.config import ProfileConfig
from cosmos.constants import BIGQUERY_PROFILE_TYPE
from cosmos.exceptions import CosmosValueError
from cosmos.operators._asynchronous.base import DbtRunAirflowAsyncFactoryOperator
from cosmos.operators.base import AbstractDbtBaseOperator
from cosmos.operators.local import (
DbtBuildLocalOperator,
@@ -58,15 +53,7 @@ class DbtSourceAirflowAsyncOperator(DbtBaseAirflowAsyncOperator, DbtSourceLocalO
pass


class DbtRunAirflowAsyncOperator(BigQueryInsertJobOperator, DbtRunLocalOperator): # type: ignore

template_fields: Sequence[str] = DbtRunLocalOperator.template_fields + ( # type: ignore[operator]
"full_refresh",
"project_dir",
"gcp_project",
"dataset",
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why did we remove these as templated fields?

"location",
)
class DbtRunAirflowAsyncOperator(DbtRunAirflowAsyncFactoryOperator): # type: ignore

def __init__( # type: ignore
self,
@@ -75,52 +62,23 @@ def __init__( # type: ignore
extra_context: dict[str, object] | None = None,
**kwargs,
) -> None:
pankajkoti marked this conversation as resolved.
Show resolved Hide resolved
# dbt task param
self.project_dir = project_dir
self.full_refresh = full_refresh
self.profile_config = profile_config
if not self.profile_config or not self.profile_config.profile_mapping:
raise CosmosValueError(f"Cosmos async support is only available when using ProfileMapping")

# Cosmos attempts to pass many kwargs that async operator simply does not accept.
# We need to pop them.
async_op_kwargs = {}
cosmos_op_kwargs = {}
clean_kwargs = {}
non_async_args = set(inspect.signature(AbstractDbtBaseOperator.__init__).parameters.keys())
non_async_args |= set(inspect.signature(DbtLocalBaseOperator.__init__).parameters.keys())

non_async_args -= {"task_id"}
pankajkoti marked this conversation as resolved.
Show resolved Hide resolved
for arg_key, arg_value in kwargs.items():
if arg_key == "task_id":
async_op_kwargs[arg_key] = arg_value
cosmos_op_kwargs[arg_key] = arg_value
elif arg_key not in non_async_args:
async_op_kwargs[arg_key] = arg_value
else:
cosmos_op_kwargs[arg_key] = arg_value

# The following are the minimum required parameters to run BigQueryInsertJobOperator using the deferrable mode
BigQueryInsertJobOperator.__init__(
self,
gcp_conn_id=self.gcp_conn_id,
configuration=self.configuration,
location=self.location,
deferrable=True,
**async_op_kwargs,
if arg_key not in non_async_args:
clean_kwargs[arg_key] = arg_value

super().__init__(
project_dir=project_dir,
profile_config=profile_config,
extra_context=extra_context,
**clean_kwargs,
)

DbtRunLocalOperator.__init__(
self,
project_dir=self.project_dir,
profile_config=self.profile_config,
**cosmos_op_kwargs,
)
self.async_context = extra_context or {}
self.async_context["profile_type"] = self.profile_type
self.async_context["async_operator"] = BigQueryInsertJobOperator

def execute(self, context: Context) -> Any | None:
return self.build_and_run_cmd(context, run_as_async=True, async_context=self.async_context)


class DbtTestAirflowAsyncOperator(DbtBaseAirflowAsyncOperator, DbtTestLocalOperator): # type: ignore
pass
You are viewing a condensed version of this merge commit. You can view the full changes here.