Move `BaseHook` class to task SDK #51873

amoghrajesh · 2025-06-18T07:51:56Z

Moving the BaseHook class into task SDK, exactly where it should live similar to other base classes in sdk like notifiers, operators etc.

Testing: Running by all the methods defined on BaseHook -- old vs new

Running with the new path

DAG:

from __future__ import annotations

from contextlib import suppress

from airflow.models.baseoperator import BaseOperator
from airflow import DAG
from airflow.sdk import BaseHook


class CustomOperator(BaseOperator):
    def execute(self, context):
        conn = BaseHook.get_connection("athena_default")
        print("got connection from basehook", conn)


        hook = BaseHook.get_hook("athena_default")
        print("got hook from basehook", hook)

        with suppress(NotImplementedError):
            BaseHook.get_conn(hook)
            print("Raising a NotImplementedError, trying to access get_conn")

        logger = BaseHook.logger()
        print("default logger is", logger)


with DAG("get_connection_basehook", schedule=None, catchup=False) as dag:
    CustomOperator(task_id="set_var")

Running with the older path to check backcompat

from __future__ import annotations

from contextlib import suppress

from airflow.models.baseoperator import BaseOperator
from airflow import DAG
from airflow.hooks.base import BaseHook


class CustomOperator(BaseOperator):
    def execute(self, context):
        conn = BaseHook.get_connection("athena_default")
        print("got connection from basehook", conn)


        hook = BaseHook.get_hook("athena_default")
        print("got hook from basehook", hook)

        with suppress(NotImplementedError):
            BaseHook.get_conn(hook)
            print("Raising a NotImplementedError, trying to access get_conn")

        logger = BaseHook.logger()
        print("default logger is", logger)


with DAG("get_connection_basehook", schedule=None, catchup=False) as dag:
    CustomOperator(task_id="set_var")

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

airflow-core/src/airflow/cli/commands/connection_command.py

kaxil · 2025-06-30T12:12:21Z

airflow-core/src/airflow/cli/commands/connection_command.py

+    # TODO: Revisit this
+    status, message = conn.test_connection()  # type: ignore


Will this work or fail currently -- with this change?

I basically haven't made any change here, just to think a bit loud, this one can use Connection from task SDK right?

Right now its not

kaxil · 2025-06-30T12:16:58Z

providers/airbyte/src/airflow/providers/airbyte/hooks/airbyte.py

+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore


This should be stricter:

Suggested change

from airflow.hooks.base import BaseHook # type: ignore

from airflow.hooks.base import BaseHook # type: ignore[no-redef]

kaxil · 2025-06-30T12:17:07Z

providers/alibaba/src/airflow/providers/alibaba/cloud/hooks/analyticdb_spark.py

+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore


Suggested change

from airflow.hooks.base import BaseHook # type: ignore

from airflow.hooks.base import BaseHook # type: ignore[no-redef]

kaxil · 2025-06-30T12:17:23Z

providers/alibaba/src/airflow/providers/alibaba/cloud/hooks/base_alibaba.py

+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore


Suggested change

from airflow.hooks.base import BaseHook # type: ignore

from airflow.hooks.base import BaseHook # type: ignore[no-redef]

kaxil · 2025-06-30T12:19:13Z

providers/alibaba/src/airflow/providers/alibaba/cloud/hooks/oss.py

