From 7733da1c5ae2e07510fa86088e045eeeb37552d2 Mon Sep 17 00:00:00 2001 From: Kevin Yang Date: Sun, 29 Jun 2025 18:27:50 -0400 Subject: [PATCH 1/2] Update BaseOperator imports for Airflow 3.0 compatibility merge updates from master --- .../unit/always/test_project_structure.py | 5 ++ .../providers/apache/druid/operators/druid.py | 4 +- .../apache/druid/transfers/hive_to_druid.py | 4 +- .../providers/apache/druid/version_compat.py | 48 +++++++++++++++++++ 4 files changed, 57 insertions(+), 4 deletions(-) create mode 100644 providers/apache/druid/src/airflow/providers/apache/druid/version_compat.py diff --git a/airflow-core/tests/unit/always/test_project_structure.py b/airflow-core/tests/unit/always/test_project_structure.py index 2a818c13ac4dd..315322cb27755 100644 --- a/airflow-core/tests/unit/always/test_project_structure.py +++ b/airflow-core/tests/unit/always/test_project_structure.py @@ -82,6 +82,11 @@ def test_providers_modules_should_have_tests(self): "providers/amazon/tests/unit/amazon/aws/utils/test_rds.py", "providers/amazon/tests/unit/amazon/aws/utils/test_sagemaker.py", "providers/amazon/tests/unit/amazon/aws/waiters/test_base_waiter.py", + "providers/amazon/tests/unit/amazon/test_version_compat.py", + "providers/apache/cassandra/tests/unit/apache/cassandra/test_version_compat.py", + "providers/apache/druid/tests/unit/apache/druid/test_version_compat.py", + "providers/apache/flink/tests/unit/apache/flink/test_version_compat.py", + "providers/apache/hdfs/tests/unit/apache/hdfs/test_version_compat.py", "providers/apache/hdfs/tests/unit/apache/hdfs/hooks/test_hdfs.py", "providers/apache/hdfs/tests/unit/apache/hdfs/log/test_hdfs_task_handler.py", "providers/apache/hdfs/tests/unit/apache/hdfs/sensors/test_hdfs.py", diff --git a/providers/apache/druid/src/airflow/providers/apache/druid/operators/druid.py b/providers/apache/druid/src/airflow/providers/apache/druid/operators/druid.py index 6867488028d29..cd9c70a06e06b 100644 --- a/providers/apache/druid/src/airflow/providers/apache/druid/operators/druid.py +++ b/providers/apache/druid/src/airflow/providers/apache/druid/operators/druid.py @@ -20,11 +20,11 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any -from airflow.models import BaseOperator from airflow.providers.apache.druid.hooks.druid import DruidHook, IngestionType +from airflow.providers.apache.druid.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.apache.druid.version_compat import Context class DruidOperator(BaseOperator): diff --git a/providers/apache/druid/src/airflow/providers/apache/druid/transfers/hive_to_druid.py b/providers/apache/druid/src/airflow/providers/apache/druid/transfers/hive_to_druid.py index 1e5280eedfe4c..5cd4db23c469e 100644 --- a/providers/apache/druid/src/airflow/providers/apache/druid/transfers/hive_to_druid.py +++ b/providers/apache/druid/src/airflow/providers/apache/druid/transfers/hive_to_druid.py @@ -22,12 +22,12 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any -from airflow.models import BaseOperator from airflow.providers.apache.druid.hooks.druid import DruidHook +from airflow.providers.apache.druid.version_compat import BaseOperator from airflow.providers.apache.hive.hooks.hive import HiveCliHook, HiveMetastoreHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.apache.druid.version_compat import Context LOAD_CHECK_INTERVAL = 5 DEFAULT_TARGET_PARTITION_SIZE = 5000000 diff --git a/providers/apache/druid/src/airflow/providers/apache/druid/version_compat.py b/providers/apache/druid/src/airflow/providers/apache/druid/version_compat.py new file mode 100644 index 0000000000000..3b1b09b65e281 --- /dev/null +++ b/providers/apache/druid/src/airflow/providers/apache/druid/version_compat.py @@ -0,0 +1,48 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# NOTE! THIS FILE IS COPIED MANUALLY IN OTHER PROVIDERS DELIBERATELY TO AVOID ADDING UNNECESSARY +# DEPENDENCIES BETWEEN PROVIDERS. IF YOU WANT TO ADD CONDITIONAL CODE IN YOUR PROVIDER THAT DEPENDS +# ON AIRFLOW VERSION, PLEASE COPY THIS FILE TO THE ROOT PACKAGE OF YOUR PROVIDER AND IMPORT +# THOSE CONSTANTS FROM IT RATHER THAN IMPORTING THEM FROM ANOTHER PROVIDER OR TEST CODE +# +from __future__ import annotations + + +def get_base_airflow_version_tuple() -> tuple[int, int, int]: + from packaging.version import Version + + from airflow import __version__ + + airflow_version = Version(__version__) + return airflow_version.major, airflow_version.minor, airflow_version.micro + + +AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk import BaseOperator + from airflow.sdk.definitions.context import Context +else: + from airflow.models import BaseOperator + from airflow.utils.context import Context + +__all__ = [ + "AIRFLOW_V_3_0_PLUS", + "BaseOperator", + "Context", +] From f255965d1e1ad52d31e2c2f609ce2f0b28d32579 Mon Sep 17 00:00:00 2001 From: Kevin Yang Date: Sun, 29 Jun 2025 18:47:26 -0400 Subject: [PATCH 2/2] remove version_compat.py update as PR 52496 --- airflow-core/tests/unit/always/test_project_structure.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/airflow-core/tests/unit/always/test_project_structure.py b/airflow-core/tests/unit/always/test_project_structure.py index 315322cb27755..2a818c13ac4dd 100644 --- a/airflow-core/tests/unit/always/test_project_structure.py +++ b/airflow-core/tests/unit/always/test_project_structure.py @@ -82,11 +82,6 @@ def test_providers_modules_should_have_tests(self): "providers/amazon/tests/unit/amazon/aws/utils/test_rds.py", "providers/amazon/tests/unit/amazon/aws/utils/test_sagemaker.py", "providers/amazon/tests/unit/amazon/aws/waiters/test_base_waiter.py", - "providers/amazon/tests/unit/amazon/test_version_compat.py", - "providers/apache/cassandra/tests/unit/apache/cassandra/test_version_compat.py", - "providers/apache/druid/tests/unit/apache/druid/test_version_compat.py", - "providers/apache/flink/tests/unit/apache/flink/test_version_compat.py", - "providers/apache/hdfs/tests/unit/apache/hdfs/test_version_compat.py", "providers/apache/hdfs/tests/unit/apache/hdfs/hooks/test_hdfs.py", "providers/apache/hdfs/tests/unit/apache/hdfs/log/test_hdfs_task_handler.py", "providers/apache/hdfs/tests/unit/apache/hdfs/sensors/test_hdfs.py",