Design: COPY IN Protocol

[SeanTAllen]

Add support for PostgreSQL's COPY ... FROM STDIN protocol for bulk data loading. Users initiate a COPY operation via a dedicated Session.copy_in() API, feed data chunks through send_copy_data(), and signal completion with finish_copy() or abort with abort_copy().

PR type: changelog - added (single type — new feature, no breaking changes).

Background

COPY table FROM STDIN is PostgreSQL's bulk loading mechanism — orders of magnitude faster than individual INSERTs. The protocol is fundamentally different from normal queries: after the initial command, the client streams data to the server.

Wire Protocol

Backend messages:

• CopyInResponse ('G'): Byte1('G') Int32(length) Int8(format) Int16(num_cols) Int16[](col_formats)
  • format: 0=text, 1=binary (overall format)
  • col_formats: per-column format codes (0=text, 1=binary)
  • Signals server is ready to receive data

Frontend messages:

• CopyData ('d'): Byte1('d') Int32(length) Byte[](data) — arbitrary chunk of copy data; message boundaries do NOT need to align with row boundaries
• CopyDone ('c'): Byte1('c') Int32(4) — signals successful end of data stream
• CopyFail ('f'): Byte1('f') Int32(length) String(error_message) — aborts the operation

Flow:

Client                              Server
COPY FROM STDIN query  ──────────→
                       ←────────── CopyInResponse ('G')
CopyData ('d')         ──────────→  (0 or more)
CopyData ('d')         ──────────→
CopyDone ('c')         ──────────→  (or CopyFail)
                       ←────────── CommandComplete ('C')  -- "COPY <count>"
                       ←────────── ReadyForQuery ('Z')

On error (malformed data, constraint violation), the server sends ErrorResponse + ReadyForQuery. After the server sends ErrorResponse, it ignores any further CopyData/CopyDone/CopyFail messages.

Current State

No COPY support exists. CopyInResponse ('G') falls through the parser to _UnsupportedMessage and is silently consumed. If a user sends a COPY FROM STDIN query via session.execute(), the session hangs — the server enters COPY mode waiting for data, while the driver waits for CommandComplete, and neither side makes progress.

Design

Public API

New behaviors on Session:

be copy_in(sql: String, receiver: CopyInReceiver)
  """
  Start a COPY ... FROM STDIN operation. The SQL string should be a COPY
  command with FROM STDIN. On success, the receiver's pg_copy_ready() is
  called, and the caller should then send data via send_copy_data(),
  finishing with finish_copy() or abort_copy().
  """

be send_copy_data(data: Array[U8] val)
  """
  Send a chunk of data to the server during a COPY IN operation. Data
  does not need to align with row boundaries — the server reassembles
  the stream. No-op if not in COPY IN mode.
  """

be finish_copy()
  """
  Signal successful completion of the COPY data stream. The server will
  validate the data and respond with pg_copy_complete() or
  pg_copy_failed(). No-op if not in COPY IN mode.
  """

be abort_copy(reason: String)
  """
  Abort the COPY operation with the given error message. The server will
  respond with pg_copy_failed(). No-op if not in COPY IN mode.
  """

New receiver interface:

interface tag CopyInReceiver
  """
  Receives the result of a Session.copy_in() call. The session is passed
  to each callback so receivers can execute follow-up operations.
  """
  be pg_copy_ready(session: Session)
    """
    Called when the server is ready to receive COPY data. The receiver
    should call session.send_copy_data() one or more times, then call
    session.finish_copy() when done, or session.abort_copy() to cancel.
    """

  be pg_copy_complete(session: Session, count: USize)
    """
    Called when the COPY operation completes successfully. The count is
    the number of rows copied.
    """

  be pg_copy_failed(session: Session,
    failure: (ErrorResponseMessage | ClientQueryError))
    """
    Called when the COPY operation fails. The failure is either a server
    error (ErrorResponseMessage) or a client-side error (ClientQueryError)
    such as the session being closed.
    """

Why a Separate API

COPY IN has fundamentally different semantics from regular queries: the client streams data to the server instead of receiving results. Using session.execute() would be ambiguous — the user would have no callback to know when to start sending data, and the result delivery (row count vs result set) is different. A dedicated copy_in + CopyInReceiver interface follows "distinct semantics deserve distinct representations."

