MCPcopy Create free account
hub / github.com/PostHog/duckgres / handleCopyInRemoteStreaming

Method handleCopyInRemoteStreaming

server/conn_copy.go:743–807  ·  view source on GitHub ↗

handleCopyInRemoteStreaming handles COPY FROM STDIN when the executor is remote (e.g. Flight, multitenant K8s worker). The control plane does not share a filesystem with the worker pod, so the legacy "spool to local tmp, run COPY FROM " approach fails with "No files found". This path streams t

(
	query string,
	opts *CopyFromOptions,
	copyStartTime time.Time,
	streamer sqlcore.CopyFromStdinExecutor,
)

Source from the content-addressed store, hash-verified

741 blobColIndices = append(blobColIndices, i)
742 }
743 }
744 if len(blobColIndices) > 0 {
745 c.logger().Debug("COPY FROM STDIN: table has BLOB columns, using CSV parse fallback.", "blob_columns", len(blobColIndices))
746 return c.handleCopyInCSVWithBlob(query, opts, cols, colTypes, blobColIndices)
747 }
748
749 // Send CopyInResponse
750 if err := wire.WriteCopyInResponse(c.writer, int16(len(cols)), true); err != nil {
751 return err
752 }
753 _ = c.flushWriter()
754 c.logger().Debug("COPY FROM STDIN sent CopyInResponse, waiting for data.")
755
756 // Remote-worker (Flight) executors implement CopyFromStdinExecutor so the
757 // CSV bytes are streamed to the worker pod via DoPut and spooled to the
758 // worker's filesystem there. The legacy local-tempfile path below works
759 // only when CP and worker share a filesystem (standalone / process
760 // backend), so prefer the streaming path when it's available.
761 if streamer, ok := c.executor.(sqlcore.CopyFromStdinExecutor); ok {
762 copySQL := BuildDuckDBCopyFromSQL(tableName, columnList, flightclient.CopyFromStdinPathPlaceholder, opts)
763 return c.handleCopyInRemoteStreaming(query, sqlcore.CopyFromStdinRequest{SQLTemplate: copySQL}, false, copyStartTime, streamer)
764 }
765
766 // Create temp file upfront and stream data directly to it (avoids memory buffering).
767 // This approach leverages DuckDB's highly optimized CSV parser which handles
768 // type conversions automatically and can load millions of rows in seconds.
769 tmpFile, err := os.CreateTemp("", "duckgres-copy-*.csv")
770 if err != nil {
771 c.logger().Error("COPY FROM STDIN failed to create temp file.", "error", err)
772 errMsg := fmt.Sprintf("failed to create temp file: %v", err)
773 c.sendError("ERROR", "58000", errMsg)
774 c.setTxError()
775 c.logQuery(copyStartTime, query, query, "COPY", 0, 0, "58000", errMsg, "simple")
776 _ = c.writeReadyForQuery(c.txStatus)
777 _ = c.flushWriter()
778 return nil
779 }
780 tmpPath := tmpFile.Name()
781 defer func() { _ = os.Remove(tmpPath) }()
782
783 // Stream COPY data directly to temp file (no memory buffering)
784 rowCount := 0
785 copyDataMessages := 0
786 bytesWritten := int64(0)
787 dataReceiveStart := time.Now()
788
789 for {
790 c.armIdleReadDeadline() // per-message: don't kill an actively-streaming COPY
791 msgType, body, err := wire.ReadMessage(c.reader)
792 if err != nil {
793 c.logger().Error("COPY FROM STDIN error reading message.", "error", err)
794 _ = tmpFile.Close()
795 return err
796 }
797
798 switch msgType {
799 case wire.MsgCopyData:
800 // Skip the PostgreSQL text COPY end-of-data marker (\.\n).

Callers 1

handleCopyInMethod · 0.95

Calls 13

loggerMethod · 0.95
logQueryStartedMethod · 0.95
logQueryFinishedMethod · 0.95
sendErrorMethod · 0.95
setTxErrorMethod · 0.95
logQueryMethod · 0.95
writeReadyForQueryMethod · 0.95
flushWriterMethod · 0.95
writeCommandCompleteMethod · 0.95
BuildDuckDBCopyFromSQLFunction · 0.85
NowMethod · 0.80
CopyFromStdinMethod · 0.65

Tested by

no test coverage detected