From 7310d8057865750938d1b12ddb961c70be30ebe8 Mon Sep 17 00:00:00 2001
From: Jarek Potiuk <jarek@potiuk.com>
Date: Mon, 20 Jan 2025 21:32:15 +0100
Subject: [PATCH] Bring back mypy checks for new-structure providers

The providers moved to the news structure have not been chedked
by mypy checks when run in "canary" or "full tests needed" builds. This
change adds possibility to pass multiple folders to mypy check and adds
special "all_new_providers" special argument that gets automatically
resolved to all new provider folders.

Also few mypy checks started to appear and they are fixed now.
---
 .pre-commit-config.yaml                       |  4 +-
 .../providers/common/sql/dialects/dialect.py  |  3 +-
 .../airflow/providers/common/sql/hooks/sql.py |  4 +-
 scripts/ci/pre_commit/mypy_folder.py          | 62 +++++++++++--------
 4 files changed, 43 insertions(+), 30 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index bf2d3feca435a..4d124c88e4aa9 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1343,14 +1343,14 @@ repos:
         name: Run mypy for providers
         language: python
         entry: ./scripts/ci/pre_commit/mypy.py --namespace-packages
-        files: ^providers/src/airflow/providers/.*\.py$|^providers/tests//.*\.py$
+        files: ^providers/src/airflow/providers/.*\.py$|^providers/tests//.*\.py$|^providers/.*/src/.*\.py$|^providers/.*/tests/.*\.py$
         require_serial: true
         additional_dependencies: ['rich>=12.4.4']
       - id: mypy-providers
         stages: ['manual']
         name: Run mypy for providers (manual)
         language: python
-        entry: ./scripts/ci/pre_commit/mypy_folder.py providers/src/airflow/providers
+        entry: ./scripts/ci/pre_commit/mypy_folder.py providers/src/airflow/providers all_new_providers
         pass_filenames: false
         files: ^.*\.py$
         require_serial: true
diff --git a/providers/src/airflow/providers/common/sql/dialects/dialect.py b/providers/src/airflow/providers/common/sql/dialects/dialect.py
index 184e6a5ce4e4c..0d4cc3d11e811 100644
--- a/providers/src/airflow/providers/common/sql/dialects/dialect.py
+++ b/providers/src/airflow/providers/common/sql/dialects/dialect.py
@@ -49,6 +49,7 @@ def __init__(self, hook, **kwargs) -> None:
     def remove_quotes(cls, value: str | None) -> str | None:
         if value:
             return cls.pattern.sub(r"\1", value)
+        return value
 
     @property
     def placeholder(self) -> str:
@@ -73,7 +74,7 @@ def _escape_column_name_format(self) -> str:
     @classmethod
     def extract_schema_from_table(cls, table: str) -> tuple[str, str | None]:
         parts = table.split(".")
-        return tuple(parts[::-1]) if len(parts) == 2 else (table, None)
+        return tuple(parts[::-1]) if len(parts) == 2 else (table, None)  # type: ignore[return-value]
 
     @lru_cache(maxsize=None)
     def get_column_names(
diff --git a/providers/src/airflow/providers/common/sql/hooks/sql.py b/providers/src/airflow/providers/common/sql/hooks/sql.py
index 25d25eaec1786..fec2b81ec128d 100644
--- a/providers/src/airflow/providers/common/sql/hooks/sql.py
+++ b/providers/src/airflow/providers/common/sql/hooks/sql.py
@@ -324,8 +324,8 @@ def dialect(self) -> Dialect:
                 return import_string(dialect_info["dialect_class_name"])(self)
             except ImportError:
                 raise AirflowOptionalProviderFeatureException(
-                    f"{dialect_info.dialect_class_name} not found, run: pip install "
-                    f"'{dialect_info.provider_name}'."
+                    f"{dialect_info['dialect_class_name']} not found, run: pip install "
+                    f"'{dialect_info['provider_name']}'."
                 )
         return Dialect(self)
 
diff --git a/scripts/ci/pre_commit/mypy_folder.py b/scripts/ci/pre_commit/mypy_folder.py
index 36a7840fff9b8..30b7809b2c214 100755
--- a/scripts/ci/pre_commit/mypy_folder.py
+++ b/scripts/ci/pre_commit/mypy_folder.py
@@ -25,51 +25,63 @@
 
 from common_precommit_utils import (
     console,
+    get_all_new_provider_ids,
     initialize_breeze_precommit,
     run_command_via_breeze_shell,
 )
 
 initialize_breeze_precommit(__name__, __file__)
 
+
 ALLOWED_FOLDERS = [
     "airflow",
     "providers/src/airflow/providers",
+    *[f"providers/{provider_id.replace('.', '/')}/src" for provider_id in get_all_new_provider_ids()],
     "dev",
     "docs",
     "task_sdk/src/airflow/sdk",
+    # TODO(potiuk): rename it to "all_providers" when we move all providers to new structure
+    "all_new_providers",
 ]
 
 if len(sys.argv) < 2:
     console.print(f"[yellow]You need to specify the folder to test as parameter: {ALLOWED_FOLDERS}\n")
     sys.exit(1)
 
-mypy_folder = sys.argv[1]
-if mypy_folder not in ALLOWED_FOLDERS:
-    console.print(f"[yellow]Wrong folder {mypy_folder}. It should be one of those: {ALLOWED_FOLDERS}\n")
-    sys.exit(1)
+mypy_folders = sys.argv[1:]
 
-arguments = [mypy_folder]
-if mypy_folder == "providers/src/airflow/providers":
-    arguments.extend(
-        [
-            "providers/tests",
-            "--namespace-packages",
-        ]
-    )
-if mypy_folder == "task_sdk/src/airflow/sdk":
-    arguments.extend(
-        [
-            "task_sdk/tests",
-            "--namespace-packages",
-        ]
-    )
+for mypy_folder in mypy_folders:
+    if mypy_folder not in ALLOWED_FOLDERS:
+        console.print(
+            f"\n[red]ERROR: Folder `{mypy_folder}` is wrong.[/]\n\n"
+            f"All folders passed should be one of those: {ALLOWED_FOLDERS}\n"
+        )
+        sys.exit(1)
 
-if mypy_folder == "airflow":
-    arguments.extend(
-        [
-            "tests",
-        ]
-    )
+arguments = mypy_folders.copy()
+namespace_packages = False
+
+for mypy_folder in mypy_folders:
+    if mypy_folder == "all_new_providers":
+        arguments.remove("all_new_providers")
+        for provider_id in get_all_new_provider_ids():
+            arguments.append(f"providers/{provider_id.replace('.', '/')}/src")
+            arguments.append(f"providers/{provider_id.replace('.', '/')}/tests")
+        namespace_packages = True
+    if mypy_folder == "providers/src/airflow/providers":
+        arguments.append("providers/tests")
+        namespace_packages = True
+    elif mypy_folder.startswith("providers/"):
+        arguments.append(f"{Path(mypy_folder).parent.as_posix()}/tests")
+        namespace_packages = True
+    if mypy_folder == "task_sdk/src/airflow/sdk":
+        arguments.append("task_sdk/tests")
+        namespace_packages = True
+    if mypy_folder == "airflow":
+        arguments.append("tests")
+
+if namespace_packages:
+    arguments.append("--namespace-packages")
 
 print("Running /opt/airflow/scripts/in_container/run_mypy.sh with arguments: ", arguments)
 
