Bring back mypy checks for new-structure providers (apache#45815)

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

@@ -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

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(

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)
 

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)