feat: Add `GCP_GKE` and `WATCHER_GCP_GKE` execution modes by vricciardulli · Pull Request #2488 · astronomer/astronomer-cosmos

vricciardulli · 2026-03-21T10:15:22Z

Description

Added GCP_GKE and WATCHER_GCP_GKE execution modes.

It's the same as KUBERNETES and WATCHER_KUBERNETES modes, but uses GKE operators (i.e. GKEStartPodOperator).

This enables a user to use cosmos in Cloud Composer 3 with a self-managed custom GKE cluster.

Tests:

Added unit tests
Tested in GCP Cloud Composer 3 with a custom GKE cluster. It works.
Did not add any integration test.

Implementation details:

Given the amount of duplication between the k8s modes and the new GKE ones, I've extracted all common code and placed it in helper functions in cosmos/operators/_k8s_common.py.
I've attempted to type hint as best as possible, but I had to ignore some type-checker errors. I've marked those with comments.
An alternative approach would have been to use a factory to generate the classes based on the pod operator provided, but I think that this approach would have been too opaque.

Docs:

I did not add separate pages documenting these modes.
I concisely added details of GKE modes in the k8s and watcher k8s pages.
Also added to the mermaid graph.

Co-Authored-By: Claude Opus 4.6 (1M context) noreply@anthropic.com

Related Issue(s)

closes #2487
closes #2379

Breaking Change?

There should not be a breaking change.

Checklist

I have made corresponding changes to the documentation (if required)
I have added tests that prove my fix is effective or that my feature works

Add two new execution modes enabling dbt execution on GKE clusters: - GCP_GKE: standard mode using GKEStartPodOperator - WATCHER_GCP_GKE: watcher mode with producer/consumer pattern Includes all operators, watcher callbacks, graph.py watcher mappings, and comprehensive unit/integration tests. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Copilot

Pull request overview

Adds new execution modes so Cosmos can run dbt tasks on self-managed GKE clusters (e.g., Cloud Composer 3 on custom GKE) by using GKEStartPodOperator and a watcher variant.

Changes:

Introduce ExecutionMode.GCP_GKE and ExecutionMode.WATCHER_GCP_GKE.
Add GKE-backed dbt operators (cosmos/operators/gcp_gke.py) and watcher counterparts (cosmos/operators/watcher_gcp_gke.py).
Wire new watcher mode into graph-building logic and add unit/integration tests.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
`cosmos/constants.py`	Adds the two new execution modes.
`cosmos/airflow/graph.py`	Ensures watcher graph building and AFTER_ALL test behavior work with `WATCHER_GCP_GKE`.
`cosmos/operators/gcp_gke.py`	Implements dbt operators backed by `GKEStartPodOperator`.
`cosmos/operators/watcher_gcp_gke.py`	Implements watcher producer/consumer behavior for GKE.
`tests/operators/test_gcp_gke.py`	Unit tests for GCP GKE operator command construction and execute path.
`tests/operators/test_watcher_gcp_gke_unit.py`	Unit tests for watcher GCP GKE retry/callback behavior.
`tests/operators/test_watcher_gcp_gke_integration.py`	Integration test for `DbtDag` with `WATCHER_GCP_GKE`.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

vricciardulli · 2026-03-29T15:56:21Z

+
+        # Get the logs from the pod
+        logs = []
+        for log in task.pod_manager.read_pod_logs(pod, "base"):  # type: ignore[attr-defined]


Type checker error:

error: "BaseOperator" has no attribute "pod_manager" [attr-defined]

was not getting this error before the refactoring

Copilot

Pull request overview

Copilot reviewed 18 out of 18 changed files in this pull request and generated 6 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 3 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-01T16:59:30Z

 .. versionadded:: 1.13.0

 The ``ExecutionMode.WATCHER_KUBERNETES`` combines the **speed of the** :ref:`watcher-execution-mode` **with the isolation of** :ref:`kubernetes`.

+A GCP GKE variant is also available as ``ExecutionMode.WATCHER_GCP_GKE``, which uses
+``GKEStartPodOperator`` instead of ``KubernetesPodOperator``. See :ref:`watcher-gcp-gke` below.
+


The page-level .. versionadded:: 1.13.0 now applies to the newly documented ExecutionMode.WATCHER_GCP_GKE as well, but this mode is introduced in this PR (later than 1.13.0). Consider adding a separate .. versionadded:: <new version> note for WATCHER_GCP_GKE (or rewording the existing directive) to avoid implying it has been available since 1.13.0.

