sgoel2be24-cyber opened a new pull request, #73324: URL: https://github.com/apache/airflow/pull/73324
With the psycopg3 driver (the default since `apache-airflow-providers-postgres` 7.0.0), `PostgresToGCSOperator(use_server_side_cursor=True)` read rows through `fetchone()`, which issues `FETCH FORWARD 1` per row and ignores `cursor_itersize`. The issue reports a 2M-row export going from ~5 to 90+ minutes. `_PostgresServerSideCursorDecorator` now iterates the cursor for both drivers, so psycopg fetches `itersize` rows per round trip, as psycopg2 already did. The decorator keeps a single iterator because psycopg < 3.3 (the provider allows `>=3.2.9`) implements `ServerCursor.__iter__` as a generator: calling `iter()` again would start a new one and drop the rest of the current batch. psycopg >= 3.3 and psycopg2 named cursors are their own iterators, so nothing changes for psycopg2. The existing tests for this operator need a Postgres backend, and they pass with any fetch size, so they couldn't catch this. The new unit test uses fake cursors that behave like both psycopg iteration styles and asserts every row is returned and rows are fetched in `itersize` batches (`[1, 100, 100, 100]` for 250 rows; the single-row fetch is the one `description` needs). On `main` both cases fail with 251 single-row fetches. Tested locally: - `pytest` on `test_postgres_to_gcs.py` and `test_sql_to_gcs.py`: 16 passed. The 18 Postgres-backend tests were skipped: there's no Postgres or Docker on my machine, so they'll run in CI - `prek` pre-commit stage passes (`check-template-fields-valid` needs Docker and was skipped; no template fields changed); `mypy` on `postgres_to_gcs.py` passes - Checked the `ServerCursor` source for psycopg 3.2.9 and 3.3.5 to confirm both iteration behaviours the fakes model closes: #72075 --- ##### Was generative AI tooling used to co-author this PR? - [X] Yes — Claude Code (Opus 5) Generated-by: Claude Code (Opus 5) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
