diff --git a/providers/postgres/README.rst b/providers/postgres/README.rst index 4e6f61a59f5ea..2a2a64e1e7ba9 100644 --- a/providers/postgres/README.rst +++ b/providers/postgres/README.rst @@ -50,17 +50,15 @@ The package supports the following python versions: 3.10,3.11,3.12,3.13,3.14 Requirements ------------ -========================================== ====================================== +========================================== ===================================== PIP package Version required -========================================== ====================================== +========================================== ===================================== ``apache-airflow`` ``>=2.11.0`` ``apache-airflow-providers-common-compat`` ``>=1.12.0`` ``apache-airflow-providers-common-sql`` ``>=1.32.0`` -``psycopg2-binary`` ``>=2.9.9; python_version < "3.13"`` -``psycopg2-binary`` ``>=2.9.10; python_version >= "3.13"`` ``psycopg[binary]`` ``>=3.2.9; python_version < "3.14"`` ``psycopg[binary]`` ``>=3.3.3; python_version >= "3.14"`` -========================================== ====================================== +========================================== ===================================== Optional cross provider package dependencies -------------------------------------------- @@ -95,6 +93,7 @@ Extra Dependencies ``openlineage`` ``apache-airflow-providers-openlineage`` ``pandas`` ``pandas>=2.1.2; python_version <"3.13"``, ``pandas>=2.2.3; python_version >="3.13" and python_version <"3.14"``, ``pandas>=2.3.3; python_version >="3.14"`` ``polars`` ``polars>=1.26.0`` +``psycopg2`` ``psycopg2-binary>=2.9.9; python_version < "3.13"``, ``psycopg2-binary>=2.9.10; python_version >= "3.13"`` ``sqlalchemy`` ``sqlalchemy>=1.4.54`` =================== ============================================================================================================================================================ diff --git a/providers/postgres/docs/changelog.rst b/providers/postgres/docs/changelog.rst index eff74f3268483..4cfd7f38668fa 100644 --- a/providers/postgres/docs/changelog.rst +++ b/providers/postgres/docs/changelog.rst @@ -42,6 +42,17 @@ Breaking changes ``apache-airflow-providers-postgres[asyncpg]`` and set ``[database] sql_alchemy_conn_async = postgresql+asyncpg://...`` explicitly. +.. note:: + The default synchronous metadata-database driver is now ``psycopg`` (psycopg3), mirroring the + async default above. ``psycopg2-binary`` is no longer installed by default; it moved from a hard + dependency to the new ``[psycopg2]`` optional extra. + + A bare ``postgresql://`` or legacy ``postgres://`` / ``postgres+psycopg2://`` ``sql_alchemy_conn`` + is now rewritten to ``postgresql+psycopg://`` instead of ``postgresql+psycopg2://``. An explicit + ``postgresql+psycopg2://`` connection string is never rewritten. To keep using psycopg2, install + ``apache-airflow-providers-postgres[psycopg2]`` and keep (or set) + ``[database] sql_alchemy_conn = postgresql+psycopg2://...`` explicitly. + 6.8.0 ..... diff --git a/providers/postgres/docs/index.rst b/providers/postgres/docs/index.rst index b10aa40df7646..91c554f999fe0 100644 --- a/providers/postgres/docs/index.rst +++ b/providers/postgres/docs/index.rst @@ -98,17 +98,15 @@ Requirements The minimum Apache Airflow version supported by this provider distribution is ``2.11.0``. -========================================== ====================================== +========================================== ===================================== PIP package Version required -========================================== ====================================== +========================================== ===================================== ``apache-airflow`` ``>=2.11.0`` ``apache-airflow-providers-common-compat`` ``>=1.12.0`` ``apache-airflow-providers-common-sql`` ``>=1.32.0`` -``psycopg2-binary`` ``>=2.9.9; python_version < "3.13"`` -``psycopg2-binary`` ``>=2.9.10; python_version >= "3.13"`` ``psycopg[binary]`` ``>=3.2.9; python_version < "3.14"`` ``psycopg[binary]`` ``>=3.3.3; python_version >= "3.14"`` -========================================== ====================================== +========================================== ===================================== Optional cross provider package dependencies -------------------------------------------- diff --git a/providers/postgres/pyproject.toml b/providers/postgres/pyproject.toml index 8adda0b57802b..e401c3a6d379f 100644 --- a/providers/postgres/pyproject.toml +++ b/providers/postgres/pyproject.toml @@ -62,10 +62,6 @@ dependencies = [ "apache-airflow>=2.11.0", "apache-airflow-providers-common-compat>=1.12.0", "apache-airflow-providers-common-sql>=1.32.0", - # psycopg2 remains the sync driver for now; its removal and the sync migration to - # psycopg3 are tracked at https://github.com/apache/airflow/issues/68453 - "psycopg2-binary>=2.9.9; python_version < '3.13'", - "psycopg2-binary>=2.9.10; python_version >= '3.13'", "psycopg[binary]>=3.2.9; python_version < '3.14'", "psycopg[binary]>=3.3.3; python_version >= '3.14'", ] @@ -93,6 +89,10 @@ dependencies = [ "polars" = [ "polars>=1.26.0" ] +"psycopg2" = [ + "psycopg2-binary>=2.9.9; python_version < '3.13'", + "psycopg2-binary>=2.9.10; python_version >= '3.13'", +] "sqlalchemy" = [ "sqlalchemy>=1.4.54" ] diff --git a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py index be6519ddade24..3821779e261fd 100644 --- a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py +++ b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py @@ -18,14 +18,12 @@ from __future__ import annotations import os -from collections.abc import Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping from contextlib import closing from copy import deepcopy -from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias, cast, overload +from typing import TYPE_CHECKING, Any, Literal, NoReturn, Protocol, TypeAlias, cast, overload from more_itertools import chunked -from psycopg2 import connect as ppg2_connect -from psycopg2.extras import DictCursor, NamedTupleCursor, RealDictCursor, execute_values from airflow.providers.common.compat.sdk import ( AirflowException, @@ -54,9 +52,35 @@ from psycopg.rows import dict_row, namedtuple_row from psycopg.types.json import register_default_adapters +try: + import psycopg2 as _psycopg2 + import psycopg2.extras as _psycopg2_extras +except (ImportError, ModuleNotFoundError): + _psycopg2 = None + _psycopg2_extras = None + +ppg2_connect: Callable[..., Any] | None = _psycopg2.connect if _psycopg2 else None +DictCursor: type | None = _psycopg2_extras.DictCursor if _psycopg2_extras else None +NamedTupleCursor: type | None = _psycopg2_extras.NamedTupleCursor if _psycopg2_extras else None +RealDictCursor: type | None = _psycopg2_extras.RealDictCursor if _psycopg2_extras else None +execute_values: Callable[..., Any] | None = _psycopg2_extras.execute_values if _psycopg2_extras else None + + +def _require_psycopg2() -> NoReturn: + raise AirflowOptionalProviderFeatureException( + "psycopg2 is not installed. Please install it with " + "`pip install apache-airflow-providers-postgres[psycopg2]`." + ) + + if TYPE_CHECKING: from pandas import DataFrame as PandasDataFrame from polars import DataFrame as PolarsDataFrame + from psycopg2.extras import ( + DictCursor as _DictCursorType, + NamedTupleCursor as _NamedTupleCursorType, + RealDictCursor as _RealDictCursorType, + ) from sqlalchemy.engine import URL from airflow.providers.common.sql.dialects.dialect import Dialect @@ -65,7 +89,7 @@ if USE_PSYCOPG3: from psycopg.errors import Diagnostic - CursorType: TypeAlias = DictCursor | RealDictCursor | NamedTupleCursor + CursorType: TypeAlias = _DictCursorType | _RealDictCursorType | _NamedTupleCursorType CursorRow: TypeAlias = dict[str, Any] | tuple[Any, ...] @@ -204,6 +228,9 @@ def _get_cursor(self, raw_cursor: str) -> CursorType: valid_cursors = "dictcursor, namedtuplecursor" raise ValueError(f"Invalid cursor passed {_cursor}. Valid options are: {valid_cursors}") + if DictCursor is None: + _require_psycopg2() + cursor_types = { "dictcursor": DictCursor, "realdictcursor": RealDictCursor, @@ -235,6 +262,9 @@ def _create_connection(self, conn_args: dict[str, Any]) -> CompatConnection: return connection + if ppg2_connect is None: + _require_psycopg2() + return ppg2_connect(**conn_args) def _generate_cursor_name(self): @@ -691,6 +721,8 @@ def insert_rows( ) # if fast_executemany is enabled with psycopg2, use optimized execute_values from psycopg + if execute_values is None: + _require_psycopg2() self._insert_statement_format = "INSERT INTO {} {} VALUES %s" nb_rows = 0 diff --git a/providers/postgres/tests/unit/postgres/hooks/test_postgres.py b/providers/postgres/tests/unit/postgres/hooks/test_postgres.py index e8c1dd5eff9d0..a9eff1bcb6f6a 100644 --- a/providers/postgres/tests/unit/postgres/hooks/test_postgres.py +++ b/providers/postgres/tests/unit/postgres/hooks/test_postgres.py @@ -58,6 +58,27 @@ import psycopg2.extras +def test_hooks_postgres_raises_clear_error_without_psycopg2(monkeypatch): + """PostgresHook must fail loudly (not with a bare ImportError or a real connection attempt) + when psycopg2-specific functionality is used but psycopg2 isn't installed.""" + import airflow.providers.postgres.hooks.postgres as postgres_module + + monkeypatch.setattr(postgres_module, "USE_PSYCOPG3", False) + monkeypatch.setattr(postgres_module, "ppg2_connect", None) + monkeypatch.setattr(postgres_module, "DictCursor", None) + monkeypatch.setattr(postgres_module, "RealDictCursor", None) + monkeypatch.setattr(postgres_module, "NamedTupleCursor", None) + monkeypatch.setattr(postgres_module, "execute_values", None) + + hook = postgres_module.PostgresHook.__new__(postgres_module.PostgresHook) + with pytest.raises(AirflowOptionalProviderFeatureException, match="psycopg2 is not installed"): + hook._get_cursor("dictcursor") + with pytest.raises(AirflowOptionalProviderFeatureException, match="psycopg2 is not installed"): + hook._create_connection({}) + with pytest.raises(AirflowOptionalProviderFeatureException, match="psycopg2 is not installed"): + hook.insert_rows(table="t", rows=[(1,)], fast_executemany=True) + + @pytest.fixture def mock_connect(mocker): """Mock the connection object according to the correct psycopg version.""" diff --git a/uv.lock b/uv.lock index 11bbace1114d2..69b4c71baf8b0 100644 --- a/uv.lock +++ b/uv.lock @@ -7151,7 +7151,6 @@ dependencies = [ { name = "apache-airflow-providers-common-compat" }, { name = "apache-airflow-providers-common-sql" }, { name = "psycopg", extra = ["binary"] }, - { name = "psycopg2-binary" }, ] [package.optional-dependencies] @@ -7173,6 +7172,9 @@ pandas = [ polars = [ { name = "polars" }, ] +psycopg2 = [ + { name = "psycopg2-binary" }, +] sqlalchemy = [ { name = "sqlalchemy" }, ] @@ -7208,11 +7210,11 @@ requires-dist = [ { name = "polars", marker = "extra == 'polars'", specifier = ">=1.26.0" }, { name = "psycopg", extras = ["binary"], marker = "python_full_version < '3.14'", specifier = ">=3.2.9" }, { name = "psycopg", extras = ["binary"], marker = "python_full_version >= '3.14'", specifier = ">=3.3.3" }, - { name = "psycopg2-binary", marker = "python_full_version < '3.13'", specifier = ">=2.9.9" }, - { name = "psycopg2-binary", marker = "python_full_version >= '3.13'", specifier = ">=2.9.10" }, + { name = "psycopg2-binary", marker = "python_full_version >= '3.13' and extra == 'psycopg2'", specifier = ">=2.9.10" }, + { name = "psycopg2-binary", marker = "python_full_version < '3.13' and extra == 'psycopg2'", specifier = ">=2.9.9" }, { name = "sqlalchemy", marker = "extra == 'sqlalchemy'", specifier = ">=1.4.54" }, ] -provides-extras = ["amazon", "asyncpg", "microsoft-azure", "openlineage", "pandas", "polars", "sqlalchemy"] +provides-extras = ["amazon", "asyncpg", "microsoft-azure", "openlineage", "psycopg2", "pandas", "polars", "sqlalchemy"] [package.metadata.requires-dev] dev = [