Check if sqlalchemy_scheme extra contains forbidden characters (#31984)

apache · Jun 18, 2023 · b683698 · b683698
1 parent b6f9331
commit b683698
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 10 deletions.
diff --git a/airflow/providers/microsoft/mssql/hooks/mssql.py b/airflow/providers/microsoft/mssql/hooks/mssql.py
@@ -63,11 +63,10 @@ def connection_extra_lower(self) -> dict:
     @property
     def sqlalchemy_scheme(self) -> str:
         """Sqlalchemy scheme either from constructor, connection extras or default."""
-        return (
-            self._sqlalchemy_scheme
-            or self.connection_extra_lower.get("sqlalchemy_scheme")
-            or self.DEFAULT_SQLALCHEMY_SCHEME
-        )
+        extra_scheme = self.connection_extra_lower.get("sqlalchemy_scheme")
+        if not self._sqlalchemy_scheme and extra_scheme and (":" in extra_scheme or "/" in extra_scheme):
+            raise RuntimeError("sqlalchemy_scheme in connection extra should not contain : or / characters")
+        return self._sqlalchemy_scheme or extra_scheme or self.DEFAULT_SQLALCHEMY_SCHEME
 
     def get_uri(self) -> str:
         from urllib.parse import parse_qs, urlencode, urlsplit, urlunsplit

diff --git a/airflow/providers/odbc/hooks/odbc.py b/airflow/providers/odbc/hooks/odbc.py
@@ -91,11 +91,10 @@ def database(self) -> str | None:
     @property
     def sqlalchemy_scheme(self) -> str:
         """SQLAlchemy scheme either from constructor, connection extras or default."""
-        return (
-            self._sqlalchemy_scheme
-            or self.connection_extra_lower.get("sqlalchemy_scheme")
-            or self.DEFAULT_SQLALCHEMY_SCHEME
-        )
+        extra_scheme = self.connection_extra_lower.get("sqlalchemy_scheme")
+        if not self._sqlalchemy_scheme and extra_scheme and (":" in extra_scheme or "/" in extra_scheme):
+            raise RuntimeError("sqlalchemy_scheme in connection extra should not contain : or / characters")
+        return self._sqlalchemy_scheme or extra_scheme or self.DEFAULT_SQLALCHEMY_SCHEME
 
     @property
     def connection_extra_lower(self) -> dict: