Fix false positive when inheriting class that inherits DbApiHook (#…

…16543) closes #14541
apache · Jun 19, 2021 · d3b0669 · d3b0669
1 parent 15c30e1
commit d3b0669
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/airflow/upgrade/rules/db_api_functions.py b/airflow/upgrade/rules/db_api_functions.py
@@ -63,8 +63,9 @@ def get_all_non_dbapi_children():
         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
@@ -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 @@ def get_records(self, sql):
         pass
 
 
-class ProperDbApiHook(DbApiHook):
+class ProperDbApiHook(DbApiHook, BigQueryHook):
     def bulk_dump(self, table, tmp_file):
         pass