Allow keeping PostgresHook SQLAlchemy engines on psycopg2

Provider 7.0.0 switched hook-built SQLAlchemy engines to psycopg (v3)
whenever SQLAlchemy 2.x is installed, with no opt-out. The psycopg
dialect renders typed bind casts, so string parameters PostgreSQL
previously coerced implicitly now fail server-side — most visibly
pandas.DataFrame.to_sql into uuid columns (#71977). Honor the existing
DbApiHook sqlalchemy_scheme connection extra (and hook parameter) so
connections that rely on psycopg2 behaviour can keep it.

Author: sehajsandhu123

providers/postgres/docs/changelog.rst

@@ -108,6 +108,9 @@ Breaking changes
     There is no connection or configuration option to keep hooks on psycopg2; the
     ``sql_alchemy_conn`` workarounds above cover the metadata database only. If your Dags rely on
     psycopg2-specific behaviour, test before upgrading or pin the provider below 7.0.0.
+    Starting with provider 7.1.0, the SQLAlchemy engines built by the hook can be kept on
+    psycopg2 by setting the ``sqlalchemy_scheme`` connection extra (or hook parameter) to
+    ``postgresql+psycopg2``.
 
 * ``Make psycopg (v3) the default synchronous Postgres driver (#69526)``
 * ``Switch the default async Postgres driver from asyncpg to psycopg3 (#69089)``

providers/postgres/docs/connections/postgres.rst

@@ -108,6 +108,15 @@ Extra (optional)
       If not specified than hostname from **Connection Host** is used.
     * ``azure_conn_id`` - Azure Connection ID to be used for authentication via Azure Entra ID. Azure Oauth token
       is retrieved from the azure connection which is used as password for PostgreSQL connection. Scope for the Azure OAuth token can be set in the config option ``azure_oauth_scope`` under the section ``[postgres]``. Requires `apache-airflow-providers-microsoft-azure>=12.8.0`.
+    * ``sqlalchemy_scheme`` - The SQLAlchemy drivername used for the URLs the hook builds
+      (``get_uri``, ``get_sqlalchemy_engine``). Must be ``postgresql`` or ``postgresql+<driver>``.
+      Since provider 7.0.0 the hook selects psycopg (v3) whenever SQLAlchemy 2.x is installed;
+      set this to ``postgresql+psycopg2`` to keep SQLAlchemy engines created from this connection
+      on psycopg2. This matters if your Dags rely on psycopg2-specific behaviour — for example,
+      psycopg2 sends string parameters untyped so PostgreSQL implicitly coerces them, while
+      the psycopg (v3) SQLAlchemy dialect renders typed casts, so inserting string values into
+      e.g. ``uuid`` columns (as ``pandas.DataFrame.to_sql`` does) fails with
+      *"column is of type uuid but expression is of type character varying"*.
 
     Example "extras" field (Amazon RDS PostgreSQL or Amazon Aurora PostgreSQL):
 

providers/postgres/src/airflow/providers/postgres/hooks/postgres.py

@@ -135,6 +135,11 @@ class PostgresHook(DbApiHook):
     :param options: Optional. Specifies command-line options to send to the server
         at connection start. For example, setting this to ``-c search_path=myschema``
         sets the session's value of the ``search_path`` to ``myschema``.
+    :param sqlalchemy_scheme: Optional. The SQLAlchemy drivername used for the URLs the hook
+        builds (``get_uri``, ``get_sqlalchemy_engine``), e.g. ``postgresql+psycopg2``. Must be
+        ``postgresql`` or ``postgresql+<driver>``. Defaults to ``postgresql+psycopg`` when
+        psycopg (v3) serves SQLAlchemy 2.x and to ``postgresql`` otherwise. Can also be set via
+        the connection extra ``sqlalchemy_scheme``; this parameter takes precedence.
     :param enable_log_db_messages: Optional. If enabled logs database messages sent to the client
         during the session. To avoid a memory leak psycopg2 only saves the last 50 messages.
         For details, see: `PostgreSQL logging configuration parameters
@@ -164,18 +169,39 @@ class PostgresHook(DbApiHook):
     default_azure_oauth_scope = "https://ossrdbms-aad.database.windows.net/.default"
 
     def __init__(
-        self, *args, options: str | None = None, enable_log_db_messages: bool = False, **kwargs
+        self,
+        *args,
+        options: str | None = None,
+        enable_log_db_messages: bool = False,
+        sqlalchemy_scheme: str | None = None,
+        **kwargs,
     ) -> None:
         super().__init__(*args, **kwargs)
         self.conn: CompatConnection | None = None
         self.database: str | None = kwargs.pop("database", None)
         self.options = options
         self.enable_log_db_messages = enable_log_db_messages
+        self._sqlalchemy_scheme = sqlalchemy_scheme
 
     @staticmethod
     def __cast_nullable(value, dst_type: type) -> Any:
         return dst_type(value) if value is not None else None
 
+    @property
+    def sqlalchemy_scheme(self) -> str:
+        """SQLAlchemy drivername used for the URLs built by this hook."""
+        scheme = self._sqlalchemy_scheme or self.connection.extra_dejson.get("sqlalchemy_scheme")
+        if not scheme:
+            return "postgresql+psycopg" if USE_PSYCOPG3 else "postgresql"
+        if ":" in scheme or "/" in scheme:
+            raise ValueError("The parameter 'sqlalchemy_scheme' must not contain ':' or '/' characters!")
+        if scheme != "postgresql" and not scheme.startswith("postgresql+"):
+            raise ValueError(
+                f"The parameter 'sqlalchemy_scheme' must be 'postgresql' or 'postgresql+<driver>', "
+                f"got: {scheme!r}"
+            )
+        return scheme
+
     @property
     def sqlalchemy_url(self) -> URL:
         try:
@@ -192,7 +218,7 @@ def sqlalchemy_url(self) -> URL:
         if conn.extra_dejson.get("iam", False):
             conn.login, conn.password, conn.port = self.get_iam_token(conn)
         return URL.create(
-            drivername="postgresql+psycopg" if USE_PSYCOPG3 else "postgresql",
+            drivername=self.sqlalchemy_scheme,
             username=self.__cast_nullable(conn.login, str),
             password=self.__cast_nullable(conn.password, str),
             host=self.__cast_nullable(conn.host, str),

providers/postgres/tests/unit/postgres/hooks/test_postgres.py

@@ -633,6 +633,48 @@ def test_get_conn_cursor(self, mocker):
         )
 
 
+class TestPostgresHookSqlalchemyScheme:
+    """Tests for overriding the SQLAlchemy drivername via the sqlalchemy_scheme extra/parameter."""
+
+    @staticmethod
+    def get_hook(extra: dict | None = None, **hook_kwargs) -> PostgresHook:
+        conn = Connection(
+            login="login-conn", password="password-conn", host="host", schema="database", extra=extra
+        )
+        return PostgresHook(connection=conn, **hook_kwargs)
+
+    @pytest.mark.parametrize("scheme", ["postgresql", "postgresql+psycopg2", "postgresql+psycopg"])
+    def test_sqlalchemy_scheme_from_extra(self, scheme):
+        hook = self.get_hook(extra=dict(sqlalchemy_scheme=scheme))
+        expected = f"{scheme}://login-conn:password-conn@host/database"
+        assert hook.sqlalchemy_url.render_as_string(hide_password=False) == expected
+
+    def test_sqlalchemy_scheme_parameter_takes_precedence_over_extra(self):
+        hook = self.get_hook(
+            extra=dict(sqlalchemy_scheme="postgresql"), sqlalchemy_scheme="postgresql+psycopg2"
+        )
+        expected = "postgresql+psycopg2://login-conn:password-conn@host/database"
+        assert hook.sqlalchemy_url.render_as_string(hide_password=False) == expected
+
+    def test_get_uri_with_sqlalchemy_scheme(self):
+        hook = self.get_hook(extra=dict(sqlalchemy_scheme="postgresql+psycopg2"))
+        assert hook.get_uri() == "postgresql+psycopg2://login-conn:password-conn@host/database"
+
+    @pytest.mark.parametrize("scheme", ["mysql", "mysql+pymysql", "postgres+psycopg2"])
+    def test_sqlalchemy_scheme_with_wrong_dialect(self, scheme):
+        hook = self.get_hook(extra=dict(sqlalchemy_scheme=scheme))
+        with pytest.raises(
+            ValueError, match="'sqlalchemy_scheme' must be 'postgresql' or 'postgresql\\+<driver>'"
+        ):
+            hook.sqlalchemy_url
+
+    @pytest.mark.parametrize("scheme", ["postgresql+psycopg2://malicious", "postgresql+psycopg2/malicious"])
+    def test_sqlalchemy_scheme_with_forbidden_characters(self, scheme):
+        hook = self.get_hook(extra=dict(sqlalchemy_scheme=scheme))
+        with pytest.raises(ValueError, match="must not contain ':' or '/' characters"):
+            hook.sqlalchemy_url
+
+
 @pytest.mark.backend("postgres")
 class TestPostgresHook:
     """Tests that are identical between psycopg2 and psycopg3."""