Skip to content

Simplify version-specific imports in the Google provider #56793

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.

Sign up for GitHub

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

Merged
kaxil merged 1 commit into from
Oct 18, 2025
Merged

Simplify version-specific imports in the Google provider #56793

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
4 changes: 4 additions & 0 deletions providers/common/compat/src/airflow/providers/common/compat/lazy_compat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,10 @@
"airflow.utils.python_virtualenv",
),
# ============================================================================
# Timeout Utilities
# ============================================================================
"timeout": ("airflow.sdk.execution_time.timeout", "airflow.utils.timeout"),
# ============================================================================
# XCom & Task Communication
# ============================================================================
"XCOM_RETURN_KEY": "airflow.models.xcom",
Expand Down
2 changes: 2 additions & 0 deletions providers/common/compat/src/airflow/providers/common/compat/lazy_compat.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,7 @@ from airflow.sdk.bases.decorator import (
)
from airflow.sdk.bases.sensor import poke_mode_only as poke_mode_only
from airflow.sdk.definitions.template import literal as literal
from airflow.sdk.execution_time.timeout import timeout as timeout
from airflow.sdk.execution_time.xcom import XCom as XCom

__all__: list[str] = [
Expand Down Expand Up @@ -225,6 +226,7 @@ __all__: list[str] = [
"task",
"task_group",
"teardown",
"timeout",
"timezone",
"write_python_script",
]
2 changes: 1 addition & 1 deletion providers/google/pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ requires-python = ">=3.10"
# After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
dependencies = [
"apache-airflow>=2.10.0",
"apache-airflow-providers-common-compat>=1.4.0",
"apache-airflow-providers-common-compat>=1.7.4", # + TODO: bump to next version
"apache-airflow-providers-common-sql>=1.27.0",
"asgiref>=3.5.2",
"dill>=0.2.3",
Expand Down
2 changes: 1 addition & 1 deletion providers/google/src/airflow/providers/google/ads/hooks/ads.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,8 @@
from google.auth.exceptions import GoogleAuthError

from airflow.exceptions import AirflowException
from airflow.providers.common.compat.lazy_compat import BaseHook
from airflow.providers.google.common.hooks.base_google import get_field
from airflow.providers.google.version_compat import BaseHook

if TYPE_CHECKING:
from google.ads.googleads.v21.services.services.customer_service import CustomerServiceClient
Expand Down
12 changes: 1 addition & 11 deletions providers/google/src/airflow/providers/google/assets/gcs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,23 +18,13 @@

from typing import TYPE_CHECKING

from airflow.providers.common.compat.lazy_compat import Asset
from airflow.providers.google.cloud.hooks.gcs import _parse_gcs_url
from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS

if TYPE_CHECKING:
from urllib.parse import SplitResult

from airflow.providers.common.compat.assets import Asset
from airflow.providers.common.compat.openlineage.facet import Dataset as OpenLineageDataset
else:
try:
from airflow.providers.common.compat.assets import Asset
except ImportError:
if AIRFLOW_V_3_0_PLUS:
from airflow.sdk.definitions.asset import Asset
else:
# dataset is renamed to asset since Airflow 3.0
from airflow.datasets import Dataset as Asset


def create_asset(*, bucket: str, key: str, extra: dict | None = None) -> Asset:
Expand Down
2 changes: 1 addition & 1 deletion providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@
else:
from airflow.models import Connection # type: ignore[assignment,attr-defined,no-redef]

from airflow.providers.common.compat.lazy_compat import BaseHook
from airflow.providers.google.cloud.hooks.secret_manager import (
GoogleCloudSecretManagerHook,
)
Expand All @@ -66,7 +67,6 @@
GoogleBaseHook,
get_field,
)
from airflow.providers.google.version_compat import BaseHook
from airflow.utils.log.logging_mixin import LoggingMixin

if TYPE_CHECKING:
Expand Down
2 changes: 1 addition & 1 deletion providers/google/src/airflow/providers/google/cloud/hooks/dataflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,12 +51,12 @@

from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
from airflow.providers.apache.beam.hooks.beam import BeamHook, BeamRunnerType, beam_options_to_args
from airflow.providers.common.compat.lazy_compat import timeout
from airflow.providers.google.common.hooks.base_google import (
PROVIDE_PROJECT_ID,
GoogleBaseAsyncHook,
GoogleBaseHook,
)
from airflow.providers.google.version_compat import timeout
from airflow.utils.log.logging_mixin import LoggingMixin

if TYPE_CHECKING:
Expand Down
2 changes: 1 addition & 1 deletion providers/google/src/airflow/providers/google/cloud/hooks/dataprep.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
from requests import HTTPError
from tenacity import retry, stop_after_attempt, wait_exponential

from airflow.providers.google.version_compat import BaseHook
from airflow.providers.common.compat.lazy_compat import BaseHook


def _get_field(extras: dict, field_name: str) -> str | None:
Expand Down
2 changes: 1 addition & 1 deletion providers/google/src/airflow/providers/google/cloud/hooks/looker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
from packaging.version import parse as parse_version

from airflow.exceptions import AirflowException
from airflow.providers.google.version_compat import BaseHook
from airflow.providers.common.compat.lazy_compat import BaseHook
from airflow.version import version

if TYPE_CHECKING:
Expand Down
13 changes: 2 additions & 11 deletions providers/google/src/airflow/providers/google/cloud/links/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,17 +20,8 @@
from typing import TYPE_CHECKING, ClassVar
from urllib.parse import urlparse

from airflow.providers.google.version_compat import (
AIRFLOW_V_3_0_PLUS,
BaseOperator,
BaseOperatorLink,
BaseSensorOperator,
)

if AIRFLOW_V_3_0_PLUS:
from airflow.sdk.execution_time.xcom import XCom
else:
from airflow.models.xcom import XCom # type: ignore[no-redef]
from airflow.providers.common.compat.lazy_compat import BaseOperatorLink, BaseSensorOperator, XCom
Copy link
Contributor

@amoghrajesh amoghrajesh Oct 22, 2025

Choose a reason for hiding this comment

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

Can we not import BaseOperator here too?

Copy link
Member Author

@kaxil kaxil Oct 22, 2025

Choose a reason for hiding this comment

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

No, BaseOperator is special cased for 3.1 due to xcom_push bug -- there should be a comment about it in version_compat file

Copy link
Contributor

@amoghrajesh amoghrajesh Oct 23, 2025

Choose a reason for hiding this comment

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

I remember now :)