tatiana · 2026-04-07T14:58:57Z

Thank you so much for these improvements, @vricciardulli ! They look really valuable and will definitely enhance the project.

We’ve had a heavier workload than expected over the past few weeks, so we haven’t been able to review everything immediately. We’ll be reviewing this PR thoroughly in the upcoming month and are considering including these changes in the upcoming 1.15.0 release.

Really appreciate your patience and contribution!

In the meantime, please, could you address Copilot's feedback and also rebase with the latest changes on main?

codecov · 2026-04-07T15:50:50Z

Codecov Report

❌ Patch coverage is 95.45455% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 97.89%. Comparing base (0ba7fde) to head (eb9c56e).
⚠️ Report is 59 commits behind head on main.

Files with missing lines	Patch %	Lines
cosmos/operators/watcher_gcp_gke.py	88.67%	6 Missing ⚠️
cosmos/operators/gcp_gke.py	91.52%	5 Missing ⚠️
cosmos/operators/_k8s_common.py	98.35%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2488      +/-   ##
==========================================
- Coverage   97.92%   97.89%   -0.03%     
==========================================
  Files         103      106       +3     
  Lines        7312     7460     +148     
==========================================
+ Hits         7160     7303     +143     
- Misses        152      157       +5

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

vricciardulli · 2026-04-10T14:56:50Z

Thank you so much for these improvements, @vricciardulli ! They look really valuable and will definitely enhance the project.

We’ve had a heavier workload than expected over the past few weeks, so we haven’t been able to review everything immediately. We’ll be reviewing this PR thoroughly in the upcoming month and are considering including these changes in the upcoming 1.15.0 release.

Really appreciate your patience and contribution!

In the meantime, please, could you address Copilot's feedback and also rebase with the latest changes on main?

@tatiana I'll rebase and clean up once #2543 is merged to main

Copilot

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

cosmos/operators/kubernetes.py:253

DbtDocsCloudKubernetesOperator.build_and_run_cmd builds the shell command from self.arguments only. After the refactor that splits the executable into self.cmds + self.arguments, this drops the dbt executable entirely (e.g., produces docs generate ...), which will fail at runtime. Build the command string from self.cmds + self.arguments (or otherwise ensure dbt is included) before overriding self.cmds/self.arguments for the bash wrapper.

vricciardulli · 2026-05-12T23:03:50Z

+from cosmos.log import get_logger
+from cosmos.operators._watcher.xcom import (
+    _backup_xcom_to_variable,
+    _delete_xcom_backup_variable,
+    _init_xcom_backup,
+    _restore_xcom_from_variable,
+)
+
+if TYPE_CHECKING:  # pragma: no cover
+    from pendulum import DateTime
+
+if TYPE_CHECKING:  # pragma: no cover
+    try:
+        from airflow.sdk.definitions.context import Context
+    except ImportError:
+        from airflow.utils.context import Context  # type: ignore[attr-defined]
+try:
+    # apache-airflow-providers-cncf-kubernetes >= 7.14.0
+    from airflow.providers.cncf.kubernetes.callbacks import KubernetesPodOperatorCallback
+except ImportError:
+
+    class KubernetesPodOperatorCallback:  # type: ignore[no-redef]
+        """Mock fallback for older versions. Should not be used in practice."""
+
+        pass
+
+
+from cosmos.dbt.parser.output import extract_log_issues
+from cosmos.operators._watcher.base import store_dbt_resource_status_from_log
+from cosmos.operators.base import AbstractDbtBase
+
+try:
+    from airflow.sdk.bases.operator import BaseOperator  # Airflow 3
+except ImportError:
+    from airflow.models import BaseOperator  # Airflow 2
+
+logger = get_logger(__name__)


Kept because it was here also before the refactoring

vricciardulli and others added 3 commits March 19, 2026 11:02

fix: try different cmd invocation

c40c7db

Merge remote-tracking branch 'upstream/main' into gcp-modes

c73aae5

vricciardulli requested review from a team, corsettigyg, dwreeves and jbandoro as code owners March 21, 2026 10:15

vricciardulli requested review from Copilot, pankajkoti and tatiana March 21, 2026 10:15

vricciardulli had a problem deploying to external March 21, 2026 10:15 — with GitHub Actions Error

🎨 [pre-commit.ci] Auto format from pre-commit.com hooks

a947590

