feat: keep a session's transaction alive across batches

[egertaia]

Companion to TabularisDB/tabularis#801, which plumbs the session id through the host. Either can merge first: the wire change is backward compatible in both directions.

Problem

execute_query_batch acquires one pooled connection for the whole batch, so BEGIN … COMMIT inside a single script works. The workflow the transaction exists for does not:

Run 1:  BEGIN; UPDATE accounts SET balance = 0 WHERE id = 42;
Run 2:  SELECT balance FROM accounts WHERE id = 42;   -- check before committing
Run 3:  COMMIT;

Each run is its own RPC call (and a one-statement run arrives as execute_query, not execute_query_batch), so the connection goes back to the pool in between. Run 2 can land on a different connection and show pre-transaction data; run 3 reports there is no transaction in progress while the real transaction stays open elsewhere.

It is also a hazard for unrelated queries. The pool is configured with RecyclingMethod::Fast (src/client.rs:376), which resets nothing, so the stranded connection returns to the pool still inside the transaction, holding its locks, and the next borrower runs inside it.

Fix

When the host sends a session_id, a batch that leaves an explicit transaction open keeps its connection pinned to that session, so the next batch from the same session continues the same transaction. Committing, rolling back, an explicit release_session, shutdown, or 30 minutes idle releases it. A pinned connection is always rolled back before it goes back to the pool, so an open transaction can no longer leak into an unrelated query.

Pinning is lazy: a session that never opens a transaction holds nothing and behaves exactly as before.

• transaction_effect() in handlers/query.rs (BEGIN, START TRANSACTION, COMMIT, END, ROLLBACK, ABORT, PREPARE TRANSACTION, and ... AND CHAIN), next to the existing strip_leading_sql_comments it builds on. It reads leading keywords only, so BEGIN in a string literal, a later clause, or a PL/pgSQL DO $$ BEGIN … END $$ body is not mistaken for transaction control, and ROLLBACK TO SAVEPOINT correctly leaves the transaction open.
• session.rs: the registry, the ROLLBACK-before-release rule, and sweep_idle(), run from the existing 10-minute run_pool_cleanup timer.
• run_batch_in_session, extracted from execute_query_batch, reuses the session's pinned connection and re-pins it if a transaction is still open afterwards. execute_query goes through it as a one-statement batch, so a single statement and a batch take the same session path. Without that, the one-at-a-time workflow above is exactly the case that stays broken. A failed ordinary statement leaves the transaction state alone (it is aborted but still open, so the user can ROLLBACK from the same tab). A failed COMMIT is different: PostgreSQL has already rolled the transaction back, so the session is released.
• release_session RPC method; shutdown releases every pinned session.

Wire compatibility

• session_id is optional on both execute_query and execute_query_batch. Without it the behaviour and the responses are byte-for-byte what they are today: a bare QueryResult and a bare array of per-statement results. Hosts older than tabularis#801 are unaffected.
• With it, the replies are { "result": {...}, "in_transaction": bool } and { "results": [...], "in_transaction": bool }. The host side accepts both shapes permanently, and recognises the wrapper by in_transaction being present so a query selecting a column named result is not mistaken for one.
• release_session is a new method; a host that never calls it loses nothing, since the idle sweep and shutdown still reclaim connections.

One deliberate behaviour difference on a reused connection

search_path is applied only when the connection is freshly acquired. Re-applying it to a pinned connection would run SET search_path inside the caller's open transaction and change what the rest of that transaction sees. The consequence is that switching schema mid-transaction does not take effect until the transaction ends, which is the safer of the two behaviours, but worth a maintainer's eye.

Testing

• cargo test --lib: 340 pass, including 9 cases for transaction_effect and in_transaction_after (opens / closes / AND CHAIN / two-phase / savepoint / leading comments / data-not-keyword / PL/pgSQL body / empty input).
• cargo clippy --all-targets -- -D warnings and cargo fmt --check clean.
• Not exercised against a live PostgreSQL instance or a real Tabularis host. tests/live_db.rs now has commit_that_fails_on_a_deferred_constraint_releases_the_session (from review), but I have not run it locally.

[aesslinger]

Thanks for this, @egertaia — this is a genuinely well-scoped fix for a real correctness/lock-leak bug, and it's unusually well self-documented for a PR of this size. The write-up connecting RecyclingMethod::Fast resetting nothing to the stranded-transaction-leaks-into-the-next-borrower hazard was exactly the right level of detail, and calling out the search_path-not-reapplied-on-reuse tradeoff yourself (rather than letting a reviewer find it) is the kind of thing that makes this easy to review with confidence.

One non-blocking note on hardening, plus a live-DB pass that turned up something I think does need fixing before this merges — not just a nice-to-have.

