From 0ed788f5fd9ada591c0f2c208b4818e561ea11ef Mon Sep 17 00:00:00 2001 From: Mingdao Yang Date: Fri, 4 Jul 2025 01:36:26 +0800 Subject: [PATCH 1/3] #52676 Move all BaseHook usages to version_compat in openai --- .../src/airflow/providers/openai/hooks/openai.py | 5 +---- .../src/airflow/providers/openai/version_compat.py | 14 +++++++++++++- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/providers/openai/src/airflow/providers/openai/hooks/openai.py b/providers/openai/src/airflow/providers/openai/hooks/openai.py index 18545aca14248..6c3758481e7ae 100644 --- a/providers/openai/src/airflow/providers/openai/hooks/openai.py +++ b/providers/openai/src/airflow/providers/openai/hooks/openai.py @@ -45,10 +45,7 @@ from openai.types.vector_stores import VectorStoreFile, VectorStoreFileBatch, VectorStoreFileDeleted from airflow.providers.openai.exceptions import OpenAIBatchJobException, OpenAIBatchTimeout -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.facebook.version_compat import BaseHook class BatchStatus(str, Enum): diff --git a/providers/openai/src/airflow/providers/openai/version_compat.py b/providers/openai/src/airflow/providers/openai/version_compat.py index 4f8d5e32bca4a..eb342ba7a11c3 100644 --- a/providers/openai/src/airflow/providers/openai/version_compat.py +++ b/providers/openai/src/airflow/providers/openai/version_compat.py @@ -33,10 +33,22 @@ 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 # type: ignore[no-redef] -__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseOperator"] +__all__ = [ + "AIRFLOW_V_3_0_PLUS", + "AIRFLOW_V_3_1_PLUS", + "BaseHook", + "BaseOperator", +] From 098269a23b9572975d4d63637d0a8b78aeca34a0 Mon Sep 17 00:00:00 2001 From: Mingdao Yang Date: Sat, 5 Jul 2025 02:27:59 +0800 Subject: [PATCH 2/3] Fix minor error --- providers/openai/src/airflow/providers/openai/hooks/openai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/openai/src/airflow/providers/openai/hooks/openai.py b/providers/openai/src/airflow/providers/openai/hooks/openai.py index 6c3758481e7ae..81c120ea40290 100644 --- a/providers/openai/src/airflow/providers/openai/hooks/openai.py +++ b/providers/openai/src/airflow/providers/openai/hooks/openai.py @@ -45,7 +45,7 @@ from openai.types.vector_stores import VectorStoreFile, VectorStoreFileBatch, VectorStoreFileDeleted from airflow.providers.openai.exceptions import OpenAIBatchJobException, OpenAIBatchTimeout -from airflow.providers.facebook.version_compat import BaseHook +from airflow.providers.openai.version_compat import BaseHook class BatchStatus(str, Enum): From 55eb1f170c1c853424d6b6259351fec4f1de0aaf Mon Sep 17 00:00:00 2001 From: Mingdao Yang Date: Sat, 5 Jul 2025 16:39:10 +0800 Subject: [PATCH 3/3] Fix a formatting problem --- providers/openai/src/airflow/providers/openai/hooks/openai.py | 1 - 1 file changed, 1 deletion(-) diff --git a/providers/openai/src/airflow/providers/openai/hooks/openai.py b/providers/openai/src/airflow/providers/openai/hooks/openai.py index 81c120ea40290..5b0566acb0b07 100644 --- a/providers/openai/src/airflow/providers/openai/hooks/openai.py +++ b/providers/openai/src/airflow/providers/openai/hooks/openai.py @@ -44,7 +44,6 @@ ) from openai.types.vector_stores import VectorStoreFile, VectorStoreFileBatch, VectorStoreFileDeleted from airflow.providers.openai.exceptions import OpenAIBatchJobException, OpenAIBatchTimeout - from airflow.providers.openai.version_compat import BaseHook