This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 2111287e087 Allow keeping PostgresHook SQLAlchemy engines on psycopg2
(#72000)
2111287e087 is described below
commit 2111287e0875e17b962833f9c8335ff7c6435813
Author: Blahaj <[email protected]>
AuthorDate: Tue Sep 22 17:50:18 2026 +0530
Allow keeping PostgresHook SQLAlchemy engines on psycopg2 (#72000)
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.
---
providers/postgres/docs/connections/postgres.rst | 9 +++
.../airflow/providers/postgres/hooks/postgres.py | 30 +++++++-
.../tests/unit/postgres/hooks/test_postgres.py | 87 ++++++++++++++++++++++
3 files changed, 124 insertions(+), 2 deletions(-)
diff --git a/providers/postgres/docs/connections/postgres.rst
b/providers/postgres/docs/connections/postgres.rst
index 3018769afbc..478fcc63085 100644
--- a/providers/postgres/docs/connections/postgres.rst
+++ b/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 without a type 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):
diff --git
a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py
b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py
index 89e972a3f82..7f8064c9b2b 100644
--- a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py
+++ b/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 @@ class PostgresHook(DbApiHook):
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),
diff --git a/providers/postgres/tests/unit/postgres/hooks/test_postgres.py
b/providers/postgres/tests/unit/postgres/hooks/test_postgres.py
index 12fe869a212..55fcb03d106 100644
--- a/providers/postgres/tests/unit/postgres/hooks/test_postgres.py
+++ b/providers/postgres/tests/unit/postgres/hooks/test_postgres.py
@@ -633,6 +633,93 @@ class TestPostgresHookConnPPG3:
)
+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
+
+
[email protected]("postgres")
+class TestPostgresHookPandasToSqlUuid:
+ """DataFrame.to_sql with string values into a uuid column: fails on
psycopg3 because SQLAlchemy
+ renders typed bind casts, works when the connection opts back into
psycopg2 via sqlalchemy_scheme."""
+
+ table = "test_pandas_to_sql_uuid_table"
+
+ PSYCOPG3_XFAIL_REASON = (
+ "The psycopg3 SQLAlchemy dialect renders typed bind casts, so string
params fail with "
+ "'column is of type uuid but expression is of type character varying'
instead of being "
+ "implicitly coerced as under psycopg2. Tracked upstream in "
+ "https://github.com/pandas-dev/pandas/issues/63511, "
+ "https://github.com/apache/arrow/pull/50325, "
+ "https://github.com/sqlalchemy/sqlalchemy/discussions/10839 and "
+ "https://github.com/sqlalchemy/sqlalchemy/issues/12060";
+ )
+
+ def teardown_method(self):
+ with PostgresHook().get_conn() as conn:
+ with conn.cursor() as cur:
+ cur.execute(f"DROP TABLE IF EXISTS {self.table}")
+
+ def insert_string_uuid_df(self, hook: PostgresHook) -> None:
+ engine = hook.get_sqlalchemy_engine()
+ with engine.begin() as conn:
+ conn.execute(sqlalchemy.text(f"CREATE TABLE {self.table} (id UUID
PRIMARY KEY, item TEXT)"))
+ df = pd.DataFrame({"id": ["b29d6cda-04ca-4073-8ef8-4f70d35e41e2"],
"item": ["laptop"]})
+ df.to_sql(self.table, engine, if_exists="append", index=False)
+
+ @pytest.mark.skipif(not USE_PSYCOPG3, reason="psycopg v3 or sqlalchemy v2
not available")
+ @pytest.mark.xfail(
+ raises=(sqlalchemy.exc.ProgrammingError, pd.errors.DatabaseError),
+ strict=True,
+ reason=PSYCOPG3_XFAIL_REASON,
+ )
+ def test_to_sql_string_uuid_fails_on_psycopg3(self):
+ self.insert_string_uuid_df(PostgresHook())
+
+ def test_to_sql_string_uuid_works_with_psycopg2_scheme(self):
+ hook = PostgresHook(sqlalchemy_scheme="postgresql+psycopg2")
+ self.insert_string_uuid_df(hook)
+ with hook.get_sqlalchemy_engine().connect() as conn:
+ assert conn.execute(sqlalchemy.text(f"SELECT COUNT(*) FROM
{self.table}")).scalar() == 1
+
+
@pytest.mark.backend("postgres")
class TestPostgresHook:
"""Tests that are identical between psycopg2 and psycopg3."""