Data Format

The driver sends raw bytes via send_copy_data(). The user is responsible for formatting data in the appropriate format (text or binary). For the default text format, this means tab-delimited columns, newline-terminated rows, \N for NULLs. This follows "it is easier to give than take away" — convenience helpers for text format encoding can be added later without changing the core API.

Internal Message Types

_backend_messages.pony — new message:

class val _CopyInResponseMessage
  """
  Message from the backend indicating it is ready to receive COPY data.
  Contains the overall format (0=text, 1=binary) and per-column format
  codes.
  """
  let format: U8
  let column_formats: Array[U8] val

  new val create(format': U8, column_formats': Array[U8] val) =>
    format = format'
    column_formats = column_formats'

Add _CopyInResponseMessage to _ResponseParserResult union.

Parser Change

Add to _response_parser.pony:

| 'G' =>
  buffer.skip(5)?
  let format = buffer.u8()?
  let num_cols = buffer.u16_be()?.usize()
  let col_fmts: Array[U8] iso = recover iso Array[U8](num_cols) end
  for i in Range(0, num_cols) do
    col_fmts.push(buffer.u16_be()?.u8())
  end
  return _CopyInResponseMessage(format, consume col_fmts)

Frontend Messages

Add to _frontend_message.pony:

fun copy_data(data: Array[U8] val): Array[U8] val =>
  """
  Build a CopyData message containing a chunk of COPY stream data.
  """

fun copy_done(): Array[U8] val =>
  """
  Build a CopyDone message signaling successful end of COPY data.
  """

fun copy_fail(reason: String): Array[U8] val =>
  """
  Build a CopyFail message aborting the COPY operation.
  """

State Machine Changes

Queue item:

class val _QueuedCopyIn
  let sql: String
  let receiver: CopyInReceiver

  new val create(sql': String, receiver': CopyInReceiver) =>
    sql = sql'
    receiver = receiver'

Add _QueuedCopyIn to _QueueItem union.

Session state interface (_SessionState):

Add four new methods:

fun ref copy_in(s: Session ref, sql: String, receiver: CopyInReceiver)
fun ref send_copy_data(s: Session ref, data: Array[U8] val)
fun ref finish_copy(s: Session ref)
fun ref abort_copy(s: Session ref, reason: String)
fun ref on_copy_in_response(s: Session ref, msg: _CopyInResponseMessage)

copy_in follows the execute/prepare pattern — each pre-auth state must call receiver.pg_copy_failed with the appropriate ClientQueryError:

• _SessionUnopened: receiver.pg_copy_failed(s, SessionNeverOpened)
• _SessionClosed: receiver.pg_copy_failed(s, SessionClosed)
• _SessionSSLNegotiating, _SessionConnected, _SessionSCRAMAuthenticating: receiver.pg_copy_failed(s, SessionNotAuthenticated)

send_copy_data, finish_copy, abort_copy default to no-op in _ConnectedState and _UnconnectedState (like cancel). _SessionSSLNegotiating also needs explicit no-op implementations for these three methods since it mixes in neither trait. on_copy_in_response defaults to _IllegalState() in _NotAuthenticated (like other post-auth callbacks).

_SessionLoggedIn:

copy_in queues a _QueuedCopyIn and tries to run. The three data methods (send_copy_data, finish_copy, abort_copy) check query_state and delegate only to _CopyInInFlight:

fun ref send_copy_data(s: Session ref, data: Array[U8] val) =>
  match query_state
  | let c: _CopyInInFlight => c.send_copy_data(s, data)
  end

fun ref finish_copy(s: Session ref) =>
  match query_state
  | let c: _CopyInInFlight => c.finish_copy(s)
  end

fun ref abort_copy(s: Session ref, reason: String) =>
  match query_state
  | let c: _CopyInInFlight => c.abort_copy(s, reason)
  end

fun ref on_copy_in_response(s: Session ref, msg: _CopyInResponseMessage) =>
  query_state.on_copy_in_response(s, this, msg)

_QueryReady dispatch:

| let ci: _QueuedCopyIn =>
  li.query_state = _CopyInInFlight.create()
  s._connection().send(_FrontendMessage.query(ci.sql))

