Fix false positive when inheriting class that inherits ``DbApiHook`` (#16543)
closes https://github.com/apache/airflow/issues/14541
diff --git a/airflow/upgrade/rules/db_api_functions.py b/airflow/upgrade/rules/db_api_functions.py
index 0ca730f..dbf72d6 100644
--- a/airflow/upgrade/rules/db_api_functions.py
+++ b/airflow/upgrade/rules/db_api_functions.py
@@ -63,8 +63,9 @@
next_generation = []
for child in basehook_children:
subclasses = child.__subclasses__()
- if subclasses:
- next_generation.extend(subclasses)
+ for subclass in subclasses:
+ if all(base_class.__name__ != 'DbApiHook' for base_class in subclass.__bases__):
+ next_generation.append(subclass)
res.extend(next_generation)
basehook_children = next_generation
return res
diff --git a/tests/upgrade/rules/test_db_api_functions.py b/tests/upgrade/rules/test_db_api_functions.py
index d73a041..156e7d5 100644
--- a/tests/upgrade/rules/test_db_api_functions.py
+++ b/tests/upgrade/rules/test_db_api_functions.py
@@ -18,6 +18,7 @@
from airflow.hooks.base_hook import BaseHook
from airflow.hooks.dbapi_hook import DbApiHook
+from airflow.contrib.hooks.bigquery_hook import BigQueryHook
from airflow.upgrade.rules.db_api_functions import DbApiRule
@@ -41,7 +42,7 @@
pass
-class ProperDbApiHook(DbApiHook):
+class ProperDbApiHook(DbApiHook, BigQueryHook):
def bulk_dump(self, table, tmp_file):
pass
