Skip to content

Commit 104f525

Browse files
authored
Move all BaseHook usages to version_compat in Microsoft PSRP (#52854)
1 parent 12f500b commit 104f525

File tree

2 files changed

+9
-5
lines changed

2 files changed

+9
-5
lines changed

providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/hooks/psrp.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,7 @@
3030
from pypsrp.wsman import WSMan
3131

3232
from airflow.exceptions import AirflowException
33-
34-
try:
35-
from airflow.sdk import BaseHook
36-
except ImportError:
37-
from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef]
33+
from airflow.providers.microsoft.psrp.version_compat import BaseHook
3834

3935
INFORMATIONAL_RECORD_LEVEL_MAP = {
4036
MessageType.DEBUG_RECORD: DEBUG,

providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/version_compat.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,13 +28,21 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
2828

2929

3030
AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)
31+
AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0)
3132

3233
if AIRFLOW_V_3_0_PLUS:
3334
from airflow.sdk import BaseOperator
3435
else:
3536
from airflow.models import BaseOperator # type: ignore[no-redef]
3637

38+
if AIRFLOW_V_3_1_PLUS:
39+
from airflow.sdk import BaseHook
40+
else:
41+
from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef]
42+
3743
__all__ = [
3844
"AIRFLOW_V_3_0_PLUS",
45+
"AIRFLOW_V_3_1_PLUS",
3946
"BaseOperator",
47+
"BaseHook",
4048
]

0 commit comments

Comments
 (0)