From 661618dcc24bba1eac7befa771e04ed688dbc981 Mon Sep 17 00:00:00 2001 From: Piotr Fus Date: Thu, 6 Aug 2026 14:20:03 +0200 Subject: [PATCH] SNOW-2499292 Implement serialisable arrow streams --- arrow_stream_serializable.go | 184 +++++++++++++++++++ arrow_stream_serializable_test.go | 285 ++++++++++++++++++++++++++++++ 2 files changed, 469 insertions(+) create mode 100644 arrow_stream_serializable.go create mode 100644 arrow_stream_serializable_test.go diff --git a/arrow_stream_serializable.go b/arrow_stream_serializable.go new file mode 100644 index 000000000..bad6c463d --- /dev/null +++ b/arrow_stream_serializable.go @@ -0,0 +1,184 @@ +package gosnowflake + +import ( + "bufio" + "bytes" + "compress/gzip" + "context" + "fmt" + "io" + "maps" + "net/http" +) + +// SerializableArrowStreamBatch is a portable, self-contained description of an +// ArrowStreamBatch (see QueryArrowStream / ArrowStreamLoader). It can be marshaled +// (e.g. with encoding/json), shipped to another process or machine that has no Snowflake +// connection, and turned back into a readable stream with Reader. This enables distributed +// fetch: partition a result set on one node and download the chunks in parallel across many +// (e.g. from an ADBC driver's ExecutePartitions / ReadPartition). +// +// The descriptor carries only the chunk address and payload format; it deliberately does not +// carry column metadata or timezone. A consumer that needs to decode the bytes obtains +// RowTypes() and Location() from the ArrowStreamLoader and carries them alongside. +// +// Security: for a remote batch, Headers carries the SSE-C decryption key and URL is a +// presigned bearer capability granting read access to the chunk. Treat a serialized batch as +// sensitive and transmit it only over secure channels. +// +// Expiry: presigned URLs are valid only for a server-controlled window (typically a few +// hours). A remote batch must be shipped and fetched within it. Inline batches (InlineData +// set) embed their bytes and never expire. +type SerializableArrowStreamBatch struct { + // Format is the payload format of the stream, "arrow" or "json". + Format string `json:"format"` + // InlineData holds the batch bytes embedded in the query response (the first batch). + // When set, the batch is local and no download is performed. + InlineData []byte `json:"inlineData,omitempty"` + // URL is the presigned cloud-storage URL for a remote chunk (empty for inline batches). + URL string `json:"url,omitempty"` + // Headers are the HTTP headers required to fetch URL, including the SSE-C key. + Headers map[string]string `json:"headers,omitempty"` + // NumRows is the number of rows the metadata reported for this batch. + NumRows int64 `json:"numRows"` + // UncompressedSize is the reported uncompressed size of the chunk in bytes. + UncompressedSize int64 `json:"uncompressedSize,omitempty"` +} + +// ToSerializable converts a stream batch into a portable descriptor for distributed fetch. +// Call it before GetStream. The batch must have been produced by ArrowStreamLoader.GetBatches +// (i.e. be backed by a downloader). +func (asb *ArrowStreamBatch) ToSerializable() (SerializableArrowStreamBatch, error) { + if asb.scd == nil { + return SerializableArrowStreamBatch{}, fmt.Errorf( + "arrow stream batch is not backed by a downloader and cannot be serialized") + } + s := SerializableArrowStreamBatch{ + Format: asb.scd.queryResultFormat, + NumRows: asb.numrows, + } + if asb.inlineData != nil { + s.InlineData = asb.inlineData + return s, nil + } + meta := asb.scd.ChunkMetas[asb.idx] + s.URL = meta.URL + s.UncompressedSize = meta.UncompressedSize + s.Headers = streamChunkHeaders(asb.scd.ChunkHeader, asb.scd.Qrmk) + return s, nil +} + +// streamChunkHeaders builds the HTTP headers used to fetch a result chunk on the stream path. +// It mirrors downloadChunkStreamHelper: use the server-provided ChunkHeader map when present +// (GCS/Azure), otherwise the S3 SSE-C headers carrying the query result master key (qrmk). +func streamChunkHeaders(chunkHeader map[string]string, qrmk string) map[string]string { + headers := make(map[string]string) + if len(chunkHeader) > 0 { + maps.Copy(headers, chunkHeader) + } else { + headers[headerSseCAlgorithm] = headerSseCAes + headers[headerSseCKey] = qrmk + } + return headers +} + +// streamOpts holds options for reconstructing a stream reader on a worker node. +type streamOpts struct { + client *http.Client +} + +// StreamOption configures SerializableArrowStreamBatch.Reader. +type StreamOption func(*streamOpts) + +// WithStreamHTTPClient injects the HTTP client used to download a remote chunk. It is +// required for a batch that references a URL; the worker has no Snowflake connection to build +// one from. This is also the seam for configuring proxy/TLS/timeouts and transport retries. +func WithStreamHTTPClient(client *http.Client) StreamOption { + return func(o *streamOpts) { o.client = client } +} + +// ArrowStreamPartitionReader reads a single serialized stream batch on a worker node. Its +// GetStream mirrors ArrowStreamBatch.GetStream so it can be consumed the same way (fed into +// ipc.NewReader, or used through the ADBC driver's stream reader). +type ArrowStreamPartitionReader struct { + inlineData []byte + url string + headers map[string]string + client *http.Client +} + +// Reader reconstructs a readable stream from a serialized batch. A remote batch (URL set) +// requires an *http.Client via WithStreamHTTPClient; an inline batch needs none. +func (s SerializableArrowStreamBatch) Reader(opts ...StreamOption) (*ArrowStreamPartitionReader, error) { + var o streamOpts + for _, opt := range opts { + opt(&o) + } + remote := len(s.InlineData) == 0 && s.URL != "" + if remote && o.client == nil { + return nil, fmt.Errorf("gosnowflake: a remote arrow stream batch requires an *http.Client (use WithStreamHTTPClient)") + } + return &ArrowStreamPartitionReader{ + inlineData: s.InlineData, + url: s.URL, + headers: s.Headers, + client: o.client, + }, nil +} + +// GetStream returns a stream of bytes for the batch, downloading the chunk first if it is +// remote. The content is Arrow IPC or JSON depending on the batch's Format. Close the +// returned stream when done. +func (r *ArrowStreamPartitionReader) GetStream(ctx context.Context) (io.ReadCloser, error) { + if len(r.inlineData) > 0 { + return newCancelableStream(ctx, io.NopCloser(bytes.NewReader(r.inlineData))), nil + } + rr, err := openStandaloneChunkStream(ctx, r.client, r.url, r.headers) + if err != nil { + return nil, err + } + return newCancelableStream(ctx, rr), nil +} + +// openStandaloneChunkStream performs an HTTP GET against a presigned result-chunk URL using +// the injected client (no Snowflake connection) and returns a reader over the (gzip-inflated +// if needed) payload. It mirrors the transport handling in downloadChunkStreamHelper. The +// caller must Close the returned stream. +func openStandaloneChunkStream(ctx context.Context, client *http.Client, url string, headers map[string]string) (io.ReadCloser, error) { + if client == nil { + return nil, fmt.Errorf("gosnowflake: no http client provided for chunk download") + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, err + } + for k, v := range headers { + req.Header.Set(k, v) + } + resp, err := client.Do(req) + if err != nil { + return nil, err + } + if resp.StatusCode != http.StatusOK { + defer func() { _ = resp.Body.Close() }() + b, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) + return nil, &SnowflakeError{ + Number: ErrFailedToGetChunk, + SQLState: SQLStateConnectionFailure, + Message: fmt.Sprintf("failed to get chunk. HTTP: %d, body: %s", resp.StatusCode, bytes.TrimSpace(b)), + } + } + + bufStream := bufio.NewReader(resp.Body) + // Peek is best-effort: a short or empty payload simply skips gzip detection. + gzipMagic, _ := bufStream.Peek(2) + if len(gzipMagic) == 2 && gzipMagic[0] == 0x1f && gzipMagic[1] == 0x8b { + gz, err := gzip.NewReader(bufStream) + if err != nil { + _ = resp.Body.Close() + return nil, err + } + return &streamWrapReader{Reader: gz, wrapped: resp.Body}, nil + } + return &streamWrapReader{Reader: bufStream, wrapped: resp.Body}, nil +} diff --git a/arrow_stream_serializable_test.go b/arrow_stream_serializable_test.go new file mode 100644 index 000000000..44cfbe7c7 --- /dev/null +++ b/arrow_stream_serializable_test.go @@ -0,0 +1,285 @@ +package gosnowflake + +import ( + "bytes" + "compress/gzip" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/apache/arrow-go/v18/arrow/ipc" + + "github.com/snowflakedb/gosnowflake/v2/internal/query" +) + +func gzipStreamBytes(t *testing.T, in []byte) []byte { + t.Helper() + var buf bytes.Buffer + w := gzip.NewWriter(&buf) + if _, err := w.Write(in); err != nil { + t.Fatalf("gzip write: %v", err) + } + if err := w.Close(); err != nil { + t.Fatalf("gzip close: %v", err) + } + return buf.Bytes() +} + +func readAllAndClose(t *testing.T, rc io.ReadCloser) []byte { + t.Helper() + got, err := io.ReadAll(rc) + if err != nil { + t.Fatalf("read stream: %v", err) + } + if err := rc.Close(); err != nil { + t.Fatalf("close stream: %v", err) + } + return got +} + +func TestArrowStreamBatchToSerializable(t *testing.T) { + t.Run("remote SSE-C", func(t *testing.T) { + scd := &snowflakeArrowStreamChunkDownloader{ + ChunkMetas: []query.ExecResponseChunk{ + {URL: "u0"}, + {URL: "https://host/chunk1", UncompressedSize: 999}, + }, + Qrmk: "the-qrmk", + queryResultFormat: "arrow", + } + asb := &ArrowStreamBatch{idx: 1, numrows: 42, scd: scd} + s, err := asb.ToSerializable() + if err != nil { + t.Fatalf("ToSerializable: %v", err) + } + if s.URL != "https://host/chunk1" || s.UncompressedSize != 999 { + t.Fatalf("unexpected remote descriptor: %+v", s) + } + if s.Format != "arrow" || s.NumRows != 42 { + t.Fatalf("unexpected format/rows: %+v", s) + } + if s.Headers[headerSseCKey] != "the-qrmk" || s.Headers[headerSseCAlgorithm] != headerSseCAes { + t.Fatalf("SSE-C headers not set: %+v", s.Headers) + } + }) + + t.Run("remote ChunkHeader", func(t *testing.T) { + scd := &snowflakeArrowStreamChunkDownloader{ + ChunkMetas: []query.ExecResponseChunk{{URL: "https://host/gcs"}}, + ChunkHeader: map[string]string{"x-goog-meta-key": "bar"}, + queryResultFormat: "arrow", + } + asb := &ArrowStreamBatch{idx: 0, scd: scd} + s, err := asb.ToSerializable() + if err != nil { + t.Fatalf("ToSerializable: %v", err) + } + if s.Headers["x-goog-meta-key"] != "bar" { + t.Fatalf("ChunkHeader not propagated: %+v", s.Headers) + } + if _, ok := s.Headers[headerSseCKey]; ok { + t.Fatalf("SSE-C header should not be present when ChunkHeader is used: %+v", s.Headers) + } + }) + + t.Run("inline", func(t *testing.T) { + asb := &ArrowStreamBatch{ + numrows: 5, + inlineData: []byte("xyz"), + scd: &snowflakeArrowStreamChunkDownloader{queryResultFormat: "json"}, + } + s, err := asb.ToSerializable() + if err != nil { + t.Fatalf("ToSerializable: %v", err) + } + if !bytes.Equal(s.InlineData, []byte("xyz")) || s.Format != "json" || s.NumRows != 5 || s.URL != "" { + t.Fatalf("unexpected inline descriptor: %+v", s) + } + }) + + t.Run("no downloader errors", func(t *testing.T) { + asb := &ArrowStreamBatch{} + if _, err := asb.ToSerializable(); err == nil { + t.Fatal("expected error for a batch not backed by a downloader") + } + }) +} + +func TestArrowStreamPartitionReaderRemote(t *testing.T) { + payload := []byte("arrow-ipc-or-json-payload-bytes") + + for _, tc := range []struct { + name string + body []byte + }{ + {name: "plain", body: payload}, + {name: "gzip", body: gzipStreamBytes(t, payload)}, + } { + t.Run(tc.name, func(t *testing.T) { + var gotHeaders http.Header + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotHeaders = r.Header.Clone() + _, _ = w.Write(tc.body) + })) + defer srv.Close() + + s := SerializableArrowStreamBatch{ + Format: "arrow", + URL: srv.URL + "/chunk", + Headers: map[string]string{headerSseCKey: "qrmk"}, + NumRows: 1, + } + reader, err := s.Reader(WithStreamHTTPClient(srv.Client())) + if err != nil { + t.Fatalf("Reader: %v", err) + } + stream, err := reader.GetStream(context.Background()) + if err != nil { + t.Fatalf("GetStream: %v", err) + } + if got := readAllAndClose(t, stream); !bytes.Equal(got, payload) { + t.Fatalf("stream = %q, want %q", got, payload) + } + if gotHeaders.Get(headerSseCKey) != "qrmk" { + t.Fatalf("SSE-C header not forwarded to storage: %v", gotHeaders) + } + }) + } +} + +func TestArrowStreamPartitionReaderInline(t *testing.T) { + s := SerializableArrowStreamBatch{Format: "arrow", InlineData: []byte("inline-bytes"), NumRows: 1} + reader, err := s.Reader() // no client required for inline + if err != nil { + t.Fatalf("Reader: %v", err) + } + stream, err := reader.GetStream(context.Background()) + if err != nil { + t.Fatalf("GetStream: %v", err) + } + if got := readAllAndClose(t, stream); !bytes.Equal(got, []byte("inline-bytes")) { + t.Fatalf("stream = %q, want inline-bytes", got) + } +} + +func TestArrowStreamReaderRequiresClient(t *testing.T) { + s := SerializableArrowStreamBatch{Format: "arrow", URL: "https://example.invalid/chunk", NumRows: 1} + if _, err := s.Reader(); err == nil { + t.Fatal("expected error when a remote batch has no HTTP client") + } +} + +func TestArrowStreamReaderExpiredURL(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "AccessDenied: Request has expired", http.StatusForbidden) + })) + defer srv.Close() + + s := SerializableArrowStreamBatch{Format: "arrow", URL: srv.URL + "/chunk", NumRows: 1} + reader, err := s.Reader(WithStreamHTTPClient(srv.Client())) + if err != nil { + t.Fatalf("Reader: %v", err) + } + if _, err := reader.GetStream(context.Background()); err == nil { + t.Fatal("expected an error fetching an expired (403) URL") + } +} + +func TestSerializableArrowStreamBatchJSONRoundTrip(t *testing.T) { + orig := SerializableArrowStreamBatch{ + Format: "arrow", + InlineData: []byte{1, 2, 3, 4}, + URL: "https://example.invalid/chunk", + Headers: map[string]string{headerSseCKey: "k"}, + NumRows: 7, + UncompressedSize: 4096, + } + data, err := json.Marshal(orig) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var got SerializableArrowStreamBatch + if err := json.Unmarshal(data, &got); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if got.Format != orig.Format || !bytes.Equal(got.InlineData, orig.InlineData) || + got.URL != orig.URL || got.NumRows != orig.NumRows || got.UncompressedSize != orig.UncompressedSize || + got.Headers[headerSseCKey] != "k" { + t.Fatalf("descriptor did not round-trip: %+v", got) + } +} + +func TestQueryArrowStreamSerializableRoundTrip(t *testing.T) { + runSnowflakeConnTest(t, func(sct *SCTest) { + numrows := 50000 + query := fmt.Sprintf(selectRandomGenerator, numrows) + + loader, err := sct.sc.QueryArrowStream(sct.sc.ctx, query) + assertNilF(t, err) + if p, ok := loader.(QueryResultFormatProvider); ok { + assertEqualE(t, p.QueryResultFormat(), "arrow") + } + + batches, err := loader.GetBatches() + assertNilF(t, err) + assertTrueF(t, len(batches) > 0, "should have at least one batch") + + ctx := context.Background() + workerClient := &http.Client{} // a worker's client: no Snowflake connection + totalRows := 0 + sawRemoteChunk := false + + for i := range batches { + // Capture the descriptor before reading the original stream. + desc, err := batches[i].ToSerializable() + assertNilF(t, err) + + // Driver path: read the original chunk bytes. + origStream, err := batches[i].GetStream(ctx) + assertNilF(t, err) + origBytes, err := io.ReadAll(origStream) + assertNilF(t, err) + assertNilF(t, origStream.Close()) + + // Ship the descriptor as JSON and decode it on the "worker". + data, err := json.Marshal(desc) + assertNilF(t, err) + var decoded SerializableArrowStreamBatch + assertNilF(t, json.Unmarshal(data, &decoded)) + if decoded.URL != "" { + sawRemoteChunk = true + } + + // Distributed path: reconstruct and read with the bare client. + reader, err := decoded.Reader(WithStreamHTTPClient(workerClient)) + assertNilF(t, err) + reconStream, err := reader.GetStream(ctx) + assertNilF(t, err) + reconBytes, err := io.ReadAll(reconStream) + assertNilF(t, err) + assertNilF(t, reconStream.Close()) + + // The distributed path must reproduce the driver path exactly. + assertTrueF(t, bytes.Equal(origBytes, reconBytes), + fmt.Sprintf("batch %d: reconstructed bytes differ from original", i)) + + // Decode the reconstructed Arrow IPC end-to-end and count rows. + if len(reconBytes) > 0 { + rr, err := ipc.NewReader(bytes.NewReader(reconBytes)) + assertNilF(t, err) + for rr.Next() { + totalRows += int(rr.Record().NumRows()) + } + assertNilF(t, rr.Err()) + rr.Release() + } + } + + assertTrueF(t, sawRemoteChunk, "expected at least one remote (downloaded) chunk; increase numrows") + assertEqualE(t, totalRows, numrows) + }) +}