Monkeypatch BiqQuery adapter to retrive SQL for async execution #1474

pankajkoti · 2025-01-21T08:23:54Z

closes: #1260
related: #1261
related: #1266
closes: #1265

netlify · 2025-01-21T08:24:16Z

✅ Deploy Preview for sunny-pastelito-5ecb04 canceled.

Name	Link
🔨 Latest commit	`379d997`
🔍 Latest deploy log	https://app.netlify.com/sites/sunny-pastelito-5ecb04/deploys/67937281c7956200082519af

cosmos/operators/local.py

cloudflare-workers-and-pages · 2025-01-21T08:59:39Z

Deploying astronomer-cosmos with Cloudflare Pages

Latest commit:	`379d997`
Status:	✅ Deploy successful!
Preview URL:	https://c520c11a.astronomer-cosmos.pages.dev
Branch Preview URL:	https://monkeypatch-bq-adapter.astronomer-cosmos.pages.dev

View logs

tatiana · 2025-01-21T11:32:13Z

cosmos/airflow/graph.py

@@ -371,32 +370,6 @@ def generate_task_or_group(
    return task_or_group


-def _add_dbt_compile_task(


As we agreed, I've added a task to follow up on this:
#1477

Ideally, we'd compare both approaches before 1.9, but I'm happy with us to go ahead with approach of #1474 in 1.9 - it is already a significant improvement when compared to 1.7.0, since it solves the existing SQL query issue.

tatiana · 2025-01-21T11:33:47Z

cosmos/mocked_dbt_adapters.py

+    ) -> Tuple[BigQueryAdapterResponse, agate.Table]:
+        return BigQueryAdapterResponse("mock_bigquery_adapter_response"), empty_table()
+
+    BigQueryConnectionManager.execute = execute


We may have to mock a different method when using dbt build. That said, this is outside of the scope of fixing the original but this PR solves, so I've logged a follow up ticket:
#1476

tatiana · 2025-01-21T11:35:44Z

cosmos/operators/airflow_async.py

-        "gcp_project",
-        "dataset",


Why did we remove these as templated fields?

tatiana · 2025-01-21T11:36:36Z

cosmos/operators/airflow_async.py

@@ -60,13 +55,11 @@ class DbtSourceAirflowAsyncOperator(DbtBaseAirflowAsyncOperator, DbtSourceLocalO
    pass


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


It would be great if we can see the SQL statement in the operator Templated Field, like we can do with the Cosmos Local operators. If this is too much work, I'm happy for it to be done in a follow up PR/ticket.

tatiana · 2025-01-21T11:46:43Z

cosmos/operators/local.py

+        return_sql: bool = False,
+        sql_context: dict[str, Any] | None = None,
+    ) -> FullOutputSubprocessResult | dbtRunnerResult | str:


I was confused with this method, given what the LocalOperatorrun_command does:

astronomer-cosmos/cosmos/operators/local.py

Lines 422 to 427 in 74c7c7d

def run_command(

self,

cmd: list[str],

env: dict[str, str | bytes | os.PathLike[Any]],

context: Context,

) -> FullOutputSubprocessResult | dbtRunnerResult:

The LocalOperator currently:

runs the dbt command

runs the actual transformation in the data warehouse

I suggest one of the following is done:
a) have both methods accomplishing the same goal (not only running the dbt command, but actually executing it in the data warehouse) or
b) name this method something else

If we decide to do approach (a), I believe we should:

stick to the existing method interface, exposing it in the cosmos BaseOperator class

change this method in this PR to also execute the SQL statement

tatiana · 2025-01-21T11:48:30Z

cosmos/operators/local.py

