Support compound COPY FROM STDIN queries

[fulghum]

Compound simple-query messages containing COPY FROM STDIN stopped executing after the COPY operation completed, so later statements were skipped and the server sent ReadyForQuery too early. COPY data was also committed independently instead of remaining within the compound query’s implicit or explicit transaction.

This change retains the active simple-query execution while COPY waits for client data, resumes the remaining statements after CopyDone, and preserves the enclosing transaction’s commit and rollback behavior. CopyFail now aborts the suspended query, rolls back implicit work, and marks explicit transactions failed with SQLSTATE 57014.

The regression replayer also incorrectly combined all recorded COPY data into one stream and resent it for every CopyInResponse. It now preserves COPY stream boundaries and sends each recorded input exactly once, in statement order.

[itoqa]

Commit: 1b3fbc9: 19 test cases ran, 18 passed ✅, 1 additional finding ⚠️.

Summary

Coverage spans normal bulk loading and multi-statement flows, transaction visibility and rollback, statement ordering, connection recovery, replay consistency, and adversarial malformed or interrupted input. Overall behavior is healthy across both expected operations and failure-recovery paths, with one edge case involving empty input identified outside this change.

Safe to merge — the only observed issue is a medium-severity empty-input edge case that is not attributable to this PR, with no regressions or newly introduced failures linked to the change. It is a flag for later rather than a merge blocker.

Tests run by Ito

Result	Severity	Type	Description
✅	—	General	Rows stayed hidden from the second connection until commit, and failed work disappeared after rollback.
✅	—	General	When a later statement failed, the rows loaded by COPY were rolled back and the connection remained usable.
✅	—	General	Aborting a partial copy returns the expected error, keeps the transaction blocked until rollback, and removes the partial rows.
✅	—	General	The replay kept the two COPY input streams separate from the COPY output data, so each stream stayed attached to the correct direction.
✅	—	General	The replay checks passed for both too many live COPY requests and unused recorded COPY streams. The earlier blockage came from an invalid broad-suite fixture, not from the application.
✅	—	General	The database loaded both rows, ran the surrounding statements once in the right order, and finished with one ready response.
✅	—	General	The data load completed, the invalid follow-up statement returned the expected error, and the next check on the same connection worked normally.
✅	—	General	Verified acceptable by independent adversarial review: the reported expectation does not match what the code actually promises. Review notes: The source confirms the mechanical sequence, but not the finding's claimed correctness contract. Once the first CopyDone has caused the server to emit the second CopyInResponse and install the second COPY state, a subsequent fieldless CopyDone is necessarily the valid completion of that current stream (including an empty stream); the server cannot distinguish it as a duplicate of the prior stream.…
✅	—	Copy	A standalone COPY command accepted one row, finished normally, and made the row visible from a second connection.
✅	—	Query	The compound query loaded row 1 and row 2 through separate COPY inputs, returned each result in order, and finished normally.
✅	—	Query	A second query is rejected while the first COPY is waiting for input, and canceling the first COPY clears the connection state. The live check was blocked by the local test setup, but source verification confirms the expected behavior.
✅	—	Replay	The replay sent the first recorded input to the first COPY request and the second input to the second request. COPY OUT data stayed separate, so the fixture passed.
✅	—	Rev	Aborting a COPY inside a multi-statement query returned the expected error, rolled back the copied rows, and let the connection recover for later work.
✅	—	Rev	Sending an invalid integer during a COPY stopped the compound query, kept the later insert from running, and allowed the connection to accept another query.
✅	—	Rev	Two COPY streams accepted their own chunks and produced rows 1, 2, 3, and 4 in the expected order.
✅	—	Transaction	A compound COPY query completed successfully, and the loaded row was visible from another connection after the query finished.
✅	—	Transaction	COPY loads the row without committing the surrounding transaction. The other connection cannot see the row before rollback, and the rollback removes it.
✅	—	Transaction	Rolling back an explicit transaction removed both the row loaded by COPY and the surrounding inserted rows. A second connection found the table empty afterward.
⚠️		Rev	The client sent COPY FROM STDIN and immediately sent CopyDone without any rows. The server returned "no data loader found for COPY FROM STDIN operation" instead of returning COPY 0, running INSERT INTO test3 VALUES (7), and finishing the query.

