Method CopyFromStdin

server/flightclient/copyfromstdin.go:43–117

1 MiB CopyFromStdin streams CSV bytes from r to a remote worker, then runs copySQL on the worker against a worker-local spool file. copySQL must contain CopyFromStdinPathPlaceholder where the file path should appear. Returns the number of rows the worker reports COPY-affected. Wire layout: frame

(ctx context.Context, copySQL string, r io.Reader)

// from the control plane to the worker during a COPY upload. Large
// enough to amortise per-frame overhead, small enough to keep memory
// bounded if the worker is slow to drain.
const CopyFromStdinChunkSize = 1 << 20 // 1 MiB

// CopyFromStdin streams COPY input bytes from r to a remote worker, then runs
// request.SQLTemplate against a worker-local spool file. The SQL template must
// contain CopyFromStdinPathPlaceholder where the file path should appear. It
// returns the number of rows the worker reports COPY-affected.
//
// Wire layout:
//
//	frame 0: FlightDescriptor{Type=PATH, Path=[CopyFromStdinDescriptorPath, optional version], Cmd=SQL or structured request}
//	frame 1..N: DataBody=<chunk of COPY input bytes>
//	(client closes send)
//	server: PutResult{AppMetadata=DoPutUpdateResult{RecordCount=N}}
func (e *FlightExecutor) CopyFromStdin(ctx context.Context, request sqlcore.CopyFromStdinRequest, r io.Reader) (int64, error) {
	if e.dead.Load() {
		return 0, ErrWorkerDead
	}

	reqCtx, cancel := e.mergedContext(ctx)
	defer cancel()
	reqCtx = e.withSession(reqCtx)

	stream, err := e.client.Client.DoPut(reqCtx)
	if err != nil {
		return 0, fmt.Errorf("flight doput: %w", err)
	}

	// Frame 0: descriptor + COPY SQL.
	path := []string{CopyFromStdinDescriptorPath}
	cmd := []byte(request.SQLTemplate)
	if len(request.PostgresBinaryDatabaseTypeNames) > 0 {
		path = append(path, CopyFromStdinPostgresBinaryPathVersion)
		cmd, err = json.Marshal(request)
		if err != nil {
			return 0, fmt.Errorf("marshal binary COPY request: %w", err)
		}
	}
	desc := &flight.FlightDescriptor{
		Type: flight.DescriptorPATH,
		Path: path,
		Cmd:  cmd,
	}
	if err := stream.Send(&flight.FlightData{FlightDescriptor: desc}); err != nil {
		return 0, fmt.Errorf("flight doput send descriptor: %w", err)
	}

	// Frames 1..N: COPY input byte chunks. We re-use one buffer; gRPC marshals the
	// proto synchronously inside Send, so reusing is safe across iterations.
	//
	// On a non-EOF read error we return immediately WITHOUT calling
	// CloseSend. The deferred cancel() then aborts the gRPC stream, so the
	// worker's Recv fails with Canceled and the worker bails before
	// running COPY on a partial spool file. This is what makes wire-level
	// COPY cancellation correct end-to-end.
	buf := make([]byte, CopyFromStdinChunkSize)
	for {
		n, readErr := r.Read(buf)