COPY FROM STDIN is sent as a regular simple query. The server recognizes it and responds with CopyInResponse.

New query state: _CopyInInFlight

class _CopyInInFlight is _QueryState
  """
  COPY IN protocol in progress. Waits for CopyInResponse from the server,
  then accepts data from the client via send_copy_data/finish_copy/
  abort_copy. Ends on CommandComplete + ReadyForQuery (success) or
  ErrorResponse + ReadyForQuery (failure).
  """
  var _complete_count: USize = 0
  var _error: Bool = false

Callbacks:

Message	Handler
on_copy_in_response	Access receiver via li.query_queue(0)? as _QueuedCopyIn. Call receiver.pg_copy_ready(s).
on_error_response	Set _error = true. Access receiver via li.query_queue(0)? as _QueuedCopyIn. Call receiver.pg_copy_failed(s, msg).
on_command_complete	Store msg.value in _complete_count
on_ready_for_query	If not error: access receiver via li.query_queue(0)? as _QueuedCopyIn, call receiver.pg_copy_complete(s, _complete_count). Shift queue. If msg.idle(), transition to _QueryReady and try_run_query; else transition to _QueryNotReady.
on_data_row	shutdown (protocol anomaly)
on_row_description	shutdown (protocol anomaly)
on_empty_query_response	shutdown (protocol anomaly)
send_copy_data	s._connection().send(_FrontendMessage.copy_data(data))
finish_copy	s._connection().send(_FrontendMessage.copy_done())
abort_copy	s._connection().send(_FrontendMessage.copy_fail(reason))

_QueryState interface changes:

Add on_copy_in_response to the _QueryState interface. Default in _QueryNoQueryInFlight: li.shutdown(s) (protocol anomaly — CopyInResponse without a query in flight). Existing in-flight states (_SimpleQueryInFlight, _ExtendedQueryInFlight, _PrepareInFlight, _CloseStatementInFlight) must also implement on_copy_in_response as li.shutdown(s) — receiving CopyInResponse during any other operation is a protocol anomaly.

The send_copy_data, finish_copy, abort_copy methods are NOT on _QueryState — they're handled by _SessionLoggedIn via direct type matching on the query state (following the cancel pattern).

Query Cancellation

