Allow keeping PostgresHook SQLAlchemy engines on psycopg2

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):
 

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,93 @@ 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 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."""

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):
 

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,93 @@ 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 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."""

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):
 

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,93 @@ 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 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."""

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):
 

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,93 @@ 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 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."""

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):
 

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,93 @@ 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 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."""

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):
 

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,93 @@ 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 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."""

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):
 

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,93 @@ 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 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."""

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):
 

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,93 @@ 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 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."""