@@ -647,11 +669,15 @@ def get_openlineage_facets_on_complete(self, task_instance: TaskInstance) -> Ope
        )

    def build_and_run_cmd(
-        self, context: Context, cmd_flags: list[str] | None = None
+        self,


Similar feedback to #1474 (comment) here.

Concerns regarding:

inconsistent interface between different ExecutionMode implementations

expectation that this method would not only build the command with necessary flags, but actually have executed the dbt transformation by the end of its run

tatiana

@pankajkoti I'm very excited that we now have a more reliable way of calculating the full dbt SQL query. This approach fixes #1260 and solves many of the async tickets we have open.

Monkey-patching always carries a risk, but it is worth it at this stage.

It would be great if - either as part of this PR - or as a priority follow-up PR, we have an efficient way of testing that the monkey patching works in multiple versions of dbt, including the latest releases, and that the transformation is not being executed when we run the dbt command. I believe this must be done before we release this feature in 1.9.0

I've logged two follow-up tickets that are relevant:

one is to consider the re-introduction of the compile task, if that means we can avoid having dbt installed in all dbt worker nodes, and executing the command in most Cosmos tasks: Re-evaluate adding compile task when using ExecutionMode.AIRFLOW_ASYNC #1477
the other is to support TestBehavior.BUILD Support TestBehavior.BUILD when using ExecutionMode.AIRFLOW_ASYNC #1476

It would be great if these could be accomplished before 1.9.0 release, but I'm also happy with us sticking to approach if time does not allow further analysis / implementation.

pankajkoti · 2025-01-22T05:24:29Z

cc: @joppevos for visibility on the ongoing work

pankajkoti temporarily deployed to internal January 21, 2025 08:24 — with GitHub Actions Inactive

pankajkoti requested review from tatiana and pankajastro January 21, 2025 08:25

pankajkoti changed the title ~~Monkeypatch BiqQuery adapter for retriveing SQL for async execution~~ Monkeypatch BiqQuery adapter to retrive SQL for async execution Jan 21, 2025

pankajkoti commented Jan 21, 2025

View reviewed changes

cosmos/operators/local.py Outdated Show resolved Hide resolved

cosmos/operators/local.py Outdated Show resolved Hide resolved

pankajkoti temporarily deployed to internal January 21, 2025 08:59 — with GitHub Actions Inactive

tatiana added this to the Cosmos 1.9.0 milestone Jan 21, 2025

This was referenced Jan 21, 2025

Support TestBehavior.BUILD when using ExecutionMode.AIRFLOW_ASYNC #1476

Open

Re-evaluate adding compile task when using ExecutionMode.AIRFLOW_ASYNC #1477

Open

tatiana reviewed Jan 21, 2025

View reviewed changes

pankajkoti force-pushed the monkeypatch-bq-adapter branch from 9590f54 to e521051 Compare January 21, 2025 11:43

pankajkoti temporarily deployed to internal January 21, 2025 11:43 — with GitHub Actions Inactive

tatiana reviewed Jan 21, 2025

View reviewed changes

pankajastro mentioned this pull request Jan 23, 2025

Add structure to support multiple db for async operator execution #1483

Open

pankajkoti temporarily deployed to internal January 23, 2025 10:25 — with GitHub Actions Inactive

pankajkoti force-pushed the monkeypatch-bq-adapter branch from 8ef4ac5 to bd85529 Compare January 23, 2025 11:08

pankajkoti temporarily deployed to internal January 23, 2025 11:08 — with GitHub Actions Inactive

pankajkoti temporarily deployed to internal January 24, 2025 08:23 — with GitHub Actions Inactive

pankajkoti added 4 commits January 24, 2025 13:55

Monkeypatch BiqQuery adapter for retriveing SQL for async execution

e657823

Update cosmos/operators/local.py

8b7b45d

Update cosmos/operators/local.py

e3ea847

Address @tatiana's review feedback

8563a8c

Refactor run_command method to reduce complexity

94eada9

pankajkoti force-pushed the monkeypatch-bq-adapter branch from c161db2 to 94eada9 Compare January 24, 2025 08:25

pankajkoti temporarily deployed to internal January 24, 2025 08:25 — with GitHub Actions Inactive

Resolve type-check errrors with respect to update method signatures

92314e8

pankajkoti temporarily deployed to internal January 24, 2025 08:50 — with GitHub Actions Inactive

Fix tests args

859f3ad

pankajkoti temporarily deployed to internal January 24, 2025 09:00 — with GitHub Actions Inactive

pankajkoti temporarily deployed to internal January 24, 2025 10:49 — with GitHub Actions Inactive

Test async dag

379d997

pankajkoti force-pushed the monkeypatch-bq-adapter branch from 1a413d6 to 379d997 Compare January 24, 2025 10:59

pankajkoti temporarily deployed to internal January 24, 2025 10:59 — with GitHub Actions Inactive

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Monkeypatch BiqQuery adapter to retrive SQL for async execution #1474

Monkeypatch BiqQuery adapter to retrive SQL for async execution #1474

pankajkoti commented Jan 21, 2025 •

edited

Loading

netlify bot commented Jan 21, 2025 •

edited

Loading

cloudflare-workers-and-pages bot commented Jan 21, 2025 •

edited

Loading

tatiana Jan 21, 2025

tatiana Jan 21, 2025

tatiana Jan 21, 2025

tatiana Jan 21, 2025

tatiana Jan 21, 2025

tatiana Jan 21, 2025

tatiana left a comment

pankajkoti commented Jan 22, 2025

		@@ -371,32 +370,6 @@ def generate_task_or_group(
		return task_or_group


		def _add_dbt_compile_task(

	def run_command(
	self,
	cmd: list[str],
	env: dict[str, str \| bytes \| os.PathLike[Any]],
	context: Context,
	) -> FullOutputSubprocessResult \| dbtRunnerResult:

Monkeypatch BiqQuery adapter to retrive SQL for async execution #1474

Are you sure you want to change the base?

Monkeypatch BiqQuery adapter to retrive SQL for async execution #1474

Conversation

pankajkoti commented Jan 21, 2025 • edited Loading

netlify bot commented Jan 21, 2025 • edited Loading

✅ Deploy Preview for sunny-pastelito-5ecb04 canceled.

cloudflare-workers-and-pages bot commented Jan 21, 2025 • edited Loading

Deploying astronomer-cosmos with Cloudflare Pages

tatiana Jan 21, 2025

Choose a reason for hiding this comment

tatiana Jan 21, 2025

Choose a reason for hiding this comment

tatiana Jan 21, 2025

Choose a reason for hiding this comment

tatiana Jan 21, 2025

Choose a reason for hiding this comment

tatiana Jan 21, 2025

Choose a reason for hiding this comment

tatiana Jan 21, 2025

Choose a reason for hiding this comment

tatiana left a comment

Choose a reason for hiding this comment

pankajkoti commented Jan 22, 2025

pankajkoti commented Jan 21, 2025 •

edited

Loading

netlify bot commented Jan 21, 2025 •

edited

Loading

cloudflare-workers-and-pages bot commented Jan 21, 2025 •

edited

Loading