---

1. Idle-sweep is opportunistic, not scheduled (non-blocking)

MAX_IDLE (30 min) is only checked inside session::take() — i.e., triggered by some other session's next call. If no other session ever calls take() again, an abandoned pinned connection holds its lock until shutdown. In practice that means one tab opening a transaction and then never being touched again (browser crash, host killed without a clean release_session, etc.) can hold a row lock indefinitely as long as nothing else happens to query through this handler.

main.rs already runs a periodic background task (run_pool_cleanup, every 10 min) for evicting idle connection pools. Would it be worth hanging the session sweep off that same timer instead of piggybacking on take()? Something like:

// main.rs
async fn run_pool_cleanup(mut shutdown_rx: watch::Receiver<bool>) {
    let mut timer = interval(POOL_CLEANUP_INTERVAL);
    loop {
        tokio::select! {
            _ = timer.tick() => {
                postgresql_plugin::client::cleanup_idle_pools();
                postgresql_plugin::session::sweep_idle().await;
            }
            _ = shutdown_rx.changed() => break,
        }
    }
}

with session::sweep_idle() factored out of the stale-scanning logic currently inline in take(). Not blocking — just want to confirm whether the current behavior (sweep-on-next-call only) was a deliberate simplification, since a busy plugin process (many sessions cycling through take()) never really needs the standalone sweep, but a mostly-idle one with one abandoned session would.

---

2. Live-database coverage — draft tests, and a bug I'd call blocking

You mentioned in the PR description you'd welcome pointers on wiring this into tests/live_db.rs. I put together 8 scenarios using the existing Plugin harness (same stdin/stdout drive-the-binary approach already in that file) and ran them against a real Postgres 16 container. 7 passed cleanly and match the PR's claims:

• BEGIN → UPDATE → SELECT (verify) → COMMIT, each as a separate execute_query call with the same session_id — confirms the core fix.
• Same sequence without session_id — documents the old bug still reproduces without opting in.
• ROLLBACK across separate calls discards the write.
• A batch that leaves a transaction open with no session_id doesn't leak its lock into an unrelated query (checked via SELECT ... FOR UPDATE NOWAIT on the same row from a fresh call).
• release_session rolls back an open transaction.
• Two distinct session_ids stay isolated from each other (separate connections, separate uncommitted state).
• execute_query_batch and execute_query interoperate on the same session (start via batch, continue via single statement, commit via batch again) — validates the "single statement and batch share one code path" claim.

Ready-to-drop-in versions of all of these, adapted to live_db.rs's existing Plugin/call_ok helpers:

fn setup_session_scratch_table(plugin: &mut Plugin, params: &Value) {
    plugin.call_ok("execute_query", json!({
        "params": params,
        "query": "CREATE TABLE IF NOT EXISTS live_db_session_scratch \
                   (id SERIAL PRIMARY KEY, value INTEGER)",
    }));
    plugin.call_ok("execute_query", json!({
        "params": params,
        "query": "TRUNCATE live_db_session_scratch RESTART IDENTITY",
    }));
    plugin.call_ok("execute_query", json!({
        "params": params,
        "query": "INSERT INTO live_db_session_scratch (value) VALUES (1)",
    }));
}

#[test]
fn session_id_keeps_a_transaction_alive_across_separate_execute_query_calls() {
    let mut plugin = Plugin::spawn();
    let params = conn_params();
    setup_session_scratch_table(&mut plugin, &params);
    let session = json!("live-db-session-1");

    let r1 = plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session, "query": "BEGIN"
    }));
    assert_eq!(r1["in_transaction"], json!(true));

    plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session,
        "query": "UPDATE live_db_session_scratch SET value = 99 WHERE id = 1"
    }));

    let r2 = plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session,
        "query": "SELECT value FROM live_db_session_scratch WHERE id = 1"
    }));
    assert_eq!(
        r2["result"]["rows"][0][0], json!(99),
        "verifying SELECT on the same session must see the uncommitted UPDATE"
    );

    let r3 = plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session, "query": "COMMIT"
    }));
    assert_eq!(r3["in_transaction"], json!(false));

    let r4 = plugin.call_ok("execute_query", json!({
        "params": params,
        "query": "SELECT value FROM live_db_session_scratch WHERE id = 1"
    }));
    assert_eq!(r4["rows"][0][0], json!(99), "committed value must be visible with no session_id");
}

