Skip to content

Switch PostgreSQL / CockroachDB adapter to psycopg v3#222

Open
brentrager wants to merge 2 commits into
Maxteabag:mainfrom
brentrager:feat/psycopg3-driver
Open

Switch PostgreSQL / CockroachDB adapter to psycopg v3#222
brentrager wants to merge 2 commits into
Maxteabag:mainfrom
brentrager:feat/psycopg3-driver

Conversation

@brentrager
Copy link
Copy Markdown

Split out from #207 at the maintainer's request. This PR contains only the driver swap and the adapter changes that come with it; the macOS Tahoe fork-safety fix stays in #207.

Why

Two motivations:

  1. macOS Tahoe (26.4+) regression. Tahoe widened the set of system libraries that abort if initialized after fork() (CoreFoundation, Security.framework, libsystem_trace's os_log, libsystem_info's getaddrinfo). psycopg2 bundles libssl and pulls Security.framework into the parent at import time, which makes a forked child more likely to hit the abort path. psycopg v3 delegates TLS to libpq instead and avoids that import-time fingerprint. This was one of the crash classes covered by issue Broken pipe/worker connection closed errors in TUI only, inconsistently #189.
  2. Maintenance. psycopg2 is in maintenance mode upstream; psycopg v3 is the actively developed branch with first-class type-hint support, native asyncio, and the modern API.

What the diff touches

File Change
sqlit/domains/connections/providers/postgresql/adapter.py import psycopg; dbname= instead of database=; install_package = "psycopg[binary]"; driver_import_names = ("psycopg",). The "blank host → localhost" fix from #205 is preserved.
sqlit/domains/connections/providers/cockroachdb/adapter.py Same swap.
sqlit/domains/connections/providers/adapters/base.py Two small psycopg v3 behavior fixes in the shared CursorBasedAdapter. See below.
sqlit/domains/connections/app/install_strategy.py Adds psycopg / psycopg[binary] to the Arch package-name mapping.
pyproject.toml psycopg2-binary>=2.9.0psycopg[binary]>=3.2.0 in the postgres, cockroachdb, and all extras.
flake.nix pyPkgs.psycopg2pyPkgs.psycopg.
README.md Driver-reference table now lists psycopg[binary].
tests/unit/test_postgresql_adapter.py, test_tls_adapters.py, test_extra_options_passthrough.py Mock psycopg and assert dbname=.
tests/fixtures/postgres.py, cockroachdb.py, ssh.py, tests/integration/test_stale_connection_reconnect.py Import psycopg and use dbname=.
tests/integration/test_transactions.py Three test-bug fixes that the driver swap surfaced (details below).

The install-strategy unit tests in tests/test_install_strategy.py and tests/test_installer_cancel.py still use the literal string "psycopg2-binary" because they're exercising detect_strategy with a sample package name, not the real driver.

Two CursorBasedAdapter changes for psycopg v3 semantics

Running tests/integration/test_transactions.py against a real Postgres container with the psycopg3 adapter surfaced two driver-behavior differences:

  • execute_non_query autocommit handling. The base adapter called conn.commit() after every non-query. psycopg2 silently makes commit() a no-op when conn.autocommit is True; psycopg v3 actively COMMITs any open transaction. So executor.execute("BEGIN") was immediately ending the transaction it just started, and the subsequent INSERT then auto-committed and ROLLBACK had nothing to undo. Skipping conn.commit() when getattr(conn, "autocommit", False) preserves intent on both drivers and is safe for adapters that don't set autocommit (mysql, mssql, sqlite, oracle, firebird — getattr returns False and the commit still fires).
  • execute_query multi-statement walk. psycopg v3 exposes each result set from a multi-statement string and positions the cursor on the first set. After cursor.execute("BEGIN; INSERT; SELECT"), psycopg v3 leaves the cursor pointed at the BEGIN (no description, no rows). psycopg2 only ever surfaced the last set. Walking cursor.nextset() until it returns falsy lands on the last set; the loop is a no-op on drivers that don't implement nextset() (some raise NotSupportedError, hence the try/except).

Three test bugs the driver swap surfaced

These also reproduce on Maxteabag/sqlit:main with psycopg2 — verified by running the upstream tree against the same docker fixture. Fixing them here keeps the integration suite green:

  • test_multi_statement_as_single_query_works sent BEGIN; INSERT; SELECT through the executor and left an open transaction on the persistent connection. The test's own DROP TABLE (on a fresh connection) then blocked waiting for the lock. Added a reset_tui_executor() before the cleanup so the persistent connection rolls back and releases the lock.
  • test_atomic_execute_rolls_back_on_error asserted that atomic_execute raises on a failing statement. Per its docstring and current implementation it actually reports the failure via MultiStatementResult(completed=False, error_index=…). Updated the assertions accordingly.
  • test_atomic_execute_commits_on_success asserted a single QueryResult for a 3-statement query. atomic_execute is documented to return a MultiStatementResult carrying each StatementResult; the trailing SELECT lands as result.results[-1].result.

