add support for psycopg3 syntax

.github/workflows/test.yml

@@ -21,6 +21,10 @@ jobs:
       matrix:
         python-version: ["3.10", "3.11", "3.12", "3.13"]
         tox-env: ["airflow2", "airflow3"]
+        include:
+          # Airflow 2 does not support Python 3.14
+          - python-version: "3.14"
+            tox-env: "airflow3"
 
     steps:
       - uses: actions/checkout@v4
@@ -40,7 +44,7 @@ jobs:
         run: .venv/bin/tox -e ${{ matrix.tox-env }}
 
       - name: Upload coverage to Codecov
-        if: matrix.python-version == '3.13' && matrix.tox-env == 'airflow3'
+        if: matrix.python-version == '3.14' && matrix.tox-env == 'airflow3'
         uses: codecov/codecov-action@v4
         with:
           files: coverage.xml

README.md

@@ -3,10 +3,10 @@
 | Category    | Badges                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        |
 |-------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
 | **License** | [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)                                                                                                                                                                                                                                                                                                                                                                                            |
-| **PyPI**    | [![python](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12%20%7C%203.13-blue.svg)](https://www.python.org/downloads/) [![airflow](https://img.shields.io/badge/airflow-2.9%2B%20%7C%203.x-blue.svg)](https://airflow.apache.org/) [![PyPI](https://img.shields.io/pypi/v/airflow-postgres-csv)](https://pypi.org/project/airflow-postgres-csv/) [![Downloads](https://img.shields.io/pypi/dm/airflow-postgres-csv)](https://pypi.org/project/airflow-postgres-csv/)                        |
+| **PyPI**    | [![python](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12%20%7C%203.13%20%7C%203.14-blue.svg)](https://www.python.org/downloads/) [![airflow](https://img.shields.io/badge/airflow-2.9%2B%20%7C%203.x-blue.svg)](https://airflow.apache.org/) [![PyPI](https://img.shields.io/pypi/v/airflow-postgres-csv)](https://pypi.org/project/airflow-postgres-csv/) [![Downloads](https://img.shields.io/pypi/dm/airflow-postgres-csv)](https://pypi.org/project/airflow-postgres-csv/)                        |
 | **CI**      | [![lint](https://github.com/Redevil10/airflow-postgres-csv/actions/workflows/lint.yml/badge.svg)](https://github.com/Redevil10/airflow-postgres-csv/actions/workflows/lint.yml) [![tests](https://github.com/Redevil10/airflow-postgres-csv/actions/workflows/test.yml/badge.svg)](https://github.com/Redevil10/airflow-postgres-csv/actions/workflows/test.yml) [![codecov](https://codecov.io/github/Redevil10/airflow-postgres-csv/graph/badge.svg)](https://codecov.io/gh/Redevil10/airflow-postgres-csv) |
 
-Airflow operators for bulk PostgreSQL <-> CSV transfers using `COPY`. Supports Airflow 2.9+ and Airflow 3.
+Airflow operators for bulk PostgreSQL <-> CSV transfers using `COPY`. Supports Airflow 2.9+ and Airflow 3, with both `psycopg2` and `psycopg` (v3) drivers.
 
 Listed on the [Apache Airflow Ecosystem](https://airflow.apache.org/ecosystem/) page.
 
@@ -92,7 +92,7 @@ CsvToPostgresOperator(
 | `conn_id` | Airflow Postgres connection ID | required |
 | `csv_file_path` | Output file path (templated) | required |
 | `sql` | SQL query string, or path to `.sql` file | required |
-| `parameters` | Dict passed to `cursor.mogrify` | `{}` |
+| `parameters` | Dict of query parameters, bound client-side via `mogrify` | `{}` |
 | `has_header` | Include CSV header row | `True` |
 | `compression` | Compression format (`"gzip"` or `None`) | `None` |
 | `timeout` | Query timeout in minutes | `60` |
@@ -137,7 +137,13 @@ Each environment installs the correct Airflow and provider versions automaticall
 |---|---|---|
 | `apache-airflow` | `>=2.9, <3.0` | `>=3.0` |
 | `apache-airflow-providers-postgres` | `>=5.0, <6.0` | `>=6.0` |
-| Python | 3.10 – 3.13 | 3.10 – 3.13 |
+| Python | 3.10 – 3.13 | 3.10 – 3.14 |
+
+### psycopg2 vs psycopg3
+
+Recent versions of `apache-airflow-providers-postgres` make `PostgresHook` return a
+`psycopg` (v3) connection when `psycopg` and SQLAlchemy 2 are installed, and a `psycopg2`
+connection otherwise. Both operators detect the driver automatically and work with either.
 
 ## License
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "airflow-postgres-csv"
-version = "0.4.0"
+version = "0.4.1"
 description = "Airflow operators for PostgreSQL <-> CSV file transfers using COPY"
 readme = "README.md"
 license = "Apache-2.0"

src/airflow_postgres_csv/__init__.py

@@ -3,4 +3,4 @@
 from airflow_postgres_csv.operators import CsvToPostgresOperator, PostgresToCsvOperator
 
 __all__ = ["PostgresToCsvOperator", "CsvToPostgresOperator"]
-__version__ = "0.4.0"
+__version__ = "0.4.1"

src/airflow_postgres_csv/operators.py

@@ -1,8 +1,10 @@
 """Custom Airflow operators for PostgreSQL <-> CSV file transfers."""
 
+import codecs
 import gzip
 import os
-from collections.abc import Sequence
+from collections.abc import Collection, Sequence
+from typing import TYPE_CHECKING, cast
 
 from airflow.exceptions import AirflowException
 from airflow.providers.postgres.hooks.postgres import PostgresHook
@@ -12,6 +14,55 @@
 except ImportError:
     from airflow.models import BaseOperator  # Airflow 2
 
+if TYPE_CHECKING:
+    from typing_extensions import LiteralString
+
+_COPY_CHUNK_SIZE = 1024 * 1024
+
+
+def _is_psycopg2(cursor) -> bool:
+    """PostgresHook returns psycopg2 or psycopg3 connections depending on the environment."""
+    return hasattr(cursor, "copy_expert")
+
+
+def _set_statement_timeout(cursor, timeout_minutes: int) -> None:
+    # ``SET`` cannot take bind parameters with psycopg3 server-side binding;
+    # ``set_config`` works with both psycopg2 and psycopg3.
+    cursor.execute(
+        "SELECT set_config('statement_timeout', %s, false)",
+        (str(int(timeout_minutes * 60 * 1000)),),
+    )
+
+
+def _mogrify(conn, cursor, sql: str, parameters) -> str:
+    if _is_psycopg2(cursor):
+        result = cursor.mogrify(sql, parameters)
+    else:
+        import psycopg
+
+        result = psycopg.ClientCursor(conn).mogrify(cast("LiteralString", sql), parameters)
+    return result.decode("utf-8") if isinstance(result, bytes) else result
+
+
+def _copy_to_file(conn, cursor, copy_command: str, file) -> None:
+    if _is_psycopg2(cursor):
+        cursor.copy_expert(copy_command, file)
+        return
+    decoder = codecs.getincrementaldecoder(conn.info.encoding)()
+    with cursor.copy(copy_command) as copy:
+        for data in copy:
+            file.write(decoder.decode(bytes(data)))
+    file.write(decoder.decode(b"", final=True))
+
+
+def _copy_from_file(cursor, copy_command: str, file) -> None:
+    if _is_psycopg2(cursor):
+        cursor.copy_expert(copy_command, file)
+        return
+    with cursor.copy(copy_command) as copy:
+        while data := file.read(_COPY_CHUNK_SIZE):
+            copy.write(data)
+
 
 class PostgresToCsvOperator(BaseOperator):
     """
@@ -29,7 +80,7 @@ class PostgresToCsvOperator(BaseOperator):
         - Absolute path to a ``.sql`` file (e.g., ``"/opt/sql/query.sql"``),
           loaded directly as a fallback
 
-    :param parameters: Parameters passed to the SQL query via ``cursor.mogrify``.
+    :param parameters: Parameters bound client-side into the SQL query (via ``mogrify``).
     :param has_header: Include a CSV header row. Defaults to ``True``.
     :param compression: Compression format. Supports ``"gzip"``.
         Defaults to ``None`` (no compression).
@@ -38,7 +89,7 @@ class PostgresToCsvOperator(BaseOperator):
         Disable for very large files to avoid the extra read. Defaults to ``True``.
     """
 
-    template_fields: Sequence[str] = (
+    template_fields: Collection[str] = (
         "sql",
         "csv_file_path",
     )
@@ -81,16 +132,16 @@ def execute(self, context):
 
         with pg_hook.get_conn() as conn:
             with conn.cursor() as cursor:
-                cursor.execute("SET statement_timeout = %s;", (self.timeout * 60 * 1000,))
-                formatted_sql = cursor.mogrify(cleaned_sql, self.parameters).decode("utf-8")
+                _set_statement_timeout(cursor, self.timeout)
+                formatted_sql = _mogrify(conn, cursor, cleaned_sql, self.parameters)
 
                 header_clause = " HEADER" if self.has_header else ""
                 copy_command = f"COPY ({formatted_sql}) TO STDOUT WITH CSV{header_clause}"
 
                 rows = 0
                 open_func = self._get_open_func()
                 with open_func(self.csv_file_path, "wt", encoding="utf-8") as csv_file:
-                    cursor.copy_expert(copy_command, csv_file)
+                    _copy_to_file(conn, cursor, copy_command, csv_file)
                     rows = cursor.rowcount
 
         if self.count_lines:
@@ -135,7 +186,7 @@ class CsvToPostgresOperator(BaseOperator):
         Disable for very large files to avoid the extra read. Defaults to ``True``.
     """
 
-    template_fields: Sequence[str] = ("csv_file_path", "table_name")
+    template_fields: Collection[str] = ("csv_file_path", "table_name")
 
     def __init__(
         self,
@@ -200,15 +251,15 @@ def execute(self, context):
 
         with pg_hook.get_conn() as conn:
             with conn.cursor() as cursor:
-                cursor.execute("SET statement_timeout = %s;", (self.timeout * 60 * 1000,))
+                _set_statement_timeout(cursor, self.timeout)
                 if self.truncate:
                     self.log.info("Truncating table %s", self.table_name)
                     cursor.execute(f"TRUNCATE {self._quote_table_name()}")
                 open_func = self._get_open_func()
                 with open_func(self.csv_file_path, "rt", encoding="utf-8") as csv_file:
                     if self.columns and self.has_header:
                         next(csv_file)
-                    cursor.copy_expert(copy_command, csv_file)
+                    _copy_from_file(cursor, copy_command, csv_file)
                 rows = cursor.rowcount
             conn.commit()
 

tests/test_operators.py

@@ -272,3 +272,97 @@ def test_csv_to_postgres_is_instance_of_models_baseoperator(self, mock_pg_hook,
         result = op.execute(context={})
         assert result == 42  # mocked rowcount
         mock_pg_hook["cursor"].copy_expert.assert_called_once()
+
+
+# ---------------------------------------------------------------------------
+# psycopg3 compatibility tests
+#
+# Newer apache-airflow-providers-postgres versions return psycopg3 connections
+# when SQLAlchemy 2 and psycopg are installed. psycopg3 has no ``copy_expert``
+# and its default cursor has no ``mogrify``.
+# ---------------------------------------------------------------------------
+
+
+class _FakeCopy:
+    def __init__(self, chunks=()):
+        self.chunks = list(chunks)
+        self.written = []
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        return False
+
+    def __iter__(self):
+        return iter(self.chunks)
+
+    def write(self, data):
+        self.written.append(data)
+
+
+@pytest.fixture
+def mock_pg3_hook(mock_pg_hook):
+    pytest.importorskip("psycopg")
+    cursor = mock_pg_hook["cursor"]
+    del cursor.copy_expert
+    del cursor.mogrify
+    mock_pg_hook["conn"].info.encoding = "utf-8"
+    with patch("psycopg.ClientCursor") as client_cursor_cls:
+        client_cursor_cls.return_value.mogrify.return_value = "SELECT 1"
+        yield {**mock_pg_hook, "client_cursor_cls": client_cursor_cls}
+
+
+class TestPsycopg3Compatibility:
+    def test_statement_timeout_does_not_use_set_with_bind_param(self, mock_pg3_hook, tmp_path):
+        mock_pg3_hook["cursor"].copy.return_value = _FakeCopy()
+        PostgresToCsvOperator(
+            task_id="test",
+            conn_id="test_conn",
+            csv_file_path=str(tmp_path / "out.csv"),
+            sql="SELECT 1",
+            timeout=2,
+        ).execute(context={})
+        sql, params = mock_pg3_hook["cursor"].execute.call_args_list[0][0]
+        assert "set_config('statement_timeout'" in sql
+        assert params == ("120000",)
+
+    def test_postgres_to_csv_uses_copy(self, mock_pg3_hook, tmp_path):
+        # Split a multi-byte UTF-8 character across chunks.
+        payload = "id,name\n1,café\n".encode()
+        chunks = [payload[:14], payload[14:]]
+        mock_pg3_hook["cursor"].copy.return_value = _FakeCopy(chunks)
+        csv_path = tmp_path / "out.csv"
+
+        PostgresToCsvOperator(
+            task_id="test",
+            conn_id="test_conn",
+            csv_file_path=str(csv_path),
+            sql="SELECT * FROM t WHERE id = %(id)s",
+            parameters={"id": 1},
+        ).execute(context={})
+
+        mock_pg3_hook["client_cursor_cls"].return_value.mogrify.assert_called_once_with(
+            "SELECT * FROM t WHERE id = %(id)s", {"id": 1}
+        )
+        copy_sql = mock_pg3_hook["cursor"].copy.call_args[0][0]
+        assert copy_sql == "COPY (SELECT 1) TO STDOUT WITH CSV HEADER"
+        assert csv_path.read_text(encoding="utf-8") == "id,name\n1,café\n"
+
+    def test_csv_to_postgres_uses_copy(self, mock_pg3_hook, tmp_path):
+        fake_copy = _FakeCopy()
+        mock_pg3_hook["cursor"].copy.return_value = fake_copy
+        csv_file = tmp_path / "data.csv"
+        csv_file.write_text("a,b\n1,2\n")
+
+        result = CsvToPostgresOperator(
+            task_id="test",
+            conn_id="test_conn",
+            table_name="my_table",
+            csv_file_path=str(csv_file),
+            columns=["a", "b"],
+        ).execute(context={})
+
+        assert result == 42
+        assert "".join(fake_copy.written) == "1,2\n"
+        mock_pg3_hook["conn"].commit.assert_called_once()