From 082ee51f3ab7e9082fc893d7396250393e6b757a Mon Sep 17 00:00:00 2001 From: Kevin Yang Date: Sat, 28 Jun 2025 00:11:38 -0400 Subject: [PATCH 1/3] Update BaseOperator imports for Airflow 3.0 compatibility --- .../providers/weaviate/operators/weaviate.py | 2 +- .../providers/weaviate/version_compat.py | 40 +++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) create mode 100644 providers/weaviate/src/airflow/providers/weaviate/version_compat.py diff --git a/providers/weaviate/src/airflow/providers/weaviate/operators/weaviate.py b/providers/weaviate/src/airflow/providers/weaviate/operators/weaviate.py index 1facdf77beb0d..fd8658e552245 100644 --- a/providers/weaviate/src/airflow/providers/weaviate/operators/weaviate.py +++ b/providers/weaviate/src/airflow/providers/weaviate/operators/weaviate.py @@ -21,8 +21,8 @@ from functools import cached_property from typing import TYPE_CHECKING, Any -from airflow.models import BaseOperator from airflow.providers.weaviate.hooks.weaviate import WeaviateHook +from airflow.providers.weaviate.version_compat import BaseOperator if TYPE_CHECKING: import pandas as pd diff --git a/providers/weaviate/src/airflow/providers/weaviate/version_compat.py b/providers/weaviate/src/airflow/providers/weaviate/version_compat.py new file mode 100644 index 0000000000000..78cd58dd016ab --- /dev/null +++ b/providers/weaviate/src/airflow/providers/weaviate/version_compat.py @@ -0,0 +1,40 @@ +# 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. + +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 +else: + from airflow.models import BaseOperator + +__all__ = [ + "AIRFLOW_V_3_0_PLUS", + "BaseOperator", +] From 98ff79bf68a8c73310dd059ba604fea8e0ecf412 Mon Sep 17 00:00:00 2001 From: Kevin Yang Date: Sat, 28 Jun 2025 16:24:28 -0400 Subject: [PATCH 2/3] update according to latest instruction --- airflow-core/tests/unit/always/test_project_structure.py | 1 + .../src/airflow/providers/weaviate/operators/weaviate.py | 6 +----- .../src/airflow/providers/weaviate/version_compat.py | 3 +++ .../weaviate/tests/unit/weaviate/operators/test_weaviate.py | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/airflow-core/tests/unit/always/test_project_structure.py b/airflow-core/tests/unit/always/test_project_structure.py index 14d63b8b2116b..26b8ef2af864f 100644 --- a/airflow-core/tests/unit/always/test_project_structure.py +++ b/airflow-core/tests/unit/always/test_project_structure.py @@ -239,6 +239,7 @@ def test_providers_modules_should_have_tests(self): "providers/standard/tests/unit/standard/utils/test_sensor_helper.py", "providers/tableau/tests/unit/tableau/test_version_compat.py", "providers/trino/tests/unit/trino/test_version_compat.py", + "providers/weaviate/tests/unit/weaviate/test_version_compat.py", ] modules_files: list[pathlib.Path] = list( AIRFLOW_PROVIDERS_ROOT_PATH.glob("**/src/airflow/providers/**/*.py") diff --git a/providers/weaviate/src/airflow/providers/weaviate/operators/weaviate.py b/providers/weaviate/src/airflow/providers/weaviate/operators/weaviate.py index fd8658e552245..793f0a1d4ac91 100644 --- a/providers/weaviate/src/airflow/providers/weaviate/operators/weaviate.py +++ b/providers/weaviate/src/airflow/providers/weaviate/operators/weaviate.py @@ -28,11 +28,7 @@ import pandas as pd from weaviate.types import UUID - 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.weaviate.version_compat import Context class WeaviateIngestOperator(BaseOperator): diff --git a/providers/weaviate/src/airflow/providers/weaviate/version_compat.py b/providers/weaviate/src/airflow/providers/weaviate/version_compat.py index 78cd58dd016ab..c267288006129 100644 --- a/providers/weaviate/src/airflow/providers/weaviate/version_compat.py +++ b/providers/weaviate/src/airflow/providers/weaviate/version_compat.py @@ -31,10 +31,13 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: 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", ] diff --git a/providers/weaviate/tests/unit/weaviate/operators/test_weaviate.py b/providers/weaviate/tests/unit/weaviate/operators/test_weaviate.py index aa3be7a3cd8f3..9b38e094d39d2 100644 --- a/providers/weaviate/tests/unit/weaviate/operators/test_weaviate.py +++ b/providers/weaviate/tests/unit/weaviate/operators/test_weaviate.py @@ -50,7 +50,7 @@ def test_constructor(self, operator): def test_execute_with_input_data(self, mock_log, operator): operator.hook.batch_data = MagicMock() - operator.execute(context=None) + operator.execute(context=None) # type: ignore[arg-type] operator.hook.batch_data.assert_called_once_with( collection_name="my_collection", From d642edf12387ece1b97be604ba2f661551de7e8e Mon Sep 17 00:00:00 2001 From: Kevin Yang Date: Sat, 28 Jun 2025 16:52:05 -0400 Subject: [PATCH 3/3] remove type ignore since it is not a mock context --- .../weaviate/tests/unit/weaviate/operators/test_weaviate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/weaviate/tests/unit/weaviate/operators/test_weaviate.py b/providers/weaviate/tests/unit/weaviate/operators/test_weaviate.py index 9b38e094d39d2..aa3be7a3cd8f3 100644 --- a/providers/weaviate/tests/unit/weaviate/operators/test_weaviate.py +++ b/providers/weaviate/tests/unit/weaviate/operators/test_weaviate.py @@ -50,7 +50,7 @@ def test_constructor(self, operator): def test_execute_with_input_data(self, mock_log, operator): operator.hook.batch_data = MagicMock() - operator.execute(context=None) # type: ignore[arg-type] + operator.execute(context=None) operator.hook.batch_data.assert_called_once_with( collection_name="my_collection",