Tests

Run against postgres:16 from infra/docker/docker-compose.test.yml:

Suite Result
tests/unit 897 passed, 1 skipped
tests/integration/test_transactions.py 13 passed
tests/integration/test_database_browsing_flow.py::TestPostgreSQLDatabaseBrowsing 1 passed
tests/integration/test_stale_connection_reconnect.py -k postgres 5 passed, 1 idle-variant skipped (pre-existing)

Tradeoffs to flag

  • Runtime dependency bump from psycopg2-binary>=2.9.0 to psycopg[binary]>=3.2.0. The adapter surface that sqlit actually uses is API-compatible, but it's still a real change for anyone bundling sqlit.
  • The autocommit branch in the base execute_non_query is the smallest change that gives both drivers consistent behavior. An alternative would be to stop setting conn.autocommit = True in the PG adapter altogether and let psycopg v3 manage transactions natively via conn.transaction() — that's a bigger surface change and would also need the RuleAggregation workers to opt in. Happy to take that path instead if you'd prefer it.

Companion to the Maxteabag#189 fork-fallback fix. psycopg3 delegates TLS to
libpq instead of bundling libssl, so importing it doesn't pull
Security.framework into the parent process on macOS — eliminating one
specific child-side crash class (psycopg2 calling SSL_CTX_new inside
the libpq init path after fork).

The change is local to the two PostgreSQL-family adapters and the
matching install metadata:

  - postgresql/adapter.py, cockroachdb/adapter.py: import psycopg
    (v3) and pass dbname= instead of database=; everything else
    (sslmode, sslrootcert, sslcert, sslkey, sslpassword, autocommit,
    extra_options) is already libpq-conninfo-compatible.
  - install_strategy.py: add psycopg / psycopg[binary] to the Arch
    package-name mapping.
  - pyproject.toml: replace psycopg2-binary>=2.9.0 with
    psycopg[binary]>=3.2.0 in the postgres, cockroachdb, and all
    extras.
  - flake.nix: pyPkgs.psycopg2 -> pyPkgs.psycopg.
  - README.md: driver-reference table now lists psycopg[binary].
  - tests + integration fixtures: import psycopg (3) and use dbname=.
    The install-strategy tests still use the literal string
    "psycopg2-binary" because they're exercising detect_strategy with
    a sample package name, not the real driver.

837 unit tests pass.
… multi-statement

Running tests/integration/test_transactions.py against a real PostgreSQL
container with the psycopg3 adapter surfaced two driver-behavior
differences that need handling in the shared CursorBasedAdapter:

1. autocommit-aware execute_non_query
   psycopg2 silently makes conn.commit() a no-op when conn.autocommit is
   True; psycopg v3 actively COMMITs any open transaction. The adapter's
   execute_non_query unconditionally called conn.commit(), which under
   psycopg3 meant `executor.execute("BEGIN")` immediately ended the
   transaction it just started — so the subsequent INSERT was committed
   and ROLLBACK had nothing to undo. Skipping the commit when
   conn.autocommit is True preserves the original intent on both
   drivers.

2. nextset() walk in execute_query
   psycopg v3 exposes each result set in a multi-statement string and
   leaves the cursor positioned at the first (e.g. on BEGIN, which has
   no columns). psycopg2 only ever surfaced the last set. Walking
   cursor.nextset() until it returns falsy lands on the last set; the
   loop is a no-op on drivers that don't implement nextset() (a number
   raise NotSupportedError, hence the try/except).

The integration tests in tests/integration/test_transactions.py needed
three fixes that are independent of the driver swap — verified to fail
or hang on upstream main with psycopg2 as well:

- test_multi_statement_as_single_query_works
  The "BEGIN; INSERT; SELECT" payload left an open transaction on the
  executor's persistent connection; the test's subsequent DROP TABLE
  (on a different connection) then blocked waiting for the lock. Added
  reset_tui_executor() before the cleanup so the persistent connection
  rolls back and releases the lock.

- test_atomic_execute_rolls_back_on_error
  The test expected atomic_execute to raise on a failing statement.
  Per its docstring and current implementation, atomic_execute reports
  the failure via a MultiStatementResult with completed=False and the
  failing index in error_index. Updated the assertions accordingly.

- test_atomic_execute_commits_on_success
  Multi-statement atomic_execute is documented to return a
  MultiStatementResult (one StatementResult per statement), not a
  single QueryResult. Asserting on result.results[-1].result picks up
  the trailing SELECT's QueryResult exactly as the test intended.

Result against the docker-compose postgres:16 fixture:
  - tests/integration/test_transactions.py: 13 passed
  - tests/integration/test_database_browsing_flow.py::TestPostgreSQLDatabaseBrowsing: 1 passed
  - tests/integration/test_stale_connection_reconnect.py -k postgres: 5 passed, 1 idle-variant skipped (pre-existing)
  - tests/unit: 842 passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant