doCopyFromStdin handles a CSV spool DoPut from the control plane. It writes streamed bytes to a worker-local tempfile, executes the COPY SQL with the placeholder substituted, and returns the row count in the standard DoPutUpdateResult AppMetadata so the existing client drain path continues to work.
( ctx context.Context, first *flight.FlightData, stream flight.FlightService_DoPutServer, )
| 40 | |
| 41 | func substituteCopyFromStdinPlaceholders(copySQL, tmpPath string, bytesWritten int64) string { |
| 42 | result := strings.ReplaceAll(copySQL, flightclient.CopyFromStdinPathPlaceholder, tmpPath) |
| 43 | return strings.ReplaceAll(result, flightclient.CopyFromStdinSizePlaceholder, strconv.FormatInt(bytesWritten, 10)) |
| 44 | } |
| 45 | |
| 46 | func copyFromStdinDescriptorRequest(desc *flight.FlightDescriptor) (string, []string, error) { |
| 47 | if desc == nil || len(desc.Cmd) == 0 { |
| 48 | return "", nil, status.Error(codes.InvalidArgument, "copy-from-stdin: missing COPY SQL in descriptor.Cmd") |
| 49 | } |
| 50 | if len(desc.Path) == 1 { |
| 51 | copySQL := string(desc.Cmd) |
| 52 | // Unversioned descriptors are the legacy text/CSV protocol and always |
| 53 | // carry a COPY statement. Native PostgreSQL binary loads use the |
| 54 | // versioned structured request below; reject any unversioned non-COPY |
| 55 | // command without inferring its format from table or function names. |
| 56 | fields := strings.Fields(copySQL) |
| 57 | if len(fields) == 0 || !strings.EqualFold(fields[0], "COPY") { |
| 58 | return "", nil, status.Error(codes.InvalidArgument, |
| 59 | "copy-from-stdin: native PostgreSQL binary COPY is missing exact schema metadata") |
| 60 | } |
| 61 | return copySQL, nil, nil |
| 62 | } |
| 63 | if len(desc.Path) != 2 || desc.Path[1] != flightclient.CopyFromStdinPostgresBinaryPathVersion { |
| 64 | return "", nil, status.Error(codes.InvalidArgument, "copy-from-stdin: unsupported binary COPY schema metadata") |
| 65 | } |
| 66 | var request sqlcore.CopyFromStdinRequest |
| 67 | if err := json.Unmarshal(desc.Cmd, &request); err != nil { |
| 68 | return "", nil, status.Errorf(codes.InvalidArgument, "copy-from-stdin: invalid binary COPY request: %v", err) |
| 69 | } |
| 70 | if request.SQLTemplate == "" || len(request.PostgresBinaryDatabaseTypeNames) == 0 { |
| 71 | return "", nil, status.Error(codes.InvalidArgument, "copy-from-stdin: incomplete binary COPY request") |
| 72 | } |
| 73 | return request.SQLTemplate, request.PostgresBinaryDatabaseTypeNames, nil |
| 74 | } |
| 75 | |
| 76 | type contextReader struct { |
| 77 | ctx context.Context |
| 78 | r io.Reader |
| 79 | } |
| 80 | |
| 81 | func (r *contextReader) Read(p []byte) (int, error) { |
| 82 | if err := r.ctx.Err(); err != nil { |
| 83 | return 0, err |
| 84 | } |
| 85 | return r.r.Read(p) |
| 86 | } |
| 87 | |
| 88 | func preparePostgresBinaryCopy(ctx context.Context, tmpPath string, databaseTypeNames []string) (string, int64, func(), error) { |
| 89 | if err := ctx.Err(); err != nil { |
| 90 | return "", 0, func() {}, err |
| 91 | } |
| 92 | schema, err := pgbinary.SchemaFromDatabaseTypes(databaseTypeNames) |
| 93 | if err != nil { |
| 94 | return "", 0, func() {}, err |
| 95 | } |
| 96 | |
| 97 | source, err := os.Open(tmpPath) |
| 98 | if err != nil { |
| 99 | return "", 0, func() {}, fmt.Errorf("open binary COPY spool: %w", err) |