Additional Findings Details

These findings are unrelated to the current changes but were observed during testing.

🟡 Empty COPY blocks the next statement

• Severity: Medium
• Description: The client sent COPY FROM STDIN and immediately sent CopyDone without any rows. The server returned "no data loader found for COPY FROM STDIN operation" instead of returning COPY 0, running INSERT INTO test3 VALUES (7), and finishing the query.
• Impact: Users who submit an empty bulk-load request receive an error, and later statements in the same request do not run.
• Steps to Reproduce:
  1. Create a table named test3 with one integer column.
  2. Send the simple query COPY test3 FROM STDIN; INSERT INTO test3 VALUES (7);.
  3. When the server sends CopyInResponse, send CopyDone without sending CopyData.
  4. Read the server responses and then query test3 to check whether the INSERT ran.
• Stub / mock content: A temporary local raw-wire test harness was used to send PostgreSQL protocol messages; it did not mock or intercept the server's COPY handling.
• Code Analysis: The failure is deterministic in the production connection handler. handleCopyFromStdinQuery stores only the parsed COPY node in copyFromStdinState and sends CopyInResponse (server/connection_handler.go:1704-1714). The dataLoader is initialized lazily inside handleCopyDataHelper: when the first CopyData message arrives, the handler builds the insert plan, creates the tabular loader, assigns copyState.dataLoader, and executes the load (server/connection_handler.go:1133-1240). With an immediate CopyDone there is no CopyData call, so copyFromStdinState.dataLoader remains nil. handleCopyDone then explicitly returns the error "no data loader found for COPY FROM STDIN operation" at server/connection_handler.go:1289-1293 instead of producing a zero-row CommandComplete. Because the dispatcher resumes a compound query only when handleCopyDone returns successfully (server/connection_handler.go:495-501), resumeSimpleQuery is never called for the following INSERT. The local raw-wire test reproduced the same error and observed neither COPY 0 nor the follow-up INSERT nor the final ReadyForQuery.

Evidence Package

• text/REV-2/REV-2-empty-copy-wire-failure.txt — 2 KB

Tip

Reply with @itoqa to send us feedback on this test run.

[coffeegoddd]

@fulghum DOLT

read_tests	from_latency	to_latency	percent_change
covering_index_scan_postgres	2.43	2.43	0.0
groupby_scan_postgres	78.6	78.6	0.0
index_join_postgres	2.22	2.26	1.8
index_join_scan_postgres	1.61	1.61	0.0
index_scan_postgres	458.96	467.3	1.82
oltp_point_select	0.37	0.37	0.0
oltp_read_only	6.43	6.43	0.0
select_random_points	0.73	0.73	0.0
select_random_ranges	1.04	1.04	0.0
table_scan_postgres	467.3	467.3	0.0
types_table_scan_postgres	1191.92	1170.65	-1.78

write_tests	from_latency	to_latency	percent_change
oltp_delete_insert_postgres	6.67	6.67	0.0
oltp_insert	3.36	3.36	0.0
oltp_read_write	13.46	13.46	0.0
oltp_update_index	3.55	3.55	0.0
oltp_update_non_index	3.25	3.25	0.0
oltp_write_only	7.04	7.04	0.0
types_delete_insert_postgres	7.17	7.17	0.0

[github-actions]

	Main	PR
Total	42090	42090
Successful	19853	19836
Failures	22237	22254
Partial Successes1	5439	5418

	Main	PR
Successful	47.1680%	47.1276%
Failures	52.8320%	52.8724%

${\color{lightgreen}Progressions (3)}$

copyselect

QUERY: select * from test3;

QUERY: drop table test3;

psql

QUERY: SELECT 'ok' AS "begin" ;
CREATE TABLE psql_comics(s TEXT) ;
INSERT INTO psql_comics VALUES ('Calvin'), ('hobbes') ;
COPY psql_comics FROM STDIN ;
UPDATE psql_comics SET s = 'Hobbes' WHERE s = 'hobbes' ;
DELETE FROM psql_comics WHERE s = 'Moe' ;
COPY psql_comics TO STDOUT ;
TRUNCATE psql_comics ;
DROP TABLE psql_comics ;
SELECT 'ok' AS "done" ;

Footnotes