@@ -93,7 +97,7 @@ def __init__(self, region: str | None = None, oss_conn_id="oss_default", *args,

    def get_conn(self) -> Connection:
        """Return connection for the hook."""
-        return self.oss_conn
+        return self.oss_conn  # type: ignore[return-value]


This can/should be fixed -- if you change L37 to be Connection from SDK

kaxil · 2025-06-30T12:19:21Z

providers/amazon/src/airflow/providers/amazon/aws/hooks/athena_sql.py

@@ -127,7 +127,10 @@ def conn_config(self) -> AwsConnectionWrapper:
                )

        return AwsConnectionWrapper(
-            conn=connection, region_name=self._region_name, botocore_config=self._config, verify=self._verify
+            conn=connection,  # type: ignore[arg-type]


kaxil · 2025-06-30T12:19:52Z

providers/amazon/src/airflow/providers/amazon/aws/hooks/base_aws.py

+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore


Worth moving to version_compat.py -- same as #52528

kaxil · 2025-06-30T12:20:29Z

providers/amazon/src/airflow/providers/amazon/aws/hooks/chime.py

@@ -69,7 +69,7 @@ def _get_webhook_endpoint(self, conn_id: str) -> str:
        token = conn.password
        if token is None:
            raise AirflowException("Webhook token field is missing and is required.")
-        url = conn.schema + "://" + conn.host
+        url = cast("str", conn.schema) + "://" + cast("str", conn.host)


You can add an error above, since you expect conn.schema and conn.host to be non-None -- similar to token field above

kaxil · 2025-06-30T12:20:55Z

providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker_unified_studio.py

+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore


Strictier typer ignore please for all

amoghrajesh · 2025-06-30T15:33:42Z

@kaxil thanks for the review, pretty positive ill be getting a green run now. Will look at your comments tomorrow!

amoghrajesh added 6 commits June 17, 2025 14:47

Moving BaseHook to task SDK

021df83

changing occurences in providers

3053da9

changing in public interface

5c88ae0

updating occurences in provider tests

b60152e

moving hooks test to task sdk

d2d4ddf

adding backcompat for basehook

8b640cf

amoghrajesh requested review from ashb, kaxil, bolkedebruin, potiuk, gopidesupavan, Fokko, hussein-awala, eladkal, mobuchowski, jscheffl, jedcunningham and o-nikolas as code owners June 18, 2025 07:51

amoghrajesh added 3 commits June 27, 2025 16:52

mypy wars: part 5 - casting

c9f2ff0

mypy wars: part 6 - casting

ff3ec38

mypy wars: part 7

2dc98dd

amoghrajesh requested a review from josh-fell as a code owner June 27, 2025 13:12

amoghrajesh added 8 commits June 27, 2025 19:45

Merge branch 'main' into move-basehook-to-task-sdk

be57175

mypy wars: part 8

0c7bad0

mypy wars: part 9

df8f331

mypy wars: part 10

a95c973

mypy wars: part 11

7c1f64a

fixing patch path for generic transfer

4311513

fixing failing provider tests

e51d5ba

Merge branch 'main' into move-basehook-to-task-sdk

f166ee4

amoghrajesh force-pushed the move-basehook-to-task-sdk branch from 4106b95 to f166ee4 Compare June 30, 2025 07:18

amoghrajesh added 7 commits June 30, 2025 12:50

Merge branch 'main' into move-basehook-to-task-sdk

552389b

final mypy?

b11bf9a

Merge branch 'main' into move-basehook-to-task-sdk

722e764

remote a bad rebase file

00567d1

fixing broken tests

ff8fbde

removing unneccessary str conversions

a0585e2

another mypy fix

5b55d78

kaxil reviewed Jun 30, 2025

View reviewed changes

amoghrajesh added 7 commits June 30, 2025 18:07

fixing compat tests

94256b5

final mypy war!

86bac28

Merge branch 'main' into move-basehook-to-task-sdk

b46bc81

fixing compat

4906f40

compat fix again

3d8c452

fixing tests again

a1f28e1

fixing docker tests

0f20865

Merge branch 'main' into move-basehook-to-task-sdk

34a6b3f

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Move `BaseHook` class to task SDK #51873

Move `BaseHook` class to task SDK #51873

amoghrajesh commented Jun 18, 2025 •

edited

Loading

Uh oh!

Uh oh!

kaxil Jun 30, 2025

Uh oh!

amoghrajesh Jun 30, 2025

Uh oh!

kaxil Jun 30, 2025

Uh oh!

kaxil Jun 30, 2025

Uh oh!

kaxil Jun 30, 2025

Uh oh!

kaxil Jun 30, 2025

Uh oh!

kaxil Jun 30, 2025

Uh oh!

kaxil Jun 30, 2025

Uh oh!

kaxil Jun 30, 2025

Uh oh!

kaxil Jun 30, 2025

Uh oh!

amoghrajesh commented Jun 30, 2025

Uh oh!

Uh oh!

		# TODO: Revisit this
		status, message = conn.test_connection() # type: ignore

	from airflow.hooks.base import BaseHook # type: ignore
	from airflow.hooks.base import BaseHook # type: ignore[no-redef]

Move BaseHook class to task SDK #51873

Are you sure you want to change the base?

Move BaseHook class to task SDK #51873

Conversation

amoghrajesh commented Jun 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Testing: Running by all the methods defined on BaseHook -- old vs new

Running with the new path

Running with the older path to check backcompat

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

amoghrajesh commented Jun 30, 2025

Uh oh!

Uh oh!

Move `BaseHook` class to task SDK #51873

Move `BaseHook` class to task SDK #51873

amoghrajesh commented Jun 18, 2025 •

edited

Loading