diff --git a/providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/operators/psrp.py b/providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/operators/psrp.py index 6927159a44ea2..3d297cf9b30f7 100644 --- a/providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/operators/psrp.py +++ b/providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/operators/psrp.py @@ -25,8 +25,8 @@ from pypsrp.serializer import TaggedValue from airflow.exceptions import AirflowException -from airflow.models import BaseOperator from airflow.providers.microsoft.psrp.hooks.psrp import PsrpHook +from airflow.providers.microsoft.psrp.version_compat import BaseOperator from airflow.settings import json from airflow.utils.helpers import exactly_one diff --git a/providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/version_compat.py b/providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/version_compat.py new file mode 100644 index 0000000000000..89ade07a76753 --- /dev/null +++ b/providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/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 # type: ignore[no-redef] + +__all__ = [ + "AIRFLOW_V_3_0_PLUS", + "BaseOperator", +] diff --git a/providers/microsoft/psrp/tests/unit/microsoft/psrp/operators/test_psrp.py b/providers/microsoft/psrp/tests/unit/microsoft/psrp/operators/test_psrp.py index 42bc6a7ea079d..c38c8ad1be8f7 100644 --- a/providers/microsoft/psrp/tests/unit/microsoft/psrp/operators/test_psrp.py +++ b/providers/microsoft/psrp/tests/unit/microsoft/psrp/operators/test_psrp.py @@ -25,8 +25,8 @@ from pypsrp.powershell import Command, PowerShell from airflow.exceptions import AirflowException -from airflow.models.baseoperator import BaseOperator from airflow.providers.microsoft.psrp.operators.psrp import PsrpOperator +from airflow.providers.microsoft.psrp.version_compat import BaseOperator from airflow.settings import json CONNECTION_ID = "conn_id"