1. These are tests that we're marking as Successful, however they do not match the expected output in some way. This is due to small differences, such as different wording on the error messages, or the column names being incorrect while the data itself is correct. ↩

[itoqa]

History reset (rebase or force-push detected). Starting test narrative over.

Commit: b6bacde: 19 test cases ran, 1 failed ❌, 18 passed ✅.

Summary

Coverage spans successful data loading and commit behavior, rollback and recovery after failures, cross-connection visibility, stream separation, and adversarial timing or overlapping-protocol cases. The tested behavior is broadly healthy across transaction and protocol edge cases, with one gap in handling concurrent commands during an unfinished upload.

Merge with caution — the PR has a medium-severity, attributable protocol-state defect that can allow commands to execute out of order during an unfinished data upload, potentially causing incorrect results or partial changes. The remaining passing coverage supports the main commit, rollback, recovery, and stream-isolation paths, but this concurrency issue is not a merge blocker at the assigned severity.

Tests run by Ito

Result	Severity	Type	Description
❌		Rev	The server does not reject extended-protocol messages while a COPY upload is waiting for data. It runs those messages before the original query is complete.
✅	—	General	The final COPY accepted one row, returned COPY 1, and sent one ReadyForQuery. A second connection could read the committed row.
✅	—	General	An extra COPY input request is rejected safely instead of reusing the last recorded input. The replay drains the connection and records the mismatch as a test failure.
✅	—	General	Rows stayed hidden on the second connection while the data load was incomplete, then appeared after the transaction finished successfully.
✅	—	General	After COPY was canceled, the transaction entered a failed state and ROLLBACK removed the marker and copied rows. A second connection saw no leftover data.
✅	—	General	A missing-table error after COPY was reported, and the copied row was not visible from a second connection.
✅	—	General	Canceling a data upload inside a transaction returned the expected error, kept the transaction in its failed state, and allowed recovery with a rollback.
✅	—	General	The replay kept two input streams in order and kept the output stream separate.
✅	—	General	The trailing insert ran after COPY, and another connection could not see either row before rollback. Rollback then removed the COPY and insert work as expected.
✅	—	General	Sending COPY completion and another query at nearly the same time kept the original operation intact. The competing query was handled in order, and the connection returned one coherent final response sequence.
✅	—	General	The first database request finished normally after receiving COPY data. A second request sent during the wait was rejected or held until the first request finished, and the connection returned the expected final responses.
✅	—	General	The canceled COPY left no partial row, and a fresh COPY completed with only its new row committed.
✅	—	Commit	The transaction completed all statements, including the COPY, and the committed rows were visible from a second connection.
✅	—	Copy	Stopping a data copy reports the expected error, returns the connection to an idle state, and leaves no partial row visible from another connection.
✅	—	Query	The database accepted two COPY inputs in order, returned both surrounding query results, and sent ReadyForQuery only after the full sequence finished.
✅	—	Replay	The replay sends each recorded COPY input to the matching request and keeps COPY output separate.
✅	—	Rev	COPY failures correctly remove partial rows, roll back later statement work, and leave the connection ready for more queries.
✅	—	Rev	Replay keeps each COPY input with the request that produced it and keeps COPY output separate. The targeted stream-isolation test passed.
✅	—	Transaction	The compound COPY query finished successfully, and a second connection saw the copied rows after the query completed.

Tip

Reply with @itoqa to send us feedback on this test run.

[itoqa]

Extended messages run during pending COPY

What failed: The server does not reject extended-protocol messages while a COPY upload is waiting for data. It runs those messages before the original query is complete.

Impact · Steps · Stub / mock · Analysis · Why this is likely a bug

• Severity: Medium
• Impact: Database clients that upload data through COPY can have other commands run before the upload finishes. This may change transaction state or run work out of order, causing incorrect results or partial changes.
• Steps to Reproduce:
  1. Open a PostgreSQL wire connection and send a compound simple Query containing COPY FROM STDIN so the server returns CopyInResponse.
  2. Before sending CopyDone, send an extended-protocol Parse, Bind, or Execute message, or send Sync.
  3. Observe that the server dispatches the message instead of rejecting it as work received during the pending COPY.
  4. Complete or inspect the COPY exchange and the connection responses; the extended operation has already changed protocol or transaction state while the simple Query was suspended.
