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 22 commits into
base: main
Choose a base branch
from
Draft
Changes from 1 commit
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
e657823
Monkeypatch BiqQuery adapter for retriveing SQL for async execution
pankajkoti Jan 21, 2025
8b7b45d
Update cosmos/operators/local.py
pankajkoti Jan 21, 2025
e3ea847
Update cosmos/operators/local.py
pankajkoti Jan 21, 2025
8563a8c
Address @tatiana's review feedback
pankajkoti Jan 23, 2025
94eada9
Refactor run_command method to reduce complexity
pankajkoti Jan 24, 2025
92314e8
Resolve type-check errrors with respect to update method signatures
pankajkoti Jan 24, 2025
859f3ad
Fix tests args
pankajkoti Jan 24, 2025
379d997
Test async dag
pankajkoti Jan 24, 2025
7a85d27
Merge branch 'main' into monkeypatch-bq-adapter
pankajkoti Jan 27, 2025
152b936
🎨 [pre-commit.ci] Auto format from pre-commit.com hooks
pre-commit-ci[bot] Jan 27, 2025
c11f614
Update cosmos/operators/airflow_async.py
pankajkoti Jan 27, 2025
685757d
Update cosmos/operators/airflow_async.py
pankajkoti Jan 27, 2025
d327b6d
Update cosmos/operators/airflow_async.py
pankajkoti Jan 27, 2025
31161bf
Moment of glory
pankajkoti Jan 28, 2025
5ea5217
Moment of glory 2
pankajkoti Jan 29, 2025
dd595f7
🎨 [pre-commit.ci] Auto format from pre-commit.com hooks
pre-commit-ci[bot] Jan 29, 2025
f6e17a5
push the progress
pankajkoti Jan 30, 2025
e4fc114
Merge branch 'main' into monkeypatch-bq-adapter
pankajkoti Jan 30, 2025
93e7a8c
Stop another call to BaseOperator init
pankajkoti Jan 30, 2025
9fc5112
Fix import
pankajkoti Jan 30, 2025
55acacc
Try changing inheritance order to see if MRO helps
pankajkoti Jan 30, 2025
57ed5a8
Remove compile task test
pankajkoti Jan 31, 2025
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
  • Loading branch information
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
Loading
You are viewing a condensed version of this merge commit. You can view the full changes here.