#[test]
fn rollback_discards_changes_made_across_separate_calls() {
    let mut plugin = Plugin::spawn();
    let params = conn_params();
    setup_session_scratch_table(&mut plugin, &params);
    let session = json!("live-db-session-2");

    plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session, "query": "BEGIN"
    }));
    plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session,
        "query": "UPDATE live_db_session_scratch SET value = 777 WHERE id = 1"
    }));
    let rollback = plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session, "query": "ROLLBACK"
    }));
    assert_eq!(rollback["in_transaction"], json!(false));

    let check = plugin.call_ok("execute_query", json!({
        "params": params,
        "query": "SELECT value FROM live_db_session_scratch WHERE id = 1"
    }));
    assert_eq!(check["rows"][0][0], json!(1), "ROLLBACK must discard the UPDATE");
}

#[test]
fn abandoned_transaction_without_session_id_does_not_leak_its_lock() {
    let mut plugin = Plugin::spawn();
    let params = conn_params();
    setup_session_scratch_table(&mut plugin, &params);

    // A batch with no session_id that leaves a transaction open.
    plugin.call_ok("execute_query_batch", json!({
        "params": params,
        "queries": json!([
            "BEGIN",
            "UPDATE live_db_session_scratch SET value = 5 WHERE id = 1"
        ]),
    }));

    // If the abandoned transaction leaked back into the pool with its
    // lock still held, this NOWAIT query on the same row would error.
    let r = plugin.call("execute_query", json!({
        "params": params,
        "query": "SELECT value FROM live_db_session_scratch WHERE id = 1 FOR UPDATE NOWAIT"
    }));
    assert!(
        r.get("error").is_none(),
        "unrelated query must not inherit the abandoned transaction's lock: {r:?}"
    );
}

#[test]
fn release_session_rolls_back_an_open_transaction() {
    let mut plugin = Plugin::spawn();
    let params = conn_params();
    setup_session_scratch_table(&mut plugin, &params);
    let session = json!("live-db-session-3");

    plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session, "query": "BEGIN"
    }));
    plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session,
        "query": "UPDATE live_db_session_scratch SET value = 333 WHERE id = 1"
    }));
    plugin.call_ok("release_session", json!({ "session_id": session }));

    let r = plugin.call_ok("execute_query", json!({
        "params": params,
        "query": "SELECT value FROM live_db_session_scratch WHERE id = 1"
    }));
    assert_eq!(r["rows"][0][0], json!(1), "release_session must roll back the uncommitted UPDATE");
}

#[test]
fn two_sessions_stay_isolated_from_each_other() {
    let mut plugin = Plugin::spawn();
    let params = conn_params();
    setup_session_scratch_table(&mut plugin, &params);
    plugin.call_ok("execute_query", json!({
        "params": params,
        "query": "INSERT INTO live_db_session_scratch (value) VALUES (0)"
    }));
    let session_a = json!("live-db-session-4a");
    let session_b = json!("live-db-session-4b");

    plugin.call_ok("execute_query", json!({ "params": params, "session_id": session_a, "query": "BEGIN" }));
    plugin.call_ok("execute_query", json!({ "params": params, "session_id": session_b, "query": "BEGIN" }));

    plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session_a,
        "query": "UPDATE live_db_session_scratch SET value = 111 WHERE id = 1"
    }));
    plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session_b,
        "query": "UPDATE live_db_session_scratch SET value = 222 WHERE id = 2"
    }));

    let ra = plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session_a,
        "query": "SELECT value FROM live_db_session_scratch WHERE id = 1"
    }));
    assert_eq!(ra["result"]["rows"][0][0], json!(111));

    plugin.call_ok("execute_query", json!({ "params": params, "session_id": session_a, "query": "ROLLBACK" }));
    plugin.call_ok("execute_query", json!({ "params": params, "session_id": session_b, "query": "COMMIT" }));

    let r1 = plugin.call_ok("execute_query", json!({
        "params": params, "query": "SELECT value FROM live_db_session_scratch WHERE id = 1"
    }));
    assert_eq!(r1["rows"][0][0], json!(1), "session_a's rollback must discard its own write");

    let r2 = plugin.call_ok("execute_query", json!({
        "params": params, "query": "SELECT value FROM live_db_session_scratch WHERE id = 2"
    }));
    assert_eq!(r2["rows"][0][0], json!(222), "session_b's commit must persist");
}

#[test]
fn batch_and_single_statement_calls_share_the_same_session() {
    let mut plugin = Plugin::spawn();
    let params = conn_params();
    setup_session_scratch_table(&mut plugin, &params);
    let session = json!("live-db-session-5");

    let r1 = plugin.call_ok("execute_query_batch", json!({
        "params": params, "session_id": session, "queries": json!(["BEGIN"])
    }));
    assert_eq!(r1["in_transaction"], json!(true));

    plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session,
        "query": "UPDATE live_db_session_scratch SET value = 55 WHERE id = 1"
    }));

    let r2 = plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session,
        "query": "SELECT value FROM live_db_session_scratch WHERE id = 1"
    }));
    assert_eq!(r2["result"]["rows"][0][0], json!(55));

    let r3 = plugin.call_ok("execute_query_batch", json!({
        "params": params, "session_id": session, "queries": json!(["COMMIT"])
    }));
    assert_eq!(r3["in_transaction"], json!(false));
}