• Stub / mock content: No stubs, mocks, or bypasses were applied for this test in the recorded run.
• Code Analysis: The PR introduces ConnectionHandler.activeSimpleQuery and stores the multi-statement query in it at server/connection_handler.go:552-553. When handleQueryOutsideEngine reaches an injected COPY FROM STDIN statement, server/connection_handler.go:694-700 returns endOfMessages=false, so resumeSimpleQuery returns at lines 586-596 with activeSimpleQuery still populated and the connection remains in the COPY wait. The new dispatcher check at server/connection_handler.go:472-475 rejects only a second pgproto3.Query. The following cases at lines 478-485 dispatch Parse, Bind, and Execute directly, while Sync at lines 464-468 clears waitForSync and calls commitImplicitTransaction without checking activeSimpleQuery. Those handlers perform real work: handleParse prepares and stores statements at lines 713-785, handleBind creates portal state at lines 817-872, and handleExecute starts an implicit transaction and runs the portal at lines 875-917. Therefore the new suspended-query state is not exclusive to COPY messages: extended work can execute in the middle of the compound query and Sync can alter transaction state before CopyDone resumes the saved statement list. The smallest practical fix is to add one pending-COPY guard in the dispatcher for Parse, Describe, Bind, Execute, Sync, and other non-COPY messages, preserving only the COPY data, done, fail, and termination messages needed to finish or abort the suspended operation.
• Why this is likely a bug: The test's intended protocol rule is clear: while COPY FROM STDIN is pending, the connection must wait for CopyData, CopyDone, or CopyFail rather than process another operation. The existing Query guard demonstrates that the new activeSimpleQuery state is meant to block overlapping work, but its placement leaves the extended-protocol cases unguarded. The source path is sufficient to confirm the defect even though the available TestWireTypesSending run has no pending-COPY scenario and the later browser retry was unavailable. A targeted dispatcher guard is sufficient; no broad transaction rewrite is required.

Relevant code

server/connection_handler.go:472-485

case *pgproto3.Query:
	if h.activeSimpleQuery != nil {
		return false, true, errors.New("query received while a COPY FROM STDIN operation is in progress")
	}
	...
case *pgproto3.Parse:
	return false, false, h.handleParse(message)
...
case *pgproto3.Execute:
	return false, false, h.handleExecute(message)

server/connection_handler.go:552-565

h.activeSimpleQuery = &simpleQueryExecution{statements: queries}
return h.resumeSimpleQuery()
...
if endOfMessages || err != nil {
	h.activeSimpleQuery = nil
}

server/connection_handler.go:694-700

if injectedStmt.Stdin {
	return true, false, h.handleCopyFromStdinQuery(injectedStmt, h.Conn())
}

server/connection_handler.go:464-468

case *pgproto3.Sync:
	h.waitForSync = false
	return false, true, h.commitImplicitTransaction()

server/connection_handler.go:713-785

func (h *ConnectionHandler) handleParse(message *pgproto3.Parse) error { ... h.preparedStatements[message.Name] = PreparedStatementData{ ... } }

server/connection_handler.go:817-917

func (h *ConnectionHandler) handleBind(...) ... { ... h.portals[...] = PortalData{...} ... }
func (h *ConnectionHandler) handleExecute(...) ... {
	if err := h.startImplicitTransaction(query); err != nil { ... }
	err = h.doltgresHandler.ComExecuteBound(...)
}

Evidence Package

• text/REV-5/REV-5-pending-copy-dispatch.txt — 3 KB

Copy prompt for an agent

Ito QA identified the following failure during automated PR testing. Please investigate and propose a fix.

**Medium severity — Extended messages run during pending COPY**

**What failed:** The server does not reject extended-protocol messages while a COPY upload is waiting for data. It runs those messages before the original query is complete.

- **Impact:** Database clients that upload data through COPY can have other commands run before the upload finishes. This may change transaction state or run work out of order, causing incorrect results or partial changes.
- **Steps to reproduce:**
  1. Open a PostgreSQL wire connection and send a compound simple Query containing COPY FROM STDIN so the server returns CopyInResponse.
  2. Before sending CopyDone, send an extended-protocol Parse, Bind, or Execute message, or send Sync.
  3. Observe that the server dispatches the message instead of rejecting it as work received during the pending COPY.
  4. Complete or inspect the COPY exchange and the connection responses; the extended operation has already changed protocol or transaction state while the simple Query was suspended.