cancel() works automatically during COPY IN with no changes needed. The existing _SessionLoggedIn.cancel() fires for any in-flight query state (anything that isn't _QueryReady or _QueryNotReady). During COPY IN, the CancelRequest tells the server to abort the operation; the server then sends ErrorResponse + ReadyForQuery, which _CopyInInFlight.on_error_response and on_ready_for_query handle correctly.

Shutdown Handling

_SessionLoggedIn.on_shutdown needs to drain _QueuedCopyIn items:

| let ci: _QueuedCopyIn =>
  ci.receiver.pg_copy_failed(s, SessionClosed)

Router Change

Add to _response_message_parser.pony:

| let msg: _CopyInResponseMessage =>
  s.state.on_copy_in_response(s, msg)

Files Changed

1. New: postgres/copy_in_receiver.pony — CopyInReceiver interface
2. Edit: postgres/session.pony — add copy_in, send_copy_data, finish_copy, abort_copy behaviors on Session; add copy_in, send_copy_data, finish_copy, abort_copy, on_copy_in_response to _SessionState interface; add copy_in with pg_copy_failed callbacks to _SessionUnopened, _SessionClosed, _SessionSSLNegotiating, _SessionConnected, _SessionSCRAMAuthenticating; add send_copy_data/finish_copy/abort_copy no-ops to _ConnectedState, _UnconnectedState, and _SessionSSLNegotiating; add on_copy_in_response default _IllegalState() to _NotAuthenticated; add _QueuedCopyIn queue item; add _CopyInInFlight query state; update _SessionLoggedIn (queue dispatch, COPY methods, shutdown drain); add on_copy_in_response to _QueryState interface, _QueryNoQueryInFlight trait, and all existing in-flight states (_SimpleQueryInFlight, _ExtendedQueryInFlight, _PrepareInFlight, _CloseStatementInFlight) as li.shutdown(s)
3. Edit: postgres/_backend_messages.pony — add _CopyInResponseMessage; add to _ResponseParserResult
4. Edit: postgres/_response_parser.pony — parse CopyInResponse ('G')
5. Edit: postgres/_response_message_parser.pony — route _CopyInResponseMessage
6. Edit: postgres/_frontend_message.pony — add copy_data, copy_done, copy_fail builders
7. New: postgres/_test_copy_in.pony — mock server COPY IN tests and pre-auth state tests
8. Edit: postgres/_test_frontend_message.pony — unit tests for new frontend messages
9. Edit: postgres/_test_response_parser.pony — parser unit test for CopyInResponse
10. Edit: postgres/_test_query.pony — COPY IN integration tests
11. Edit: postgres/_test.pony — register new test classes
12. New: examples/copy-in/copy-in-example.pony — COPY IN example program
13. Edit: CLAUDE.md — update architecture, query execution flow, public API types, file layout, supported features, test organization

Tests

Frontend Message Unit Tests

In _test_frontend_message.pony:

• _TestFrontendMessageCopyData — verify CopyData wire format: type byte 'd', length, data payload
• _TestFrontendMessageCopyDone — verify CopyDone wire format: type byte 'c', length = 4
• _TestFrontendMessageCopyFail — verify CopyFail wire format: type byte 'f', length, null-terminated error string

Parser Unit Test

In _test_response_parser.pony:

• _TestResponseParserCopyInResponseMessage — parse CopyInResponse bytes with text format and multiple column format codes; verify fields.

Unit Tests (Mock Servers)

New test file: postgres/_test_copy_in.pony

Mock port assignments: 7688, 7689, 7690, 7691

• _TestCopyInSuccess (port 7688) — mock server authenticates, becomes ready. Client sends copy_in. Server responds with CopyInResponse. Client sends two CopyData chunks + CopyDone. Server verifies CopyDone received, responds with CommandComplete("COPY 2") + ReadyForQuery. Client verifies pg_copy_ready and pg_copy_complete(2) callbacks.

• _TestCopyInAbort (port 7689) — mock server authenticates, becomes ready. Client sends copy_in. Server responds with CopyInResponse. Client sends CopyFail. Server responds with ErrorResponse + ReadyForQuery. Client verifies pg_copy_ready and pg_copy_failed callbacks.

• _TestCopyInServerError (port 7690) — mock server authenticates, becomes ready. Client sends copy_in. Server responds with CopyInResponse. Client sends CopyData. Server responds with ErrorResponse (constraint violation) + ReadyForQuery. Client verifies pg_copy_failed fires and session remains usable (send a follow-up query).

• _TestCopyInShutdownDrainsCopyQueue (port 7691) — mock server authenticates but never becomes ready. Client queues copy_in. Close session. Verify pg_copy_failed with SessionClosed. Follows existing patterns _TestUnansweredQueriesFailOnShutdown and _TestPrepareShutdownDrainsPrepareQueue.

Pre-Auth State Tests

In _test_copy_in.pony (no mock server needed):

• _TestCopyInAfterSessionClosed — connect and authenticate, then close session. Call copy_in. Verify pg_copy_failed with SessionClosed. Follows Query/AfterSessionHasBeenClosed pattern.

Integration Tests

Added to postgres/_test_query.pony:

• _TestCopyInInsert — create a temp table, COPY 3 rows of tab-delimited text data, verify row count via SELECT.
• _TestCopyInAbortRollback — create a temp table, start COPY, send some data, abort with abort_copy, verify table remains empty.

Existing Tests

No regressions expected. The new 'G' parser case handles a message type that previously fell through to _UnsupportedMessage. The new queue item and query state are additive. The on_copy_in_response callback is added to interfaces with appropriate defaults.

Example

New example: examples/copy-in/copy-in-example.pony

Demonstrates:

• Creating a table with execute
• Using copy_in to bulk load text-format data
• Verifying the load with a SELECT query
• Console output showing row count

[SeanTAllen]

Backpressure for COPY IN

Problem

The current design has send_copy_data as a push-based behavior with no flow control. A caller loading a large dataset in a loop would queue unbounded messages in Session's mailbox, each carrying an Array[U8] val. Nothing tells the caller to slow down, so memory grows proportionally to the dataset size — defeating the purpose of streaming.

Decision: Add backpressure callbacks to CopyInReceiver

Add throttle/unthrottle callbacks to CopyInReceiver:

interface tag CopyInReceiver
  be pg_copy_ready(session: Session)
  be pg_copy_complete(session: Session, count: USize)
  be pg_copy_failed(session: Session,
    failure: (ErrorResponseMessage | ClientQueryError))

  be pg_throttled(session: Session)
    """
    Called when the session is experiencing TCP backpressure. The receiver
    should stop calling send_copy_data() until pg_unthrottled() is called.
    """

  be pg_unthrottled(session: Session)
    """
    Called when TCP backpressure is released. The receiver may resume
    calling send_copy_data().
    """

Internally, Session implements lori's _on_throttled()/_on_unthrottled() callbacks. During COPY IN (when query_state is _CopyInInFlight), these are forwarded to the CopyInReceiver. Outside of COPY IN, they're no-ops for now.

Lori's send() already returns (SendToken | SendErrorNotWriteable) to signal per-call pressure, and is_writeable() provides a synchronous check. The current session ignores these return values entirely (not just for COPY — all sends). Addressing that is a broader issue tracked separately; the throttle/unthrottle callbacks are sufficient for COPY IN flow control.

The caller is responsible for respecting throttle/unthrottle. This is cooperative backpressure — the same model used throughout Pony's ecosystem.

Alternatives considered

Pull-based model — instead of the caller pushing chunks, Session calls something like pg_send_next on the receiver after each chunk is written. The receiver provides the next chunk or signals done. This is inherently bounded but adds one actor round-trip per chunk, which is significant overhead for bulk loading where throughput matters. The backpressure approach allows the caller to send as fast as the network can consume without per-chunk synchronization.

Single payload — copy_in takes all the data upfront as one Array[U8] val. Simplest API but forces the caller to materialize the entire dataset in memory, which is exactly the memory bloat problem COPY IN's streaming protocol is designed to avoid. Rejected.

[SeanTAllen]

Reconsidering: the throttle/unthrottle approach could lead to surprising results for programmers who aren't familiar with backpressure systems, and could lead to livelocks. The pull-based model (Session calls pg_send_next on the receiver after each chunk) trades some performance and adds per-chunk round-trip overhead, but it's much easier to reason about — the programming model is inherently bounded without requiring the caller to understand cooperative flow control. That tradeoff is probably worth it.

Leaning toward pull-based for now.

[SeanTAllen]

Pull-Based Data Flow

Concrete sketch of the pull-based model from the previous comment.

Change from the Original Design

The original design uses a push model: after pg_copy_ready, the receiver calls send_copy_data in a loop, with no flow control. As identified above, this allows unbounded message accumulation in Session's mailbox.

The pull model inverts control: Session requests each chunk by calling pg_copy_ready after each send_copy_data is processed. The data flow is inherently bounded without throttle/unthrottle callbacks.

API Surface

The public API types are identical to the original design. The change is behavioral — the contract of pg_copy_ready is different.

CopyInReceiver:

interface tag CopyInReceiver
  be pg_copy_ready(session: Session)
    """
    Called when the session is ready to accept the next chunk of COPY data.
    Called once when the server enters COPY mode, and again after each
    send_copy_data() call.

    The receiver should respond by calling exactly one of:
    - session.send_copy_data(data) to send a chunk (triggers another
      pg_copy_ready)
    - session.finish_copy() to signal successful end of data
    - session.abort_copy(reason) to abort the operation

    The flow control guarantee depends on the receiver calling exactly one
    of these methods per pg_copy_ready. Calling send_copy_data without a
    preceding pg_copy_ready bypasses the bounded-memory guarantee.
    """

  be pg_copy_complete(session: Session, count: USize)
    """
    Called when the COPY operation completes successfully. The count is the
    number of rows copied.
    """

  be pg_copy_failed(session: Session,
    failure: (ErrorResponseMessage | ClientQueryError))
    """
    Called when the COPY operation fails.
    """

No pg_throttled/pg_unthrottled needed — the pull loop is the flow control.

Session behaviors (unchanged from original):

be copy_in(sql: String, receiver: CopyInReceiver)
be send_copy_data(data: Array[U8] val)
be finish_copy()
be abort_copy(reason: String)

Pull Loop

The loop lives in _CopyInInFlight. After sending each CopyData to the server, it calls pg_copy_ready to request the next chunk:

CopyInResponse received
  → receiver.pg_copy_ready(session)                 [first pull]

receiver calls session.send_copy_data(data)
  → send CopyData to server
  → receiver.pg_copy_ready(session)                 [next pull]

receiver calls session.finish_copy()
  → send CopyDone to server → wait for server response

receiver calls session.abort_copy(reason)
  → send CopyFail to server → wait for server response

Implementation Detail

send_copy_data, finish_copy, and abort_copy are dispatched by _SessionLoggedIn via direct type matching on query_state (following the cancel pattern — not on the _QueryState interface). _SessionLoggedIn is a thin dispatcher; the pull-loop logic lives in _CopyInInFlight, which receives li (as all _QueryState methods do) and accesses the receiver from the queue:

// _SessionLoggedIn — thin dispatcher
fun ref send_copy_data(s: Session ref, data: Array[U8] val) =>
  match query_state
  | let c: _CopyInInFlight => c.send_copy_data(s, this, data)
  end

// _CopyInInFlight — owns the pull loop
fun ref send_copy_data(s: Session ref, li: _SessionLoggedIn ref,
  data: Array[U8] val)
=>
  s._connection().send(_FrontendMessage.copy_data(data))
  try
    (li.query_queue(0)? as _QueuedCopyIn).receiver.pg_copy_ready(s)
  else
    _Unreachable()
  end

This is consistent with how on_copy_in_response (from the original design) issues the first pull — both the first and subsequent pulls go through _CopyInInFlight using the same queue access pattern that other in-flight states use for receiver access.

Backpressure Properties

At steady state, the pipeline contains at most:

• 1 pg_copy_ready message in the receiver's mailbox, OR
• 1 send_copy_data message in Session's mailbox

This is O(1) memory regardless of dataset size. The actor round-trip (Session → Receiver → Session) is the rate limiter.

If the network is slower than actor messaging, CopyData messages accumulate in lori's write buffer at the round-trip rate (microseconds per chunk) rather than all at once. For most workloads, the network keeps up and no accumulation occurs.

Performance

The pull model adds one extra actor message per chunk (the pg_copy_ready callback). For N chunks, that's N additional messages at microseconds each.

Chunk size is the receiver's choice. PostgreSQL's COPY protocol doesn't require message boundaries to align with rows, so receivers can batch many rows per send_copy_data call to amortize overhead.

Edge Cases

Server error during COPY (e.g., constraint violation): The server sends ErrorResponse + ReadyForQuery. A pg_copy_ready may already be in the receiver's mailbox from the preceding send_copy_data. The receiver processes the stale pg_copy_ready, calls send_copy_data, which is a no-op (query state has already transitioned). The receiver then processes pg_copy_failed. This is standard asynchronous message ordering.

Cancellation: Unchanged from the original design. cancel() sends CancelRequest; the server responds with ErrorResponse + ReadyForQuery. A pending pg_copy_ready becomes stale — same handling as server errors.

Shutdown: Unchanged. Queued _QueuedCopyIn items drain with pg_copy_failed(SessionClosed). A pending pg_copy_ready from an in-flight COPY may trigger a stale send_copy_data, which is a no-op.

Example

actor BulkLoader is CopyInReceiver
  let _env: Env
  let _rows: Array[String] val
  var _index: USize = 0

  new create(env: Env, rows: Array[String] val) =>
    _env = env
    _rows = rows

  be pg_copy_ready(session: Session) =>
    if _index < _rows.size() then
      try
        session.send_copy_data(_rows(_index)?.array())
        _index = _index + 1
      end
    else
      session.finish_copy()
    end

  be pg_copy_complete(session: Session, count: USize) =>
    _env.out.print("Loaded " + count.string() + " rows")

  be pg_copy_failed(session: Session,
    failure: (ErrorResponseMessage | ClientQueryError)) =>
    _env.out.print("COPY failed")

Relationship to Original Design

Everything not mentioned above is unchanged: internal message types (_CopyInResponseMessage, _QueuedCopyIn), parser changes, frontend messages (copy_data, copy_done, copy_fail), router changes, state machine additions (pre-auth state handlers, _QueryState interface changes, shutdown drain), test categories, example program, file list.

[SeanTAllen]

Implemented in #112.