The 8th scenario found a bug that I think should block merge, not follow up later, because it undermines the exact thing this PR sets out to fix — stray pinned connections outliving the transaction that justified holding them.

run_batch_in_session only advances in_transaction via transaction_effect() when outcome.is_ok() (query.rs, the if outcome.is_ok() { match transaction_effect(query) ... } block). That's correct for an ordinary statement failing mid-transaction (e.g. a bad INSERT — the transaction is aborted but still open, so staying pinned is right). But it's not correct for a COMMIT that itself fails, which happens whenever a deferred constraint is violated:

#[test]
fn commit_that_fails_on_a_deferred_constraint_does_not_leave_a_stale_pin() {
    let mut plugin = Plugin::spawn();
    let params = conn_params();

    plugin.call_ok("execute_query", json!({
        "params": params, "query": "DROP TABLE IF EXISTS live_db_deferred_fk_scratch"
    }));
    plugin.call_ok("execute_query", json!({
        "params": params,
        "query": "CREATE TABLE live_db_deferred_fk_scratch (id INT PRIMARY KEY, other_id INT)"
    }));
    plugin.call_ok("execute_query", json!({
        "params": params,
        "query": "ALTER TABLE live_db_deferred_fk_scratch ADD CONSTRAINT fk_self \
                   FOREIGN KEY (other_id) REFERENCES live_db_deferred_fk_scratch(id) \
                   DEFERRABLE INITIALLY DEFERRED"
    }));

    let session = json!("live-db-session-commit-fails");
    plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session, "query": "BEGIN"
    }));
    // The FK check is deferred, so this INSERT succeeds; only COMMIT fails.
    plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session,
        "query": "INSERT INTO live_db_deferred_fk_scratch (id, other_id) VALUES (1, 999)"
    }));

    let commit = plugin.call("execute_query", json!({
        "params": params, "session_id": session, "query": "COMMIT"
    }));
    assert!(commit.get("error").is_some(), "COMMIT violating a deferred FK must error");

    // PostgreSQL implicitly ends the transaction attempt once COMMIT
    // itself fails — a later ROLLBACK on the same backend just warns
    // "no transaction in progress". So the session should NOT still
    // report in_transaction: true here.
    let after = plugin.call_ok("execute_query", json!({
        "params": params, "session_id": session, "query": "SELECT 1 AS still_usable"
    }));
    assert_eq!(
        after["in_transaction"], json!(false),
        "a failed COMMIT already ended the transaction server-side; \
         the session must not stay pinned believing one is still open"
    );
}

Ran this against the branch locally and it fails:

assertion `left == right` failed: a failed COMMIT already ended the transaction
server-side; the session must not stay pinned believing one is still open
  left: Bool(true)
 right: Bool(false)

I verified the underlying Postgres behavior directly with psql first, since it surprised me too: a COMMIT on a transaction merely aborted by an earlier statement (e.g. SELECT 1/0) does not error at all — Postgres silently treats it as a ROLLBACK and reports success. It's specifically a COMMIT that fails on its own (deferred constraint, SET CONSTRAINTS ... IMMEDIATE, etc.) that errors and leaves no transaction behind. Since the current code treats "COMMIT returned Err" as "no change to in_transaction," it keeps the connection pinned to the session indefinitely with no transaction to justify it — the exact failure mode this PR exists to close, just via a different trigger.

A COMMIT/ROLLBACK failure probably needs to be treated as TransactionEffect::Closes regardless of outcome, rather than folded into the generic "only apply on success" rule that's correct for every other statement.

---

Really appreciate the depth here — happy to take another pass as soon as you've had a chance to look at the COMMIT-failure case.

[github-actions]

Version suggestion

Based on this PR's title (feat) and the prerelease:rc label:

Current	1.0.0-rc.4
Suggested next tag	v1.0.0-rc.5

This is informational only — no tag or release is created automatically yet.

[egertaia]

Thanks, both confirmed and fixed on the branch (09849ce, 14548f4).

COMMIT failure: you were right, the "failed COMMIT leaves the transaction open" premise was wrong. TransactionEffect::in_transaction_after() now treats a close as ending the transaction regardless of outcome. Auditing the classifier turned up three more of the same kind, fixed in the same commit:

• COMMIT AND CHAIN was classified as a plain close, so on success the connection went back to the pool with the new chained transaction still open. This was the worst of them. It is now TransactionEffect::Chains: it stays pinned on success and releases on failure.
• ABORT (a synonym for ROLLBACK) was not recognised.
• PREPARE TRANSACTION now counts as a close, and COMMIT PREPARED / ROLLBACK PREPARED are ignored, since they never touch the session's own transaction.

Your deferred-FK test is in tests/live_db.rs as commit_that_fails_on_a_deferred_constraint_releases_the_session. I have no local Postgres to run it against, so it would be great if you could run it on your setup.

Idle sweep: not deliberate, it was an oversight. sweep_idle() is factored out of take() and now runs on the run_pool_cleanup timer, as you suggested.

[aesslinger]

Thanks for the quick turnaround, @egertaia — both fixes are correct and the audit you did beyond just my one repro case is the right instinct.

What I checked

Rather than just reading the diff, I built the release binary off this branch and ran it against a real, freshly-initialized PostgreSQL 16:

• cargo test --lib: 340/340 pass. clippy --all-targets -- -D warnings and fmt --check: clean.
• cargo test --test live_db: 25/26 pass (1 ignored, needs the pgvector extension). Your new commit_that_fails_on_a_deferred_constraint_releases_the_session test passes for real against a live database, not just the classifier unit tests. The one failure (connecting_with_a_wrong_password_...) is not a regression — I ran the identical test against main in a separate worktree with the same scratch DB and it fails there too, since my scratch instance uses trust auth and doesn't check passwords at all.
• The TransactionEffect::Chains handling is right, and I like that you went and audited the whole classifier rather than patching just the deferred-FK case — ABORT, PREPARE TRANSACTION, and COMMIT/ROLLBACK PREPARED were real gaps COMMIT AND CHAIN in particular was the nastiest of the three, since on success it silently left a new transaction open on a connection about to go back to the pool.
• The idle-sweep move to run_pool_cleanup's timer is exactly what I suggested, and store()'s last_used refresh means there's no risk of the timer sweeping a session that's actually mid-flight (it's absent from the map while a batch is running against it).

No blockers from any of that.

One new issue, found doing a deeper concurrency pass — I think this should block merge

The plugin dispatches requests across WORKER_POOL_SIZE = 4 concurrent workers (main.rs). session::take() / session::store() (src/session.rs:59, src/session.rs:93) don't serialize access per session_id — nothing stops two overlapping RPC calls for the same session from both calling take(), both getting None (because neither has stored yet), and both grabbing separate fresh connections. I reproduced this against the actual built binary, not just by reading the code:

Reproduction 1 — stale read. Session S, call A = execute_query_batch with ["BEGIN", "UPDATE t SET v=99 WHERE id=1", "SELECT pg_sleep(1)"] (the sleep holds A's worker). ~150ms later, call B = execute_query on the same session, "SELECT v FROM t WHERE id=1". Call B returned before call A, and read v=1 (the pre-transaction value) — it ran on a brand-new connection, not the one holding A's uncommitted UPDATE. This is the exact "run 2 sees pre-transaction data" bug this PR exists to fix, just triggered by worker-pool concurrency instead of pool-checkout timing.

Reproduction 2 — silent data loss, worse. Same setup, but both calls open a transaction and write to different rows (so they don't block on Postgres row locks — this is a plugin-side session-map race, not a database lock issue). Call B (fast) does BEGIN; UPDATE t SET v=200 WHERE id=2, finishes, and calls store(S, connB) — the map now holds connB. ~850ms later call A (still sleeping) finishes and calls store(S, connA) — store()'s collision handling (src/session.rs:93-111, the if let Some(stale) = previous { rollback_and_release(stale.client) } block) sees connB already there, silently rolls it back, and overwrites the map with connA. Call B was told in_transaction: true (success) but its transaction and write were silently discarded. Then a COMMIT sent on session S commits connA's data, not connB's — the host has no way to know which transaction it's actually talking to at that point. Actual output from the run:

Committed rows in DB: [[1, 100], [2, 0]]
CONFIRMED: call A's update (row1=100) was committed; call B's update (row2=200)
was silently discarded even though call B was told in_transaction:true (success).

This is reachable in real usage, not just theoretical. I checked the host side (tabularis PR #801, src/pages/Editor.tsx) to see whether the editor guarantees one in-flight call per tab before recommending a fix. It doesn't:

• queryGenerationRef's own comments describe "running two different statements/tables back to back via cursor-driven Run" as an anticipated, normal scenario — the generation counter exists specifically to let the UI discard stale results from an earlier overlapping run, which means overlapping backend calls for one tab are an accepted case, not a bug the host thinks it prevents.
• The global Ctrl/Cmd+F5 shortcut handler calls handleRunButton() unconditionally — no activeTab.isLoading check.
• The Monaco context-menu "Execute Selection" action (editor.addAction({ id: "run-selection", ... })) calls runQueryRef.current(...) directly, also with no isLoading check.
• Only the toolbar Run button itself is guarded (it swaps to a Stop button while isLoading) — the keyboard shortcut and context-menu paths bypass that entirely.

So a user pressing the run shortcut twice quickly, or right-clicking "Execute Selection" while a previous statement on the same tab is still executing, sends two overlapping execute_query/execute_query_batch calls with the same sessionId today (once #801 lands). That's exactly the race above.

Suggested fix direction: serialize per-session_id access inside the plugin rather than relying on the host — other hosts/callers of this RPC in the future would hit the same hole otherwise. Something like an async Mutex (or a per-session lock map) around the take → run batch → store sequence in run_batch_in_session (src/handlers/query.rs:103), so a second concurrent call for the same session either waits for the first to finish and release/store, or gets a clear error back instead of silently grabbing an unrelated connection or clobbering a live transaction. Happy to pair on this or take a pass myself if useful — wanted to hand over the full repro first since it's a subtle one.

[egertaia]

Thanks for the repro, confirmed and fixed in 7005c15.

Each session now has its own lock. session::lock() returns an OwnedMutexGuard<Slot> that run_batch_in_session holds across take, run and pin, so an overlapping call for the same session waits, then runs on the same connection and sees the first call's uncommitted changes. The collision rollback in store() is gone, since two runs can no longer pin at once. release_session takes the same lock, so it waits for a run still in flight instead of missing it. sweep_idle() skips slots a run is holding, and forgets empty ones nothing references.

I went with waiting rather than returning an error: statements then run in the order the user sent them, which is what the editor expects. src/session_tests.rs covers the waiting, sessions not blocking each other, and the sweep not dropping a held slot. It would be great if you could rerun your two repros against it.

Also on the branch: 6e9857c makes a successful ... AND CHAIN always count as inside a transaction, from a review on tabularis#801. The host side has the same lock (tabularis#801, bc5301d6).

[aesslinger]

@egertaia — ran this through the same review process I use on my own work — eleven independent lenses (correctness, security, concurrency, error-handling, contract-consistency, docs-vs-code, edge-cases, test-coverage, validation, backward-compat, test-structure), each looking only through its own lens. That fan-out raised a batch of raw findings, but the ones below are the ones that survived actually running them against the real binary and the real function — not just re-reading the diff and reasoning about what "should" happen. That distinction mattered a lot here: the most serious one (#1) only became undeniable once I instrumented the live code and watched it happen; a plausible-sounding version of it could easily have been argued away on paper. Sorry to come back with more after saying this looked close; better now than after merge, and happy to help however's useful — pair on any of these, or I can just take a pass at one/all of the fixes myself if that's faster for you.

1. release_all() can orphan a busy session's connection with no ROLLBACK — new, distinct from the race you already fixed, and the one live testing actually caught red-handed

src/session.rs:138 (release_all). It does sessions().lock().await.drain() first — unconditionally removing every session's slot from the map, including ones a run is currently using — and only afterward calls try_lock() per slot to decide what to roll back. Contrast with sweep_idle(), which uses map.retain() and explicitly keeps a busy slot tracked (return true) instead of removing it.

I want to flag how this one was found, because it's the strongest evidence in this batch. Reading the code, the race is plausible but arguable — you could talk yourself into "surely something protects against that." So instead of stopping at the reasoning, I built the actual release binary, ran it against a real Postgres, and drove the exact race live: session S has an in-flight execute_query_batch (BEGIN, then a 2s pg_sleep) holding session::lock(S)'s guard; shutdown fires concurrently while that's still running. Watching pg_stat_activity in real time, the backend running pg_sleep(2) simply vanished the moment the in-flight call returned — not present, not idle-in-transaction, just gone, with no ROLLBACK ever sent. To pin down exactly why, I added three temporary eprintln!s to the real session.rs (drop them, this isn't a suggested patch) and reran the same live race:

release_all: drained 1 slot(s)
release_all: slot[0] strong_count=2 try_lock=false
release_all: took 0 client(s)
...
Slot::pin called
Slot dropped WHILE STILL HOLDING A PINNED CLIENT (no rollback issued)

That's the actual code, mid-flight, catching itself in the act: release_all correctly sees the slot is busy and skips it (as designed) — but by then it's already removed the slot from the map, and once the in-flight run finishes and re-pins its connection into that now-orphaned slot, nothing else references it, so it deallocates without ever calling rollback_and_release.

session.rs's own module doc says "a pinned connection is always rolled back before it goes back to the pool" — this path breaks that. Fix direction: make release_all retain-and-skip busy entries the same way sweep_idle does, instead of draining unconditionally.

2. ROLLBACK WORK/TRANSACTION TO SAVEPOINT misclassified as closing the transaction

transaction_effect() (src/handlers/query.rs:452) only recognizes ROLLBACK TO .... Postgres's actual grammar is ROLLBACK [ WORK | TRANSACTION ] TO [ SAVEPOINT ] savepoint_name — the WORK/TRANSACTION keyword is optional. When present, it falls through to the unconditional Closes arm instead of None, contradicting the function's own doc comment. This one didn't need a live database at all — I ran the actual transaction_effect() from the branch directly against these inputs and it fails immediately:

assert_eq!(transaction_effect("ROLLBACK WORK TO SAVEPOINT sp1"), TransactionEffect::None); // fails: got Closes
assert_eq!(transaction_effect("ROLLBACK TRANSACTION TO SAVEPOINT sp1"), TransactionEffect::None); // fails: got Closes

Consequence: a session that does BEGIN; ...; SAVEPOINT sp1; ...; ROLLBACK TRANSACTION TO SAVEPOINT sp1 gets in_transaction forced to false even though the outer transaction is still open server-side — the connection then falls through neither the re-pin nor the rollback branch at the end of run_batch_in_session, and drops back to the pool via plain Drop still mid-transaction. Same leak class as the COMMIT AND CHAIN bug from your last pass.

3. A trailing comment mentioning "and chain" misclassifies COMMIT as chaining

Also in transaction_effect() (src/handlers/query.rs:434): only leading comments are stripped (strip_leading_sql_comments), so a query like "COMMIT -- and chain later if needed" has its trailing comment tokenized right alongside the real keyword — -- is just another delimiter character to the word-splitter, so AND and CHAIN end up adjacent and the chains check fires. Same approach as #2 — ran it directly, no live DB needed:

assert_eq!(transaction_effect("COMMIT -- and chain later if needed"), TransactionEffect::Closes); // fails: got Chains

Consequence: the plugin believes the transaction is still open (reports in_transaction: true) when Postgres actually closed it. A later explicit COMMIT/ROLLBACK the client sends becomes a silent no-op, and a host UI reading in_transaction off every response shows a stale "in transaction" indicator indefinitely. This is a pretty ordinary thing for a user to type, not an exotic input.

Related, same root cause, lower likelihood of hitting it in practice: "START /* explicit */ TRANSACTION" shifts the second-word position and defeats the Opens check entirely (second == Some("EXPLICIT"), not "TRANSACTION") — also confirmed directly. Probably worth stripping/ignoring all comments before tokenizing, not just leading ones, rather than patching each shape individually.

4. execute_query's session path drops in_transaction on an error response

src/handlers/query.rs:28-46. When the one statement in a session-scoped execute_query call fails, the handler does error_response(id, -32603, error) and discards the in_transaction value run_batch_in_session just computed. execute_query_batch doesn't have this problem — it always returns { results, in_transaction }, with the per-statement error nested inside results[i].error instead of promoted to an RPC-level error.

This bites on exactly the scenario your own new test covers: BEGIN, an INSERT against a deferred FK, then COMMIT via execute_query with a session_id. The COMMIT fails and (correctly, per your fix) closes the transaction — but the response for that failing call is just {"error": {...}}, no in_transaction at all. The test only catches the corrected state on the next call (SELECT 1 → in_transaction: false), which is why this slipped through. A host reading in_transaction off every response to drive a "transaction open" indicator gets no signal on the call that actually needs it most.

One more, lower confidence — flagging for awareness, not blocking

run_batch_in_session reuses a session's pinned connection unconditionally and ignores whatever conn_params the current call supplied. If a session_id ever got repointed at a different host/user/password mid-transaction, the statement would silently run against the stale connection with no error. I couldn't confirm the real host ever does this (Editor.tsx ties session_id to a tab and doesn't seem to repoint a tab's connection mid-transaction), so this may not be reachable in practice — just wanted it on your radar.

Happy to help move any of these along — I can write the fix for #1 (mirror sweep_idle's retain-and-skip pattern in release_all) and #2/#3 (either handle the WORK/TRANSACTION keyword explicitly, or the more robust fix: strip all comments before tokenizing, not just leading ones) and open a PR against your branch if that's useful, or just talk through approach if you'd rather drive it yourself. #1 and #2/#3 feel like the ones worth fixing before merge; #4 and the conn_params note could reasonably be fast-follows depending on how you'd like to sequence it. Just say the word on any of it.

[egertaia]

Thanks, and thanks for driving #1 live, that trace made it unambiguous. All four are fixed on the branch, each with a test that fails on the previous code.

1. release_all orphaning a busy session (66ee2f7). It now uses retain like sweep_idle: a slot a run still holds stays in the map, so the run pins into a tracked slot. As a structural backstop, Slot now has a Drop impl that closes a connection it still holds (Client::take) instead of returning it to the pool. So no path that drops a pinned slot can hand an open transaction to the next borrower, including ones we have not thought of. New test: release_all_keeps_tracking_a_slot_a_run_holds.

2 and 3. Classifier (54b5b13). I took the robust route: transaction_effect now reads its keywords through a small leading_keywords scanner. The scanner skips every -- and (nested) /* */ comment wherever it appears, and stops at a string literal, dollar quote or ;. The optional WORK / TRANSACTION after COMMIT/END/ROLLBACK/ABORT is dropped before matching. Your inputs are all in transaction_effect_reads_optional_noise_words_and_inner_comments: ROLLBACK WORK TO SAVEPOINT, ROLLBACK TRANSACTION TO SAVEPOINT, COMMIT -- and chain later if needed, START /* explicit */ TRANSACTION, plus a nested comment and PREPARE TRANSACTION 'and chain'.

4. in_transaction lost on a failed execute_query (df5e448). With a session_id, a failing statement is now a reply, { "error": ..., "in_transaction": ... }, rather than an RPC error, so the call that needs the signal carries it. Without a session_id the reply is unchanged. commit_that_fails_on_a_deferred_constraint_releases_the_session now asserts on the failing COMMIT's own reply. The host side is in tabularis#801 (ff132136): it reads that shape, remembers each session's last reported state, and answers session_in_transaction from it, so a plugin tab's TX badge now clears after a failed COMMIT too. That also closes a gap Kilo raised on #801.

conn_params on a reused connection: no change. The host's session id is the editor tab's id, and a tab stays on the connection it was opened for, so a pinned session is never repointed. I agree it is worth knowing if another host ever reuses session ids differently.

cargo test --lib (346), clippy -D warnings and fmt --check are clean. I have not rerun tests/live_db.rs, as I have no Postgres here, so a pass on your setup would be very welcome, especially your shutdown race.

[aesslinger]

@egertaia — thank you for the work on this, genuinely. This was a long back-and-forth and you tracked down and fixed every single thing that came up, including the two subtle ones (the busy-slot orphan and the classifier's comment/keyword handling) that needed a live repro to even pin down. The Drop for Slot backstop on #1 is a nice touch — it closes the gap for cases neither of us has thought of yet, not just the one we found.

I re-verified all four fixes independently rather than just reading the diffs:

• 1 (release_all orphaning a busy session): re-ran the exact same live race (in-flight BEGIN + pg_sleep on a session, shutdown fired concurrently) against the fixed binary. The slot now stays correctly tracked and the connection sits legitimately pinned, not orphaned. Also traced the Drop for Slot → Client::take path into deadpool's internals to confirm it genuinely detaches the connection before dropping it, so it can never re-enter the pool. And I pushed on the one loose end this raised — does a plain process exit (no explicit shutdown call) clean up a still-tracked session? — and confirmed it does, matching the "closes with the process" doc.
• 2/3 (classifier): reran my exact three failing assertions directly against the new transaction_effect() — all three now pass. The comment-aware tokenizer is the right fix, not a per-case patch.
• 4 (in_transaction on a failed execute_query): confirmed via the updated live_db.rs test, and — since this changes the wire contract (a failed statement is now a successful reply instead of an RPC error) — I pulled the host commit you referenced (tabularis#801, ff132136) to make sure the host actually round-trips it back into a proper error rather than silently swallowing it. It does.

Full suite: cargo test --lib 346/346, clippy/fmt clean, tests/live_db.rs 25/26 (the one failure is the same pre-existing, unrelated trust-auth artifact from every prior run on my scratch DB, not a regression).

This is ready to merge from my side. Nice work getting it across the line.

[aesslinger]

Closing this fork-based PR in favor of a fresh PR opened from TabularisDB/tabularis-postgresql-plugin:feat/pinned-transaction-session (commit 73745b0), which merges this work onto current main (with #119/#123) and resolves the conflicts that blocked this fork branch from merging.

All of @egertaia's original commits are preserved with full attribution in the replacement branch. The replacement is fully verified: 83/83 cross-repo parity, 356 unit tests, 30 live-DB tests (including the deferred-constraint session test, now confirmed passing against a live PostgreSQL instance), clippy/fmt/markdownlint clean.

No reflection on the work here — this is purely the fork-PR mechanics (the head branch can't be updated from the upstream side). The replacement PR will link here.