from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS, BaseOperator

if TYPE_CHECKING:
from airflow.models.taskinstancekey import TaskInstanceKey
Expand Down
12 changes: 2 additions & 10 deletions providers/google/src/airflow/providers/google/cloud/links/dataproc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,22 +25,14 @@
import attr

from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.common.compat.lazy_compat import BaseOperatorLink, XCom
from airflow.providers.google.cloud.links.base import BASE_LINK, BaseGoogleLink
from airflow.providers.google.version_compat import (
AIRFLOW_V_3_0_PLUS,
BaseOperator,
BaseOperatorLink,
)

if TYPE_CHECKING:
from airflow.models.taskinstancekey import TaskInstanceKey
from airflow.providers.google.version_compat import BaseOperator
from airflow.utils.context import Context

if AIRFLOW_V_3_0_PLUS:
from airflow.sdk.execution_time.xcom import XCom
else:
from airflow.models.xcom import XCom # type: ignore[no-redef]


def __getattr__(name: str) -> Any:
# PEP-562: deprecate module-level variable
Expand Down
2 changes: 1 addition & 1 deletion providers/google/src/airflow/providers/google/cloud/operators/cloud_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,14 +28,14 @@

from airflow.configuration import conf
from airflow.exceptions import AirflowException
from airflow.providers.common.compat.lazy_compat import BaseHook
from airflow.providers.google.cloud.hooks.cloud_sql import CloudSQLDatabaseHook, CloudSQLHook
from airflow.providers.google.cloud.links.cloud_sql import CloudSQLInstanceDatabaseLink, CloudSQLInstanceLink
from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
from airflow.providers.google.cloud.triggers.cloud_sql import CloudSQLExportTrigger
from airflow.providers.google.cloud.utils.field_validator import GcpBodyFieldValidator
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, get_field
from airflow.providers.google.common.links.storage import FileDetailsLink
from airflow.providers.google.version_compat import BaseHook

if TYPE_CHECKING:
from airflow.models import Connection
Expand Down
2 changes: 1 addition & 1 deletion providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,12 +26,12 @@

from airflow.configuration import conf
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
from airflow.providers.common.compat.lazy_compat import BaseSensorOperator
from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
from airflow.providers.google.cloud.triggers.bigquery import (
BigQueryTableExistenceTrigger,
BigQueryTablePartitionExistenceTrigger,
)
from airflow.providers.google.version_compat import BaseSensorOperator

if TYPE_CHECKING:
from airflow.utils.context import Context
Expand Down
7 changes: 1 addition & 6 deletions providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,14 +26,9 @@
from google.cloud.bigquery_datatransfer_v1 import TransferState

from airflow.exceptions import AirflowException
from airflow.providers.common.compat.lazy_compat import BaseSensorOperator
from airflow.providers.google.cloud.hooks.bigquery_dts import BiqQueryDataTransferServiceHook
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS

if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import BaseSensorOperator
else:
from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef]

if TYPE_CHECKING:
from google.api_core.retry import Retry
Expand Down
7 changes: 1 addition & 6 deletions providers/google/src/airflow/providers/google/cloud/sensors/bigtable.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,16 +26,11 @@
from google.cloud.bigtable import enums
from google.cloud.bigtable.table import ClusterState

from airflow.providers.common.compat.lazy_compat import BaseSensorOperator
from airflow.providers.google.cloud.hooks.bigtable import BigtableHook
from airflow.providers.google.cloud.links.bigtable import BigtableTablesLink
from airflow.providers.google.cloud.operators.bigtable import BigtableValidationMixin
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS

if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import BaseSensorOperator
else:
from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef]

if TYPE_CHECKING:
from airflow.utils.context import Context
Expand Down
7 changes: 1 addition & 6 deletions providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,15 +31,10 @@