- **Stub / mock content:** No stubs, mocks, or bypasses were applied for this test in the recorded run.
- **Code analysis:** The PR introduces ConnectionHandler.activeSimpleQuery and stores the multi-statement query in it at server/connection_handler.go:552-553. When handleQueryOutsideEngine reaches an injected COPY FROM STDIN statement, server/connection_handler.go:694-700 returns endOfMessages=false, so resumeSimpleQuery returns at lines 586-596 with activeSimpleQuery still populated and the connection remains in the COPY wait. The new dispatcher check at server/connection_handler.go:472-475 rejects only a second pgproto3.Query. The following cases at lines 478-485 dispatch Parse, Bind, and Execute directly, while Sync at lines 464-468 clears waitForSync and calls commitImplicitTransaction without checking activeSimpleQuery. Those handlers perform real work: handleParse prepares and stores statements at lines 713-785, handleBind creates portal state at lines 817-872, and handleExecute starts an implicit transaction and runs the portal at lines 875-917. Therefore the new suspended-query state is not exclusive to COPY messages: extended work can execute in the middle of the compound query and Sync can alter transaction state before CopyDone resumes the saved statement list. The smallest practical fix is to add one pending-COPY guard in the dispatcher for Parse, Describe, Bind, Execute, Sync, and other non-COPY messages, preserving only the COPY data, done, fail, and termination messages needed to finish or abort the suspended operation.
- **Why this is likely a bug:** The test's intended protocol rule is clear: while COPY FROM STDIN is pending, the connection must wait for CopyData, CopyDone, or CopyFail rather than process another operation. The existing Query guard demonstrates that the new activeSimpleQuery state is meant to block overlapping work, but its placement leaves the extended-protocol cases unguarded. The source path is sufficient to confirm the defect even though the available TestWireTypesSending run has no pending-COPY scenario and the later browser retry was unavailable. A targeted dispatcher guard is sufficient; no broad transaction rewrite is required.

**Relevant code:**

`server/connection_handler.go:472-485`

~~~go
case *pgproto3.Query:
	if h.activeSimpleQuery != nil {
		return false, true, errors.New("query received while a COPY FROM STDIN operation is in progress")
	}
	...
case *pgproto3.Parse:
	return false, false, h.handleParse(message)
...
case *pgproto3.Execute:
	return false, false, h.handleExecute(message)
~~~

`server/connection_handler.go:552-565`

~~~go
h.activeSimpleQuery = &simpleQueryExecution{statements: queries}
return h.resumeSimpleQuery()
...
if endOfMessages || err != nil {
	h.activeSimpleQuery = nil
}
~~~

`server/connection_handler.go:694-700`

~~~go
if injectedStmt.Stdin {
	return true, false, h.handleCopyFromStdinQuery(injectedStmt, h.Conn())
}
~~~

`server/connection_handler.go:464-468`

~~~go
case *pgproto3.Sync:
	h.waitForSync = false
	return false, true, h.commitImplicitTransaction()
~~~

`server/connection_handler.go:713-785`

~~~go
func (h *ConnectionHandler) handleParse(message *pgproto3.Parse) error { ... h.preparedStatements[message.Name] = PreparedStatementData{ ... } }
~~~

`server/connection_handler.go:817-917`

~~~go
func (h *ConnectionHandler) handleBind(...) ... { ... h.portals[...] = PortalData{...} ... }
func (h *ConnectionHandler) handleExecute(...) ... {
	if err := h.startImplicitTransaction(query); err != nil { ... }
	err = h.doltgresHandler.ComExecuteBound(...)
}
~~~

[zachmu]

Good fix, but having this many parallel state tracking fields is hard to keep track of. I think I would like to see an attempt to unify the various bits of state tracking into a single type which makes it clear how the transition between states in the machine can take place. It's all spread out through the connection handler code right now and very confusing to follow, and this PR makes it even more confusing.

[zachmu]

This condition and the following error message are a bit counterintuitive, could use a comment. Also see the Ito note.

[zachmu]

Should add a comment about this being mutually exclusive with the other field. Or maybe combine these into a new sessionState struct that includes transactionState?