Copilot started reviewing on behalf of vricciardulli March 21, 2026 10:15 View session

pre-commit-ci Bot had a problem deploying to external March 21, 2026 10:15 Error

Copilot AI reviewed Mar 21, 2026

View reviewed changes

Comment thread tests/operators/test_gcp_gke.py

Comment thread cosmos/operators/gcp_gke.py Outdated

Comment thread cosmos/operators/watcher_gcp_gke.py Outdated

fix: bots comments

8416d7e

vricciardulli had a problem deploying to external March 21, 2026 10:26 — with GitHub Actions Error

fix: tests

d046ae2

Copilot AI review requested due to automatic review settings March 21, 2026 10:41

vricciardulli had a problem deploying to external March 21, 2026 10:41 — with GitHub Actions Error

Copilot started reviewing on behalf of vricciardulli March 21, 2026 10:41 View session

Copilot AI reviewed Mar 21, 2026

View reviewed changes

Comment thread tests/operators/test_watcher_gcp_gke_integration.py Outdated

Comment thread cosmos/operators/gcp_gke.py Outdated

fix: bots comments

885e2e2

vricciardulli had a problem deploying to external March 21, 2026 10:53 — with GitHub Actions Error

vricciardulli commented Mar 21, 2026

View reviewed changes

Comment thread cosmos/operators/watcher_gcp_gke.py Outdated

Merge branch 'main' into gcp-modes

2e64343

Copilot AI review requested due to automatic review settings March 21, 2026 12:47

vricciardulli had a problem deploying to external March 21, 2026 12:47 — with GitHub Actions Error

Merge remote-tracking branch 'upstream/main' into gcp-modes

ce703e9

Copilot started reviewing on behalf of vricciardulli March 21, 2026 12:47 View session

vricciardulli commented Mar 29, 2026

View reviewed changes

Copilot AI reviewed Mar 29, 2026

View reviewed changes

vricciardulli added 2 commits March 29, 2026 18:04

docs: fix pip install

8ec48cc

fix: remove integration tests

446dbf8

vricciardulli temporarily deployed to external March 29, 2026 16:05 — with GitHub Actions Inactive

vricciardulli added 3 commits March 30, 2026 21:08

Merge remote-tracking branch 'upstream/main' into gcp-modes

2b458a5

Merge remote-tracking branch 'upstream/main' into gcp-modes

d41cea2

fix: bring back inits

eb9c56e

Copilot AI review requested due to automatic review settings April 1, 2026 16:52

vricciardulli temporarily deployed to external April 1, 2026 16:52 — with GitHub Actions Inactive

Copilot started reviewing on behalf of vricciardulli April 1, 2026 16:52 View session

Copilot AI reviewed Apr 1, 2026

View reviewed changes

Merge remote-tracking branch 'upstream/main' into gcp-modes

6139d46

vricciardulli mentioned this pull request Apr 7, 2026

[Feature] Support emitting Airflow assets/datasets with ExecutionMode.KUBERNETES #2329

Open

1 task

Merge remote-tracking branch 'upstream/main' into gcp-modes

7a4ac8c

vricciardulli had a problem deploying to external May 12, 2026 22:34 — with GitHub Actions Error

🎨 [pre-commit.ci] Auto format from pre-commit.com hooks

0d3f2cf

Copilot AI review requested due to automatic review settings May 12, 2026 22:34

vricciardulli review requested due to automatic review settings May 12, 2026 22:34

pre-commit-ci Bot had a problem deploying to external May 12, 2026 22:35 Error

vricciardulli added 2 commits May 13, 2026 00:54

fix: enforce pre-commit checks

737c036

fix: simplify

0ba5e9b

Copilot AI review requested due to automatic review settings May 12, 2026 22:55

vricciardulli temporarily deployed to external May 12, 2026 22:55 — with GitHub Actions Inactive

Copilot started reviewing on behalf of vricciardulli May 12, 2026 22:56 View session

Copilot AI reviewed May 12, 2026

View reviewed changes

tatiana assigned pankajkoti May 27, 2026

Conversation

vricciardulli commented Mar 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue(s)

Breaking Change?

Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vricciardulli Mar 29, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Copilot AI Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

tatiana commented Apr 7, 2026

Uh oh!

codecov Bot commented Apr 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

vricciardulli commented Apr 10, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

vricciardulli May 12, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

vricciardulli commented Mar 21, 2026 •

edited

Loading

codecov Bot commented Apr 7, 2026 •

edited

Loading