from airflow.configuration import conf
from airflow.exceptions import AirflowException
from airflow.providers.common.compat.lazy_compat import BaseSensorOperator
from airflow.providers.google.cloud.hooks.cloud_composer import CloudComposerHook
from airflow.providers.google.cloud.triggers.cloud_composer import CloudComposerDAGRunTrigger
from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME
from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS

if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import BaseSensorOperator
else:
from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef]
from airflow.utils.state import TaskInstanceState

if TYPE_CHECKING:
Expand Down
7 changes: 1 addition & 6 deletions ...iders/google/src/airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@

from airflow.configuration import conf
from airflow.exceptions import AirflowException
from airflow.providers.common.compat.lazy_compat import BaseSensorOperator
from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
COUNTERS,
METADATA,
Expand All @@ -35,12 +36,6 @@
CloudStorageTransferServiceCheckJobStatusTrigger,
)
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS

if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import BaseSensorOperator
else:
from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef]

if TYPE_CHECKING:
from airflow.utils.context import Context
Expand Down
2 changes: 1 addition & 1 deletion providers/google/src/airflow/providers/google/cloud/sensors/dataflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@

from airflow.configuration import conf
from airflow.exceptions import AirflowException
from airflow.providers.common.compat.lazy_compat import BaseSensorOperator, PokeReturnValue
from airflow.providers.google.cloud.hooks.dataflow import (
DEFAULT_DATAFLOW_LOCATION,
DataflowHook,
Expand All @@ -37,7 +38,6 @@
DataflowJobStatusTrigger,
)
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
from airflow.providers.google.version_compat import BaseSensorOperator, PokeReturnValue

if TYPE_CHECKING:
from airflow.utils.context import Context
Expand Down
7 changes: 1 addition & 6 deletions providers/google/src/airflow/providers/google/cloud/sensors/dataform.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,13 +23,8 @@
from typing import TYPE_CHECKING

from airflow.exceptions import AirflowException
from airflow.providers.common.compat.lazy_compat import BaseSensorOperator
from airflow.providers.google.cloud.hooks.dataform import DataformHook
from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS

if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import BaseSensorOperator
else:
from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef]

if TYPE_CHECKING:
from airflow.utils.context import Context
Expand Down
7 changes: 1 addition & 6 deletions providers/google/src/airflow/providers/google/cloud/sensors/datafusion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,9 @@
from typing import TYPE_CHECKING

from airflow.exceptions import AirflowException, AirflowNotFoundException
from airflow.providers.common.compat.lazy_compat import BaseSensorOperator
from airflow.providers.google.cloud.hooks.datafusion import DataFusionHook
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS

if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import BaseSensorOperator
else:
from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef]

if TYPE_CHECKING:
from airflow.utils.context import Context
Expand Down
7 changes: 1 addition & 6 deletions providers/google/src/airflow/providers/google/cloud/sensors/dataplex.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,17 +32,12 @@
from google.cloud.dataplex_v1.types import DataScanJob

from airflow.exceptions import AirflowException
from airflow.providers.common.compat.lazy_compat import BaseSensorOperator
from airflow.providers.google.cloud.hooks.dataplex import (
AirflowDataQualityScanException,
AirflowDataQualityScanResultTimeoutException,
DataplexHook,
)
from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS

if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import BaseSensorOperator
else:
from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef]


class TaskState:
Expand Down
7 changes: 1 addition & 6 deletions providers/google/src/airflow/providers/google/cloud/sensors/dataprep.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,13 +22,8 @@
from collections.abc import Sequence
from typing import TYPE_CHECKING

from airflow.providers.common.compat.lazy_compat import BaseSensorOperator
from airflow.providers.google.cloud.hooks.dataprep import GoogleDataprepHook, JobGroupStatuses
from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS

if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import BaseSensorOperator
else:
from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef]

if TYPE_CHECKING:
from airflow.utils.context import Context
Expand Down
7 changes: 1 addition & 6 deletions providers/google/src/airflow/providers/google/cloud/sensors/dataproc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,14 +27,9 @@
from google.cloud.dataproc_v1.types import Batch, JobStatus

from airflow.exceptions import AirflowException
from airflow.providers.common.compat.lazy_compat import BaseSensorOperator
from airflow.providers.google.cloud.hooks.dataproc import DataprocHook
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS

if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import BaseSensorOperator
else:
from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef]

if TYPE_CHECKING:
from airflow.utils.context import Context
Expand Down
7 changes: 1 addition & 6 deletions providers/google/src/airflow/providers/google/cloud/sensors/dataproc_metastore.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,14 +21,9 @@
from typing import TYPE_CHECKING

from airflow.exceptions import AirflowException
from airflow.providers.common.compat.lazy_compat import BaseSensorOperator
from airflow.providers.google.cloud.hooks.dataproc_metastore import DataprocMetastoreHook
from airflow.providers.google.cloud.hooks.gcs import parse_json_from_gcs
from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS

if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import BaseSensorOperator
else:
from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef]

if TYPE_CHECKING:
from google.api_core.operation import Operation
Expand Down
Loading