diff --git a/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py b/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py index 8810ff454e04b..f7bf7fc3b8c3b 100644 --- a/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py +++ b/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py @@ -35,7 +35,7 @@ from packaging.version import Version from airflow.exceptions import AirflowConfigException, AirflowException -from airflow.providers.apache.beam.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook from airflow.providers.common.compat.standard.utils import prepare_virtualenv if TYPE_CHECKING: diff --git a/providers/apache/beam/src/airflow/providers/apache/beam/operators/beam.py b/providers/apache/beam/src/airflow/providers/apache/beam/operators/beam.py index 7713bbf067b1c..da189d66c4fec 100644 --- a/providers/apache/beam/src/airflow/providers/apache/beam/operators/beam.py +++ b/providers/apache/beam/src/airflow/providers/apache/beam/operators/beam.py @@ -34,14 +34,13 @@ from airflow.exceptions import AirflowException, AirflowOptionalProviderFeatureException from airflow.providers.apache.beam.hooks.beam import BeamHook, BeamRunnerType from airflow.providers.apache.beam.triggers.beam import BeamJavaPipelineTrigger, BeamPythonPipelineTrigger -from airflow.providers.apache.beam.version_compat import BaseOperator +from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers_manager import ProvidersManager from airflow.utils.helpers import convert_camel_to_snake, exactly_one from airflow.version import version if TYPE_CHECKING: - from airflow.utils.context import Context - + from airflow.providers.common.compat.sdk import Context GOOGLE_PROVIDER = ProvidersManager().providers.get("apache-airflow-providers-google") diff --git a/providers/apache/beam/src/airflow/providers/apache/beam/version_compat.py b/providers/apache/beam/src/airflow/providers/apache/beam/version_compat.py index 9f69b38930166..68321a4035b55 100644 --- a/providers/apache/beam/src/airflow/providers/apache/beam/version_compat.py +++ b/providers/apache/beam/src/airflow/providers/apache/beam/version_compat.py @@ -31,14 +31,6 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0) -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook, BaseOperator -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - from airflow.models import BaseOperator - __all__ = [ "AIRFLOW_V_3_1_PLUS", - "BaseHook", - "BaseOperator", ] diff --git a/providers/apache/cassandra/pyproject.toml b/providers/apache/cassandra/pyproject.toml index 0db5d1be56116..46627f02033fc 100644 --- a/providers/apache/cassandra/pyproject.toml +++ b/providers/apache/cassandra/pyproject.toml @@ -61,11 +61,19 @@ dependencies = [ "cassandra-driver>=3.29.1", ] +# The optional dependencies should be modified in place in the generated file +# Any change in the dependencies is preserved when the file is regenerated +[project.optional-dependencies] +"common.compat" = [ + "apache-airflow-providers-common-compat" +] + [dependency-groups] dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", # Additional devel dependencies (do not remove this line and add extra development dependencies) ] diff --git a/providers/apache/cassandra/src/airflow/providers/apache/cassandra/hooks/cassandra.py b/providers/apache/cassandra/src/airflow/providers/apache/cassandra/hooks/cassandra.py index 821e0678213ae..f29a1c66c54a5 100644 --- a/providers/apache/cassandra/src/airflow/providers/apache/cassandra/hooks/cassandra.py +++ b/providers/apache/cassandra/src/airflow/providers/apache/cassandra/hooks/cassandra.py @@ -31,7 +31,7 @@ WhiteListRoundRobinPolicy, ) -from airflow.providers.apache.cassandra.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook from airflow.utils.log.logging_mixin import LoggingMixin Policy: TypeAlias = DCAwareRoundRobinPolicy | RoundRobinPolicy | TokenAwarePolicy | WhiteListRoundRobinPolicy diff --git a/providers/apache/cassandra/src/airflow/providers/apache/cassandra/sensors/record.py b/providers/apache/cassandra/src/airflow/providers/apache/cassandra/sensors/record.py index cb46283fb2631..6147b33544540 100644 --- a/providers/apache/cassandra/src/airflow/providers/apache/cassandra/sensors/record.py +++ b/providers/apache/cassandra/src/airflow/providers/apache/cassandra/sensors/record.py @@ -21,10 +21,10 @@ from typing import TYPE_CHECKING, Any from airflow.providers.apache.cassandra.hooks.cassandra import CassandraHook -from airflow.providers.apache.cassandra.version_compat import BaseSensorOperator +from airflow.providers.common.compat.sdk import BaseSensorOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CassandraRecordSensor(BaseSensorOperator): diff --git a/providers/apache/cassandra/src/airflow/providers/apache/cassandra/sensors/table.py b/providers/apache/cassandra/src/airflow/providers/apache/cassandra/sensors/table.py index 89bd81def30dd..8792c8dc03b88 100644 --- a/providers/apache/cassandra/src/airflow/providers/apache/cassandra/sensors/table.py +++ b/providers/apache/cassandra/src/airflow/providers/apache/cassandra/sensors/table.py @@ -21,10 +21,10 @@ from typing import TYPE_CHECKING, Any from airflow.providers.apache.cassandra.hooks.cassandra import CassandraHook -from airflow.providers.apache.cassandra.version_compat import BaseSensorOperator +from airflow.providers.common.compat.sdk import BaseSensorOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CassandraTableSensor(BaseSensorOperator): diff --git a/providers/apache/cassandra/src/airflow/providers/apache/cassandra/version_compat.py b/providers/apache/cassandra/src/airflow/providers/apache/cassandra/version_compat.py index a0de76600540f..a4fae50317e81 100644 --- a/providers/apache/cassandra/src/airflow/providers/apache/cassandra/version_compat.py +++ b/providers/apache/cassandra/src/airflow/providers/apache/cassandra/version_compat.py @@ -35,20 +35,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) - -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseSensorOperator -else: - from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] - __all__ = [ "AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", - "BaseHook", - "BaseSensorOperator", ] diff --git a/providers/apache/flink/pyproject.toml b/providers/apache/flink/pyproject.toml index f6c4dc206ea30..9f1030313e34b 100644 --- a/providers/apache/flink/pyproject.toml +++ b/providers/apache/flink/pyproject.toml @@ -62,12 +62,20 @@ dependencies = [ "apache-airflow-providers-cncf-kubernetes>=5.1.0", ] +# The optional dependencies should be modified in place in the generated file +# Any change in the dependencies is preserved when the file is regenerated +[project.optional-dependencies] +"common.compat" = [ + "apache-airflow-providers-common-compat" +] + [dependency-groups] dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", "apache-airflow-providers-cncf-kubernetes", + "apache-airflow-providers-common-compat", # Additional devel dependencies (do not remove this line and add extra development dependencies) ] diff --git a/providers/apache/flink/src/airflow/providers/apache/flink/operators/flink_kubernetes.py b/providers/apache/flink/src/airflow/providers/apache/flink/operators/flink_kubernetes.py index 2c90ce9cc0e42..4a846a078c08e 100644 --- a/providers/apache/flink/src/airflow/providers/apache/flink/operators/flink_kubernetes.py +++ b/providers/apache/flink/src/airflow/providers/apache/flink/operators/flink_kubernetes.py @@ -21,13 +21,13 @@ from functools import cached_property from typing import TYPE_CHECKING -from airflow.providers.apache.flink.version_compat import BaseOperator from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook +from airflow.providers.common.compat.sdk import BaseOperator if TYPE_CHECKING: from kubernetes.client import CoreV1Api - from airflow.providers.apache.flink.version_compat import Context + from airflow.providers.common.compat.sdk import Context class FlinkKubernetesOperator(BaseOperator): diff --git a/providers/apache/flink/src/airflow/providers/apache/flink/sensors/flink_kubernetes.py b/providers/apache/flink/src/airflow/providers/apache/flink/sensors/flink_kubernetes.py index 4d86aa3d5ae06..0ed11ada0e977 100644 --- a/providers/apache/flink/src/airflow/providers/apache/flink/sensors/flink_kubernetes.py +++ b/providers/apache/flink/src/airflow/providers/apache/flink/sensors/flink_kubernetes.py @@ -23,11 +23,11 @@ from kubernetes import client from airflow.exceptions import AirflowException -from airflow.providers.apache.flink.version_compat import BaseSensorOperator from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook +from airflow.providers.common.compat.sdk import BaseSensorOperator if TYPE_CHECKING: - from airflow.providers.apache.flink.version_compat import Context + from airflow.providers.common.compat.sdk import Context class FlinkKubernetesSensor(BaseSensorOperator): diff --git a/providers/apache/flink/src/airflow/providers/apache/flink/version_compat.py b/providers/apache/flink/src/airflow/providers/apache/flink/version_compat.py index 5faab587ffcef..0956edd21112f 100644 --- a/providers/apache/flink/src/airflow/providers/apache/flink/version_compat.py +++ b/providers/apache/flink/src/airflow/providers/apache/flink/version_compat.py @@ -34,17 +34,6 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseOperator, BaseSensorOperator - from airflow.sdk.definitions.context import Context -else: - from airflow.models import BaseOperator - from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] - from airflow.utils.context import Context - __all__ = [ "AIRFLOW_V_3_0_PLUS", - "BaseOperator", - "BaseSensorOperator", - "Context", ] diff --git a/providers/apache/hdfs/pyproject.toml b/providers/apache/hdfs/pyproject.toml index cbe20c49c4d88..d78dc157dfadc 100644 --- a/providers/apache/hdfs/pyproject.toml +++ b/providers/apache/hdfs/pyproject.toml @@ -64,11 +64,19 @@ dependencies = [ 'pandas>=2.2.3; python_version >="3.13"', ] +# The optional dependencies should be modified in place in the generated file +# Any change in the dependencies is preserved when the file is regenerated +[project.optional-dependencies] +"common.compat" = [ + "apache-airflow-providers-common-compat" +] + [dependency-groups] dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", # Additional devel dependencies (do not remove this line and add extra development dependencies) ] diff --git a/providers/apache/hdfs/src/airflow/providers/apache/hdfs/hooks/webhdfs.py b/providers/apache/hdfs/src/airflow/providers/apache/hdfs/hooks/webhdfs.py index fe63b09527375..da1d55c6436d0 100644 --- a/providers/apache/hdfs/src/airflow/providers/apache/hdfs/hooks/webhdfs.py +++ b/providers/apache/hdfs/src/airflow/providers/apache/hdfs/hooks/webhdfs.py @@ -26,7 +26,7 @@ from airflow.configuration import conf from airflow.exceptions import AirflowException -from airflow.providers.apache.hdfs.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook log = logging.getLogger(__name__) diff --git a/providers/apache/hdfs/src/airflow/providers/apache/hdfs/sensors/web_hdfs.py b/providers/apache/hdfs/src/airflow/providers/apache/hdfs/sensors/web_hdfs.py index 342cb5b526ba6..57e2d1d5a5b6e 100644 --- a/providers/apache/hdfs/src/airflow/providers/apache/hdfs/sensors/web_hdfs.py +++ b/providers/apache/hdfs/src/airflow/providers/apache/hdfs/sensors/web_hdfs.py @@ -20,13 +20,13 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any -from airflow.providers.apache.hdfs.version_compat import BaseSensorOperator +from airflow.providers.common.compat.sdk import BaseSensorOperator if TYPE_CHECKING: from hdfs import InsecureClient from hdfs.ext.kerberos import KerberosClient - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class WebHdfsSensor(BaseSensorOperator): diff --git a/providers/apache/hdfs/src/airflow/providers/apache/hdfs/version_compat.py b/providers/apache/hdfs/src/airflow/providers/apache/hdfs/version_compat.py index cbdf7f8f859d9..a4fae50317e81 100644 --- a/providers/apache/hdfs/src/airflow/providers/apache/hdfs/version_compat.py +++ b/providers/apache/hdfs/src/airflow/providers/apache/hdfs/version_compat.py @@ -35,20 +35,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseSensorOperator -else: - from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] - - __all__ = [ "AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", - "BaseHook", - "BaseSensorOperator", ] diff --git a/providers/apache/hive/pyproject.toml b/providers/apache/hive/pyproject.toml index 9ce27e988e9d2..61af9f4828770 100644 --- a/providers/apache/hive/pyproject.toml +++ b/providers/apache/hive/pyproject.toml @@ -58,6 +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.7.4", # + TODO: bump to next version "apache-airflow-providers-common-sql>=1.26.0", "hmsclient>=0.1.0", 'pandas>=2.1.2; python_version <"3.13"', @@ -98,6 +99,7 @@ dev = [ "apache-airflow-task-sdk", "apache-airflow-devel-common", "apache-airflow-providers-amazon", + "apache-airflow-providers-common-compat", "apache-airflow-providers-common-sql", "apache-airflow-providers-microsoft-mssql", "apache-airflow-providers-mysql", diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py b/providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py index 4c8275d36cd0c..9b43bfb0e6743 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py @@ -33,7 +33,7 @@ from airflow.configuration import conf from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning -from airflow.providers.apache.hive.version_compat import ( +from airflow.providers.common.compat.sdk import ( AIRFLOW_VAR_NAME_FORMAT_MAPPING, BaseHook, ) diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive.py b/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive.py index 1880a7ad2e5d0..0a2763a658ca9 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive.py @@ -25,14 +25,14 @@ from airflow.configuration import conf from airflow.providers.apache.hive.hooks.hive import HiveCliHook -from airflow.providers.apache.hive.version_compat import ( +from airflow.providers.common.compat.sdk import ( AIRFLOW_VAR_NAME_FORMAT_MAPPING, BaseOperator, context_to_airflow_vars, ) if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class HiveOperator(BaseOperator): diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive_stats.py b/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive_stats.py index 85a6d22ddda91..77751816df303 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive_stats.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive_stats.py @@ -23,12 +23,12 @@ from airflow.exceptions import AirflowException from airflow.providers.apache.hive.hooks.hive import HiveMetastoreHook -from airflow.providers.apache.hive.version_compat import BaseOperator +from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.mysql.hooks.mysql import MySqlHook from airflow.providers.presto.hooks.presto import PrestoHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class HiveStatsCollectionOperator(BaseOperator): diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/sensors/hive_partition.py b/providers/apache/hive/src/airflow/providers/apache/hive/sensors/hive_partition.py index b50dcfa1ff6a0..6b2141600ae37 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/sensors/hive_partition.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/sensors/hive_partition.py @@ -21,10 +21,10 @@ from typing import TYPE_CHECKING, Any from airflow.providers.apache.hive.hooks.hive import HiveMetastoreHook -from airflow.providers.apache.hive.version_compat import BaseSensorOperator +from airflow.providers.common.compat.sdk import BaseSensorOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class HivePartitionSensor(BaseSensorOperator): diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/sensors/metastore_partition.py b/providers/apache/hive/src/airflow/providers/apache/hive/sensors/metastore_partition.py index 7f2d7dafe3159..6b3abee572c49 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/sensors/metastore_partition.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/sensors/metastore_partition.py @@ -23,7 +23,7 @@ from airflow.providers.common.sql.sensors.sql import SqlSensor if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class MetastorePartitionSensor(SqlSensor): diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/sensors/named_hive_partition.py b/providers/apache/hive/src/airflow/providers/apache/hive/sensors/named_hive_partition.py index 51d0fb86b4159..250e207e79c2c 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/sensors/named_hive_partition.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/sensors/named_hive_partition.py @@ -20,10 +20,10 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any -from airflow.providers.apache.hive.version_compat import BaseSensorOperator +from airflow.providers.common.compat.sdk import BaseSensorOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class NamedHivePartitionSensor(BaseSensorOperator): diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_mysql.py b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_mysql.py index 8764b74b694bd..a22c5ee09f312 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_mysql.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_mysql.py @@ -24,11 +24,11 @@ from typing import TYPE_CHECKING from airflow.providers.apache.hive.hooks.hive import HiveServer2Hook -from airflow.providers.apache.hive.version_compat import BaseOperator, context_to_airflow_vars +from airflow.providers.common.compat.sdk import BaseOperator, context_to_airflow_vars from airflow.providers.mysql.hooks.mysql import MySqlHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class HiveToMySqlOperator(BaseOperator): diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_samba.py b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_samba.py index a1b60a493dc4b..000b8ecf2eb1f 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_samba.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_samba.py @@ -24,11 +24,11 @@ from typing import TYPE_CHECKING from airflow.providers.apache.hive.hooks.hive import HiveServer2Hook -from airflow.providers.apache.hive.version_compat import BaseOperator, context_to_airflow_vars +from airflow.providers.common.compat.sdk import BaseOperator, context_to_airflow_vars from airflow.providers.samba.hooks.samba import SambaHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class HiveToSambaOperator(BaseOperator): diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/mssql_to_hive.py b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/mssql_to_hive.py index c09bf1f5a639d..3bb6cee2de2b2 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/mssql_to_hive.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/mssql_to_hive.py @@ -27,11 +27,11 @@ import pymssql from airflow.providers.apache.hive.hooks.hive import HiveCliHook -from airflow.providers.apache.hive.version_compat import BaseOperator +from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class MsSqlToHiveOperator(BaseOperator): diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/mysql_to_hive.py b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/mysql_to_hive.py index 0cbdeaca2366b..915518b8a1680 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/mysql_to_hive.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/mysql_to_hive.py @@ -37,11 +37,11 @@ from airflow.providers.apache.hive.hooks.hive import HiveCliHook -from airflow.providers.apache.hive.version_compat import BaseOperator +from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.mysql.hooks.mysql import MySqlHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class MySqlToHiveOperator(BaseOperator): diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/s3_to_hive.py b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/s3_to_hive.py index 636ffabd34f9d..d441761eca549 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/s3_to_hive.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/s3_to_hive.py @@ -31,10 +31,10 @@ from airflow.exceptions import AirflowException from airflow.providers.amazon.aws.hooks.s3 import S3Hook from airflow.providers.apache.hive.hooks.hive import HiveCliHook -from airflow.providers.apache.hive.version_compat import BaseOperator +from airflow.providers.common.compat.sdk import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class S3ToHiveOperator(BaseOperator): diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/vertica_to_hive.py b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/vertica_to_hive.py index 1c2c76bd0ea4d..87672b113019c 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/vertica_to_hive.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/vertica_to_hive.py @@ -25,11 +25,11 @@ from typing import TYPE_CHECKING, Any from airflow.providers.apache.hive.hooks.hive import HiveCliHook -from airflow.providers.apache.hive.version_compat import BaseOperator +from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.vertica.hooks.vertica import VerticaHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class VerticaToHiveOperator(BaseOperator): diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/version_compat.py b/providers/apache/hive/src/airflow/providers/apache/hive/version_compat.py index 832df86ea391d..bc4db65f4cac9 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/version_compat.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/version_compat.py @@ -22,6 +22,14 @@ # from __future__ import annotations +# Re-export from common.compat for backward compatibility +from airflow.providers.common.compat.sdk import ( + AIRFLOW_VAR_NAME_FORMAT_MAPPING, + BaseOperator, + BaseSensorOperator, + context_to_airflow_vars, +) + def get_base_airflow_version_tuple() -> tuple[int, int, int]: from packaging.version import Version @@ -35,27 +43,9 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseOperator, BaseSensorOperator - from airflow.sdk.execution_time.context import AIRFLOW_VAR_NAME_FORMAT_MAPPING, context_to_airflow_vars -else: - from airflow.models import BaseOperator - from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] - from airflow.utils.operator_helpers import ( # type: ignore[no-redef, attr-defined] - AIRFLOW_VAR_NAME_FORMAT_MAPPING, - context_to_airflow_vars, - ) - - __all__ = [ "AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", - "BaseHook", "BaseOperator", "BaseSensorOperator", "AIRFLOW_VAR_NAME_FORMAT_MAPPING", diff --git a/providers/apache/hive/tests/system/apache/hive/example_twitter_dag.py b/providers/apache/hive/tests/system/apache/hive/example_twitter_dag.py index ff7f722ac27de..65b8d77f35ce1 100644 --- a/providers/apache/hive/tests/system/apache/hive/example_twitter_dag.py +++ b/providers/apache/hive/tests/system/apache/hive/example_twitter_dag.py @@ -26,16 +26,9 @@ from airflow import DAG from airflow.providers.apache.hive.operators.hive import HiveOperator +from airflow.providers.common.compat.sdk import task from airflow.providers.standard.operators.bash import BashOperator -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import task -else: - # Airflow 2 path - from airflow.decorators import task # type: ignore[attr-defined,no-redef] - # -------------------------------------------------------------------------------- # Caveat: This Dag will not run because of missing scripts. # The purpose of this is to give you a sample of a real world example DAG! diff --git a/providers/apache/hive/tests/unit/apache/hive/hooks/test_hive.py b/providers/apache/hive/tests/unit/apache/hive/hooks/test_hive.py index b7d87f49046c3..5dd898a669cc1 100644 --- a/providers/apache/hive/tests/unit/apache/hive/hooks/test_hive.py +++ b/providers/apache/hive/tests/unit/apache/hive/hooks/test_hive.py @@ -31,6 +31,7 @@ from airflow.models.connection import Connection from airflow.models.dag import DAG from airflow.providers.apache.hive.hooks.hive import HiveCliHook, HiveMetastoreHook, HiveServer2Hook +from airflow.providers.common.compat.sdk import AIRFLOW_VAR_NAME_FORMAT_MAPPING from airflow.secrets.environment_variables import CONN_ENV_PREFIX from airflow.utils import timezone @@ -44,13 +45,6 @@ MockSubProcess, ) -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.execution_time.context import AIRFLOW_VAR_NAME_FORMAT_MAPPING -else: - from airflow.utils.operator_helpers import ( # type: ignore[no-redef, attr-defined] - AIRFLOW_VAR_NAME_FORMAT_MAPPING, - ) - DEFAULT_DATE = timezone.datetime(2015, 1, 1) DEFAULT_DATE_ISO = DEFAULT_DATE.isoformat() DEFAULT_DATE_DS = DEFAULT_DATE_ISO[:10] diff --git a/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_mysql.py b/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_mysql.py index 3c11dd6682f32..1dc181c071536 100644 --- a/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_mysql.py +++ b/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_mysql.py @@ -24,16 +24,11 @@ import pytest from airflow.providers.apache.hive.transfers.hive_to_mysql import HiveToMySqlOperator +from airflow.providers.common.compat.sdk import context_to_airflow_vars from airflow.utils import timezone -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS from unit.apache.hive import MockHiveServer2Hook, MockMySqlHook, TestHiveEnvironment -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.execution_time.context import context_to_airflow_vars -else: - from airflow.utils.operator_helpers import context_to_airflow_vars # type: ignore[no-redef, attr-defined] - DEFAULT_DATE = timezone.datetime(2015, 1, 1) diff --git a/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_samba.py b/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_samba.py index ec621bb59f7fb..cd0bf5e76923a 100644 --- a/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_samba.py +++ b/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_samba.py @@ -23,9 +23,9 @@ import pytest from airflow.providers.apache.hive.transfers.hive_to_samba import HiveToSambaOperator +from airflow.providers.common.compat.sdk import context_to_airflow_vars from airflow.providers.samba.hooks.samba import SambaHook -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS from unit.apache.hive import ( DEFAULT_DATE, MockConnectionCursor, @@ -33,11 +33,6 @@ TestHiveEnvironment, ) -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.execution_time.context import context_to_airflow_vars -else: - from airflow.utils.operator_helpers import context_to_airflow_vars # type: ignore[no-redef, attr-defined] - class MockSambaHook(SambaHook): def __init__(self, *args, **kwargs): diff --git a/providers/apache/iceberg/pyproject.toml b/providers/apache/iceberg/pyproject.toml index 87e6d9f78fd75..7215b707dcc10 100644 --- a/providers/apache/iceberg/pyproject.toml +++ b/providers/apache/iceberg/pyproject.toml @@ -60,11 +60,19 @@ dependencies = [ "apache-airflow>=2.10.0", ] +# The optional dependencies should be modified in place in the generated file +# Any change in the dependencies is preserved when the file is regenerated +[project.optional-dependencies] +"common.compat" = [ + "apache-airflow-providers-common-compat" +] + [dependency-groups] dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", # Additional devel dependencies (do not remove this line and add extra development dependencies) "pyiceberg>=0.5.0", ] diff --git a/providers/apache/iceberg/src/airflow/providers/apache/iceberg/hooks/iceberg.py b/providers/apache/iceberg/src/airflow/providers/apache/iceberg/hooks/iceberg.py index 3a38ca824d96d..38a0dbcc6ac82 100644 --- a/providers/apache/iceberg/src/airflow/providers/apache/iceberg/hooks/iceberg.py +++ b/providers/apache/iceberg/src/airflow/providers/apache/iceberg/hooks/iceberg.py @@ -21,7 +21,7 @@ import requests from requests import HTTPError -from airflow.providers.apache.iceberg.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook TOKENS_ENDPOINT = "oauth/tokens" diff --git a/providers/apache/iceberg/src/airflow/providers/apache/iceberg/version_compat.py b/providers/apache/iceberg/src/airflow/providers/apache/iceberg/version_compat.py index a525254fd993f..b05a7ff4d6b51 100644 --- a/providers/apache/iceberg/src/airflow/providers/apache/iceberg/version_compat.py +++ b/providers/apache/iceberg/src/airflow/providers/apache/iceberg/version_compat.py @@ -34,12 +34,6 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - __all__ = [ "AIRFLOW_V_3_1_PLUS", - "BaseHook", ] diff --git a/providers/apache/impala/src/airflow/providers/apache/impala/version_compat.py b/providers/apache/impala/src/airflow/providers/apache/impala/version_compat.py index 48d122b669696..0956edd21112f 100644 --- a/providers/apache/impala/src/airflow/providers/apache/impala/version_compat.py +++ b/providers/apache/impala/src/airflow/providers/apache/impala/version_compat.py @@ -33,3 +33,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) + +__all__ = [ + "AIRFLOW_V_3_0_PLUS", +] diff --git a/providers/apache/kafka/pyproject.toml b/providers/apache/kafka/pyproject.toml index 12b914ebe19af..edc0b3151cb90 100644 --- a/providers/apache/kafka/pyproject.toml +++ b/providers/apache/kafka/pyproject.toml @@ -70,12 +70,16 @@ dependencies = [ "common.messaging" = [ "apache-airflow-providers-common-messaging>=2.0.0" ] +"common.compat" = [ + "apache-airflow-providers-common-compat" +] [dependency-groups] dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", "apache-airflow-providers-common-messaging", "apache-airflow-providers-google", # Additional devel dependencies (do not remove this line and add extra development dependencies) diff --git a/providers/apache/kafka/src/airflow/providers/apache/kafka/hooks/base.py b/providers/apache/kafka/src/airflow/providers/apache/kafka/hooks/base.py index 95c323caf8b16..907500b3ef98d 100644 --- a/providers/apache/kafka/src/airflow/providers/apache/kafka/hooks/base.py +++ b/providers/apache/kafka/src/airflow/providers/apache/kafka/hooks/base.py @@ -21,7 +21,7 @@ from confluent_kafka.admin import AdminClient -from airflow.providers.apache.kafka.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook class KafkaBaseHook(BaseHook): diff --git a/providers/apache/kafka/src/airflow/providers/apache/kafka/operators/consume.py b/providers/apache/kafka/src/airflow/providers/apache/kafka/operators/consume.py index bf33a66f5b654..551ccbf2f17d8 100644 --- a/providers/apache/kafka/src/airflow/providers/apache/kafka/operators/consume.py +++ b/providers/apache/kafka/src/airflow/providers/apache/kafka/operators/consume.py @@ -22,7 +22,7 @@ from airflow.exceptions import AirflowException from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook -from airflow.providers.apache.kafka.version_compat import BaseOperator +from airflow.providers.common.compat.sdk import BaseOperator from airflow.utils.module_loading import import_string VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"} diff --git a/providers/apache/kafka/src/airflow/providers/apache/kafka/operators/produce.py b/providers/apache/kafka/src/airflow/providers/apache/kafka/operators/produce.py index f5b8a45fd7fea..79a888a39aaa2 100644 --- a/providers/apache/kafka/src/airflow/providers/apache/kafka/operators/produce.py +++ b/providers/apache/kafka/src/airflow/providers/apache/kafka/operators/produce.py @@ -23,7 +23,7 @@ from airflow.exceptions import AirflowException from airflow.providers.apache.kafka.hooks.produce import KafkaProducerHook -from airflow.providers.apache.kafka.version_compat import BaseOperator +from airflow.providers.common.compat.sdk import BaseOperator from airflow.utils.module_loading import import_string local_logger = logging.getLogger("airflow") diff --git a/providers/apache/kafka/src/airflow/providers/apache/kafka/sensors/kafka.py b/providers/apache/kafka/src/airflow/providers/apache/kafka/sensors/kafka.py index 7bf1886ea262c..5a55530ebd50b 100644 --- a/providers/apache/kafka/src/airflow/providers/apache/kafka/sensors/kafka.py +++ b/providers/apache/kafka/src/airflow/providers/apache/kafka/sensors/kafka.py @@ -20,7 +20,7 @@ from typing import Any from airflow.providers.apache.kafka.triggers.await_message import AwaitMessageTrigger -from airflow.providers.apache.kafka.version_compat import BaseOperator +from airflow.providers.common.compat.sdk import BaseOperator VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"} diff --git a/providers/apache/kafka/src/airflow/providers/apache/kafka/version_compat.py b/providers/apache/kafka/src/airflow/providers/apache/kafka/version_compat.py index 57b7277fde97b..a4fae50317e81 100644 --- a/providers/apache/kafka/src/airflow/providers/apache/kafka/version_compat.py +++ b/providers/apache/kafka/src/airflow/providers/apache/kafka/version_compat.py @@ -35,18 +35,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseOperator -else: - from airflow.models import BaseOperator - __all__ = [ "AIRFLOW_V_3_0_PLUS", - "BaseHook", - "BaseOperator", + "AIRFLOW_V_3_1_PLUS", ] diff --git a/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py b/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py index 57f90832bd27a..18525ab8ca923 100644 --- a/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py +++ b/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py @@ -25,14 +25,14 @@ from airflow.exceptions import AirflowException from airflow.providers.apache.livy.hooks.livy import BatchState, LivyHook from airflow.providers.apache.livy.triggers.livy import LivyTrigger -from airflow.providers.apache.livy.version_compat import BaseOperator from airflow.providers.common.compat.openlineage.utils.spark import ( inject_parent_job_information_into_spark_properties, inject_transport_information_into_spark_properties, ) +from airflow.providers.common.compat.sdk import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class LivyOperator(BaseOperator): diff --git a/providers/apache/livy/src/airflow/providers/apache/livy/sensors/livy.py b/providers/apache/livy/src/airflow/providers/apache/livy/sensors/livy.py index ca0b352d7ee1a..9084763f8a3db 100644 --- a/providers/apache/livy/src/airflow/providers/apache/livy/sensors/livy.py +++ b/providers/apache/livy/src/airflow/providers/apache/livy/sensors/livy.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING, Any from airflow.providers.apache.livy.hooks.livy import LivyHook -from airflow.providers.apache.livy.version_compat import BaseSensorOperator +from airflow.providers.common.compat.sdk import BaseSensorOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class LivySensor(BaseSensorOperator): diff --git a/providers/apache/livy/src/airflow/providers/apache/livy/version_compat.py b/providers/apache/livy/src/airflow/providers/apache/livy/version_compat.py index 126f95c848abd..55c0b32b890d1 100644 --- a/providers/apache/livy/src/airflow/providers/apache/livy/version_compat.py +++ b/providers/apache/livy/src/airflow/providers/apache/livy/version_compat.py @@ -34,10 +34,4 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseOperator, BaseSensorOperator -else: - from airflow.models import BaseOperator - from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] - -__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseOperator", "BaseSensorOperator"] +__all__ = ["AIRFLOW_V_3_0_PLUS"] diff --git a/providers/apache/pig/pyproject.toml b/providers/apache/pig/pyproject.toml index 26498379927b5..047f933d31b01 100644 --- a/providers/apache/pig/pyproject.toml +++ b/providers/apache/pig/pyproject.toml @@ -60,11 +60,19 @@ dependencies = [ "apache-airflow>=2.10.0", ] +# The optional dependencies should be modified in place in the generated file +# Any change in the dependencies is preserved when the file is regenerated +[project.optional-dependencies] +"common.compat" = [ + "apache-airflow-providers-common-compat" +] + [dependency-groups] dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", # Additional devel dependencies (do not remove this line and add extra development dependencies) ] diff --git a/providers/apache/pig/src/airflow/providers/apache/pig/hooks/pig.py b/providers/apache/pig/src/airflow/providers/apache/pig/hooks/pig.py index 99a97c4cd5e12..21fc14da0ea8e 100644 --- a/providers/apache/pig/src/airflow/providers/apache/pig/hooks/pig.py +++ b/providers/apache/pig/src/airflow/providers/apache/pig/hooks/pig.py @@ -22,7 +22,7 @@ from typing import Any from airflow.exceptions import AirflowException -from airflow.providers.apache.pig.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook class PigCliHook(BaseHook): diff --git a/providers/apache/pig/src/airflow/providers/apache/pig/operators/pig.py b/providers/apache/pig/src/airflow/providers/apache/pig/operators/pig.py index 742e15acd742e..888c1059ad2a3 100644 --- a/providers/apache/pig/src/airflow/providers/apache/pig/operators/pig.py +++ b/providers/apache/pig/src/airflow/providers/apache/pig/operators/pig.py @@ -22,10 +22,10 @@ from typing import TYPE_CHECKING, Any from airflow.providers.apache.pig.hooks.pig import PigCliHook -from airflow.providers.apache.pig.version_compat import BaseOperator +from airflow.providers.common.compat.sdk import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class PigOperator(BaseOperator): diff --git a/providers/apache/pig/src/airflow/providers/apache/pig/version_compat.py b/providers/apache/pig/src/airflow/providers/apache/pig/version_compat.py index 615e9c2f23db7..a4fae50317e81 100644 --- a/providers/apache/pig/src/airflow/providers/apache/pig/version_compat.py +++ b/providers/apache/pig/src/airflow/providers/apache/pig/version_compat.py @@ -35,19 +35,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseOperator -else: - from airflow.models import BaseOperator - __all__ = [ "AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", - "BaseHook", - "BaseOperator", ] diff --git a/providers/apache/pinot/pyproject.toml b/providers/apache/pinot/pyproject.toml index f305e5891ca98..4f0361912daae 100644 --- a/providers/apache/pinot/pyproject.toml +++ b/providers/apache/pinot/pyproject.toml @@ -62,11 +62,19 @@ dependencies = [ "pinotdb>=5.1.0", ] +# The optional dependencies should be modified in place in the generated file +# Any change in the dependencies is preserved when the file is regenerated +[project.optional-dependencies] +"common.compat" = [ + "apache-airflow-providers-common-compat" +] + [dependency-groups] dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", "apache-airflow-providers-common-sql", # Additional devel dependencies (do not remove this line and add extra development dependencies) "apache-airflow-providers-common-sql[pandas,polars]" diff --git a/providers/apache/pinot/src/airflow/providers/apache/pinot/hooks/pinot.py b/providers/apache/pinot/src/airflow/providers/apache/pinot/hooks/pinot.py index 6394d0e524b1c..12960171a0699 100644 --- a/providers/apache/pinot/src/airflow/providers/apache/pinot/hooks/pinot.py +++ b/providers/apache/pinot/src/airflow/providers/apache/pinot/hooks/pinot.py @@ -26,7 +26,7 @@ from pinotdb import connect from airflow.exceptions import AirflowException -from airflow.providers.apache.pinot.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook from airflow.providers.common.sql.hooks.sql import DbApiHook if TYPE_CHECKING: diff --git a/providers/apache/pinot/src/airflow/providers/apache/pinot/version_compat.py b/providers/apache/pinot/src/airflow/providers/apache/pinot/version_compat.py index a525254fd993f..b05a7ff4d6b51 100644 --- a/providers/apache/pinot/src/airflow/providers/apache/pinot/version_compat.py +++ b/providers/apache/pinot/src/airflow/providers/apache/pinot/version_compat.py @@ -34,12 +34,6 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - __all__ = [ "AIRFLOW_V_3_1_PLUS", - "BaseHook", ] diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/decorators/pyspark.py b/providers/apache/spark/src/airflow/providers/apache/spark/decorators/pyspark.py index e59af8cd3d83e..e5080c151b549 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/decorators/pyspark.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/decorators/pyspark.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any from airflow.providers.apache.spark.hooks.spark_connect import SparkConnectHook -from airflow.providers.apache.spark.version_compat import ( +from airflow.providers.common.compat.sdk import ( BaseHook, DecoratedOperator, TaskDecorator, @@ -31,8 +31,7 @@ from airflow.providers.common.compat.standard.operators import PythonOperator if TYPE_CHECKING: - from airflow.utils.context import Context - + from airflow.providers.common.compat.sdk import Context SPARK_CONTEXT_KEYS = ["spark", "sc"] diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_connect.py b/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_connect.py index de5e39bed8d4a..67ff43b5e84f1 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_connect.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_connect.py @@ -20,7 +20,7 @@ from typing import Any, cast from urllib.parse import quote, urlparse, urlunparse -from airflow.providers.apache.spark.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook from airflow.utils.log.logging_mixin import LoggingMixin diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_sql.py b/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_sql.py index 119e265dc0492..9b7b485af25ec 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_sql.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_sql.py @@ -21,7 +21,7 @@ from typing import TYPE_CHECKING, Any from airflow.exceptions import AirflowException, AirflowNotFoundException -from airflow.providers.apache.spark.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook if TYPE_CHECKING: try: diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py b/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py index eba5292dd2d75..4a32e93d0cdb0 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py @@ -32,7 +32,7 @@ from airflow.configuration import conf as airflow_conf from airflow.exceptions import AirflowException -from airflow.providers.apache.spark.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook from airflow.security.kerberos import renew_from_kt from airflow.utils.log.logging_mixin import LoggingMixin diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_jdbc.py b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_jdbc.py index 465e23ab53c4d..a4a334fb45283 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_jdbc.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_jdbc.py @@ -23,7 +23,7 @@ from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class SparkJDBCOperator(SparkSubmitOperator): diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_sql.py b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_sql.py index e4fc07e6ec4f5..ccb6b58e5879f 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_sql.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_sql.py @@ -21,10 +21,10 @@ from typing import TYPE_CHECKING, Any from airflow.providers.apache.spark.hooks.spark_sql import SparkSqlHook -from airflow.providers.apache.spark.version_compat import BaseOperator +from airflow.providers.common.compat.sdk import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class SparkSqlOperator(BaseOperator): diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py index 44fd5f3f43bf7..88097284f3e9e 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py @@ -22,15 +22,15 @@ from airflow.configuration import conf from airflow.providers.apache.spark.hooks.spark_submit import SparkSubmitHook -from airflow.providers.apache.spark.version_compat import BaseOperator from airflow.providers.common.compat.openlineage.utils.spark import ( inject_parent_job_information_into_spark_properties, inject_transport_information_into_spark_properties, ) +from airflow.providers.common.compat.sdk import BaseOperator from airflow.settings import WEB_COLORS if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class SparkSubmitOperator(BaseOperator): diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/version_compat.py b/providers/apache/spark/src/airflow/providers/apache/spark/version_compat.py index 46dd58a433ac9..146a5c0c1e909 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/version_compat.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/version_compat.py @@ -22,6 +22,15 @@ # from __future__ import annotations +# Re-export from common.compat for backward compatibility +from airflow.providers.common.compat.sdk import ( + BaseHook, + BaseOperator, + DecoratedOperator, + TaskDecorator, + task_decorator_factory, +) + def get_base_airflow_version_tuple() -> tuple[int, int, int]: from packaging.version import Version @@ -35,22 +44,6 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseOperator - from airflow.sdk.bases.decorator import DecoratedOperator, TaskDecorator, task_decorator_factory -else: - from airflow.decorators.base import ( # type: ignore[no-redef] - DecoratedOperator, - TaskDecorator, - task_decorator_factory, - ) - from airflow.models import BaseOperator - __all__ = [ "AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", diff --git a/providers/apache/tinkerpop/pyproject.toml b/providers/apache/tinkerpop/pyproject.toml index 69f5fba657b67..055512e5a84ed 100644 --- a/providers/apache/tinkerpop/pyproject.toml +++ b/providers/apache/tinkerpop/pyproject.toml @@ -61,11 +61,19 @@ dependencies = [ "gremlinpython>=3.7.3", ] +# The optional dependencies should be modified in place in the generated file +# Any change in the dependencies is preserved when the file is regenerated +[project.optional-dependencies] +"common.compat" = [ + "apache-airflow-providers-common-compat" +] + [dependency-groups] dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", # Additional devel dependencies (do not remove this line and add extra development dependencies) ] diff --git a/providers/apache/tinkerpop/src/airflow/providers/apache/tinkerpop/hooks/gremlin.py b/providers/apache/tinkerpop/src/airflow/providers/apache/tinkerpop/hooks/gremlin.py index e6bfd75242afb..8fd35acb1f9ec 100644 --- a/providers/apache/tinkerpop/src/airflow/providers/apache/tinkerpop/hooks/gremlin.py +++ b/providers/apache/tinkerpop/src/airflow/providers/apache/tinkerpop/hooks/gremlin.py @@ -23,7 +23,7 @@ from gremlin_python.driver.client import Client -from airflow.providers.apache.tinkerpop.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook if TYPE_CHECKING: from airflow.models import Connection diff --git a/providers/apache/tinkerpop/src/airflow/providers/apache/tinkerpop/operators/gremlin.py b/providers/apache/tinkerpop/src/airflow/providers/apache/tinkerpop/operators/gremlin.py index 7a38a220d979e..4050be00972b3 100644 --- a/providers/apache/tinkerpop/src/airflow/providers/apache/tinkerpop/operators/gremlin.py +++ b/providers/apache/tinkerpop/src/airflow/providers/apache/tinkerpop/operators/gremlin.py @@ -19,14 +19,10 @@ from typing import TYPE_CHECKING, Any from airflow.providers.apache.tinkerpop.hooks.gremlin import GremlinHook -from airflow.providers.apache.tinkerpop.version_compat import BaseOperator +from airflow.providers.common.compat.sdk import BaseOperator if TYPE_CHECKING: - try: - from airflow.sdk.definitions.context import Context - except ImportError: - # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GremlinOperator(BaseOperator): diff --git a/providers/apache/tinkerpop/src/airflow/providers/apache/tinkerpop/version_compat.py b/providers/apache/tinkerpop/src/airflow/providers/apache/tinkerpop/version_compat.py index 25e10387f37fb..f67f0a5fa1860 100644 --- a/providers/apache/tinkerpop/src/airflow/providers/apache/tinkerpop/version_compat.py +++ b/providers/apache/tinkerpop/src/airflow/providers/apache/tinkerpop/version_compat.py @@ -29,18 +29,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseOperator -else: - from airflow.models import BaseOperator - __all__ = [ "AIRFLOW_V_3_0_PLUS", - "BaseHook", - "BaseOperator", + "AIRFLOW_V_3_1_PLUS", ] diff --git a/providers/common/compat/src/airflow/providers/common/compat/sdk.py b/providers/common/compat/src/airflow/providers/common/compat/sdk.py index 3ccda1e1929d4..080c2658bd5b2 100644 --- a/providers/common/compat/src/airflow/providers/common/compat/sdk.py +++ b/providers/common/compat/src/airflow/providers/common/compat/sdk.py @@ -77,7 +77,10 @@ from airflow.sdk.definitions.context import context_merge as context_merge from airflow.sdk.definitions.mappedoperator import MappedOperator as MappedOperator from airflow.sdk.definitions.template import literal as literal - from airflow.sdk.execution_time.context import context_to_airflow_vars as context_to_airflow_vars + from airflow.sdk.execution_time.context import ( + AIRFLOW_VAR_NAME_FORMAT_MAPPING as AIRFLOW_VAR_NAME_FORMAT_MAPPING, + context_to_airflow_vars as context_to_airflow_vars, + ) from airflow.sdk.execution_time.timeout import timeout as timeout from airflow.sdk.execution_time.xcom import XCom as XCom @@ -182,6 +185,10 @@ "Context": ("airflow.sdk", "airflow.utils.context"), "context_merge": ("airflow.sdk.definitions.context", "airflow.utils.context"), "context_to_airflow_vars": ("airflow.sdk.execution_time.context", "airflow.utils.operator_helpers"), + "AIRFLOW_VAR_NAME_FORMAT_MAPPING": ( + "airflow.sdk.execution_time.context", + "airflow.utils.operator_helpers", + ), "get_current_context": ("airflow.sdk", "airflow.operators.python"), "get_parsing_context": ("airflow.sdk", "airflow.utils.dag_parsing_context"), # ============================================================================ diff --git a/providers/elasticsearch/pyproject.toml b/providers/elasticsearch/pyproject.toml index c4674ec0faf73..3baa2ba6d2cc1 100644 --- a/providers/elasticsearch/pyproject.toml +++ b/providers/elasticsearch/pyproject.toml @@ -58,6 +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.7.4", # + TODO: bump to next version "apache-airflow-providers-common-sql>=1.27.0", "elasticsearch>=8.10,<9", ] @@ -67,6 +68,7 @@ dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", "apache-airflow-providers-common-sql", # Additional devel dependencies (do not remove this line and add extra development dependencies) "apache-airflow-providers-common-sql[pandas,polars]", diff --git a/providers/elasticsearch/src/airflow/providers/elasticsearch/hooks/elasticsearch.py b/providers/elasticsearch/src/airflow/providers/elasticsearch/hooks/elasticsearch.py index e0bcc07683042..2bbc5bab17614 100644 --- a/providers/elasticsearch/src/airflow/providers/elasticsearch/hooks/elasticsearch.py +++ b/providers/elasticsearch/src/airflow/providers/elasticsearch/hooks/elasticsearch.py @@ -25,8 +25,8 @@ from elasticsearch import Elasticsearch +from airflow.providers.common.compat.sdk import BaseHook from airflow.providers.common.sql.hooks.sql import DbApiHook -from airflow.providers.elasticsearch.version_compat import BaseHook if TYPE_CHECKING: from elastic_transport import ObjectApiResponse diff --git a/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py b/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py index ef10047eeb486..c38469bc17230 100644 --- a/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py +++ b/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py @@ -41,23 +41,15 @@ from airflow.configuration import conf from airflow.exceptions import AirflowException from airflow.models.dagrun import DagRun +from airflow.providers.common.compat.sdk import timezone from airflow.providers.elasticsearch.log.es_json_formatter import ElasticsearchJSONFormatter from airflow.providers.elasticsearch.log.es_response import ElasticSearchResponse, Hit -from airflow.providers.elasticsearch.version_compat import ( - AIRFLOW_V_3_0_PLUS, - AIRFLOW_V_3_1_PLUS, - EsLogMsgType, -) +from airflow.providers.elasticsearch.version_compat import AIRFLOW_V_3_0_PLUS from airflow.utils.log.file_task_handler import FileTaskHandler from airflow.utils.log.logging_mixin import ExternalLoggingMixin, LoggingMixin from airflow.utils.module_loading import import_string from airflow.utils.session import create_session -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import timezone -else: - from airflow.utils import timezone # type: ignore[attr-defined,no-redef] - if TYPE_CHECKING: from datetime import datetime @@ -65,6 +57,14 @@ from airflow.utils.log.file_task_handler import LogMetadata +if AIRFLOW_V_3_0_PLUS: + from airflow.utils.log.file_task_handler import StructuredLogMessage + + EsLogMsgType = list[StructuredLogMessage] | str +else: + EsLogMsgType = list[tuple[str, str]] # type: ignore[assignment,misc] + + LOG_LINE_DEFAULTS = {"exc_text": "", "stack_info": ""} # Elasticsearch hosted log type diff --git a/providers/elasticsearch/src/airflow/providers/elasticsearch/version_compat.py b/providers/elasticsearch/src/airflow/providers/elasticsearch/version_compat.py index 679a54714add7..f5bb3ae555c1f 100644 --- a/providers/elasticsearch/src/airflow/providers/elasticsearch/version_compat.py +++ b/providers/elasticsearch/src/airflow/providers/elasticsearch/version_compat.py @@ -35,16 +35,4 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - -if AIRFLOW_V_3_0_PLUS: - from airflow.utils.log.file_task_handler import StructuredLogMessage - - EsLogMsgType = list[StructuredLogMessage] | str -else: - EsLogMsgType = list[tuple[str, str]] # type: ignore[assignment,misc] - -__all__ = ["AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", "BaseHook", "EsLogMsgType"] +__all__ = ["AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS"] diff --git a/providers/elasticsearch/tests/system/elasticsearch/example_elasticsearch_query.py b/providers/elasticsearch/tests/system/elasticsearch/example_elasticsearch_query.py index b0cc58a553aa6..6dd0203d09b07 100644 --- a/providers/elasticsearch/tests/system/elasticsearch/example_elasticsearch_query.py +++ b/providers/elasticsearch/tests/system/elasticsearch/example_elasticsearch_query.py @@ -24,17 +24,10 @@ from datetime import datetime from airflow import models +from airflow.providers.common.compat.sdk import task from airflow.providers.elasticsearch.hooks.elasticsearch import ElasticsearchPythonHook, ElasticsearchSQLHook from airflow.providers.standard.operators.python import PythonOperator -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import task -else: - # Airflow 2 path - from airflow.decorators import task # type: ignore[attr-defined,no-redef] - ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID") DAG_ID = "elasticsearch_dag" CONN_ID = "elasticsearch_default"