batchloader_test.go

package arrowbased

import (
	"bytes"
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	dbsqlerr "github.com/databricks/databricks-sql-go/errors"
	"github.com/databricks/databricks-sql-go/internal/cli_service"
	"github.com/databricks/databricks-sql-go/internal/config"
	"github.com/pkg/errors"

	"github.com/apache/arrow/go/v12/arrow"
	"github.com/apache/arrow/go/v12/arrow/array"
	"github.com/apache/arrow/go/v12/arrow/ipc"
	"github.com/apache/arrow/go/v12/arrow/memory"
	"github.com/stretchr/testify/assert"
)

func TestCloudFetchIterator(t *testing.T) {

	var handler func(w http.ResponseWriter, r *http.Request)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		handler(w, r)
	}))
	defer server.Close()

	t.Run("should fetch all the links", func(t *testing.T) {
		cloudFetchHeaders := map[string]string{
			"foo": "bar",
		}

		handler = func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
			for name, value := range cloudFetchHeaders {
				if values, ok := r.Header[name]; ok {
					if values[0] != value {
						panic(errors.New("Missing auth headers"))
					}
				}
			}
			_, err := w.Write(generateMockArrowBytes(generateArrowRecord()))
			if err != nil {
				panic(err)
			}
		}

		startRowOffset := int64(100)

		links := []*cli_service.TSparkArrowResultLink{
			{
				FileLink:       server.URL,
				ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
				StartRowOffset: startRowOffset,
				RowCount:       1,
				HttpHeaders:    cloudFetchHeaders,
			},
			{
				FileLink:       server.URL,
				ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
				StartRowOffset: startRowOffset + 1,
				RowCount:       1,
				HttpHeaders:    cloudFetchHeaders,
			},
		}

		cfg := config.WithDefaults()
		cfg.UseLz4Compression = false
		cfg.MaxDownloadThreads = 1

		bi, err := NewCloudBatchIterator(
			context.Background(),
			links,
			startRowOffset,
			nil,
			cfg,
		)
		if err != nil {
			panic(err)
		}

		// Access the internal structure through the wrapper
		wrapper, ok := bi.(*batchIterator)
		assert.True(t, ok)
		cbi, ok := wrapper.ipcIterator.(*cloudIPCStreamIterator)
		assert.True(t, ok)

		assert.True(t, bi.HasNext())
		assert.Equal(t, cbi.pendingLinks.Len(), len(links))
		assert.Equal(t, cbi.downloadTasks.Len(), 0)

		// get first link - should succeed
		sab1, err2 := bi.Next()
		if err2 != nil {
			panic(err2)
		}

		assert.Equal(t, cbi.pendingLinks.Len(), len(links)-1)
		assert.Equal(t, cbi.downloadTasks.Len(), 0)
		assert.Equal(t, sab1.Start(), startRowOffset)

		// get second link - should succeed
		sab2, err3 := bi.Next()
		if err3 != nil {
			panic(err3)
		}

		assert.Equal(t, cbi.pendingLinks.Len(), len(links)-2)
		assert.Equal(t, cbi.downloadTasks.Len(), 0)
		assert.Equal(t, sab2.Start(), startRowOffset+sab1.Count())

		// all links downloaded, should be no more data
		assert.False(t, bi.HasNext())
	})

	t.Run("should fail on expired link", func(t *testing.T) {
		handler = func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
			_, err := w.Write(generateMockArrowBytes(generateArrowRecord()))
			if err != nil {
				panic(err)
			}
		}

		startRowOffset := int64(100)

		links := []*cli_service.TSparkArrowResultLink{
			{
				FileLink:       server.URL,
				ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
				StartRowOffset: startRowOffset,
				RowCount:       1,
			},
			{
				FileLink:       server.URL,
				ExpiryTime:     time.Now().Add(-10 * time.Minute).Unix(), // expired link
				StartRowOffset: startRowOffset + 1,
				RowCount:       1,
			},
		}

		cfg := config.WithDefaults()
		cfg.UseLz4Compression = false
		cfg.MaxDownloadThreads = 1

		bi, err := NewCloudBatchIterator(
			context.Background(),
			links,
			startRowOffset,
			nil,
			cfg,
		)
		if err != nil {
			panic(err)
		}

		// Access the internal structure through the wrapper
		wrapper, ok := bi.(*batchIterator)
		assert.True(t, ok)
		cbi, ok := wrapper.ipcIterator.(*cloudIPCStreamIterator)
		assert.True(t, ok)

		assert.True(t, bi.HasNext())
		assert.Equal(t, cbi.pendingLinks.Len(), len(links))
		assert.Equal(t, cbi.downloadTasks.Len(), 0)

		// get first link - should succeed
		sab1, err2 := bi.Next()
		if err2 != nil {
			panic(err2)
		}

		assert.Equal(t, cbi.pendingLinks.Len(), len(links)-1)
		assert.Equal(t, cbi.downloadTasks.Len(), 0)
		assert.Equal(t, sab1.Start(), startRowOffset)

		// get second link - should fail
		_, err3 := bi.Next()
		assert.NotNil(t, err3)
		assert.ErrorContains(t, err3, dbsqlerr.ErrLinkExpired)
	})

	t.Run("should fail on HTTP errors", func(t *testing.T) {
		startRowOffset := int64(100)

		links := []*cli_service.TSparkArrowResultLink{
			{
				FileLink:       server.URL,
				ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
				StartRowOffset: startRowOffset,
				RowCount:       1,
			},
			{
				FileLink:       server.URL,
				ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
				StartRowOffset: startRowOffset + 1,
				RowCount:       1,
			},
		}

		cfg := config.WithDefaults()
		cfg.UseLz4Compression = false
		cfg.MaxDownloadThreads = 1

		bi, err := NewCloudBatchIterator(
			context.Background(),
			links,
			startRowOffset,
			nil,
			cfg,
		)
		if err != nil {
			panic(err)
		}

		// Access the internal structure through the wrapper
		wrapper, ok := bi.(*batchIterator)
		assert.True(t, ok)
		cbi, ok := wrapper.ipcIterator.(*cloudIPCStreamIterator)
		assert.True(t, ok)

		assert.True(t, bi.HasNext())
		assert.Equal(t, cbi.pendingLinks.Len(), len(links))
		assert.Equal(t, cbi.downloadTasks.Len(), 0)

		// set handler for the first link, which returns some data
		handler = func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
			_, err := w.Write(generateMockArrowBytes(generateArrowRecord()))
			if err != nil {
				panic(err)
			}
		}

		// get first link - should succeed
		sab1, err2 := bi.Next()
		if err2 != nil {
			panic(err2)
		}

		assert.Equal(t, cbi.pendingLinks.Len(), len(links)-1)
		assert.Equal(t, cbi.downloadTasks.Len(), 0)
		assert.Equal(t, sab1.Start(), startRowOffset)

		// set handler for the first link, which fails with some non-retryable HTTP error
		handler = func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusNotFound)
		}

		// get second link - should fail
		_, err3 := bi.Next()
		assert.NotNil(t, err3)
		assert.ErrorContains(t, err3, fmt.Sprintf("%s %d", "HTTP error", http.StatusNotFound))
	})

	t.Run("should use custom HTTPClient when provided", func(t *testing.T) {
		handler = func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
			_, err := w.Write(generateMockArrowBytes(generateArrowRecord()))
			if err != nil {
				panic(err)
			}
		}

		startRowOffset := int64(100)
		customHTTPClient := &http.Client{
			Transport: &http.Transport{MaxIdleConns: 10},
		}

		cfg := config.WithDefaults()
		cfg.UseLz4Compression = false
		cfg.MaxDownloadThreads = 1
		cfg.UserConfig.CloudFetchConfig.HTTPClient = customHTTPClient

		bi, err := NewCloudBatchIterator(
			context.Background(),
			[]*cli_service.TSparkArrowResultLink{{
				FileLink:       server.URL,
				ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
				StartRowOffset: startRowOffset,
				RowCount:       1,
			}},
			startRowOffset,
			nil,
			cfg,
		)
		assert.Nil(t, err)

		cbi := bi.(*batchIterator).ipcIterator.(*cloudIPCStreamIterator)
		assert.Equal(t, customHTTPClient, cbi.httpClient)

		// Verify fetch works
		sab, nextErr := bi.Next()
		assert.Nil(t, nextErr)
		assert.NotNil(t, sab)
	})

	t.Run("should fallback to http.DefaultClient when HTTPClient is nil", func(t *testing.T) {
		handler = func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
			_, err := w.Write(generateMockArrowBytes(generateArrowRecord()))
			if err != nil {
				panic(err)
			}
		}

		startRowOffset := int64(100)
		cfg := config.WithDefaults()
		cfg.UseLz4Compression = false
		cfg.MaxDownloadThreads = 1
		// Explicitly set HTTPClient to nil to verify fallback behavior
		cfg.UserConfig.CloudFetchConfig.HTTPClient = nil

		bi, err := NewCloudBatchIterator(
			context.Background(),
			[]*cli_service.TSparkArrowResultLink{{
				FileLink:       server.URL,
				ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
				StartRowOffset: startRowOffset,
				RowCount:       1,
			}},
			startRowOffset,
			nil,
			cfg,
		)
		assert.Nil(t, err)

		cbi := bi.(*batchIterator).ipcIterator.(*cloudIPCStreamIterator)
		assert.Equal(t, http.DefaultClient, cbi.httpClient)

		// Verify fetch works with default client
		sab, nextErr := bi.Next()
		assert.Nil(t, nextErr)
		assert.NotNil(t, sab)
	})
}

func TestCloudFetchSchemaOverride(t *testing.T) {
	// Reproduces ES-1804970: When the server result cache serves Arrow IPC files
	// from a prior query, the embedded schema has stale column names. The
	// authoritative schema from GetResultSetMetadata must override them.

	// IPC data has columns ["id", "name"] (stale, from cached query)
	staleRecord := generateArrowRecord()
	staleIPCBytes := generateMockArrowBytes(staleRecord)

	// Authoritative schema has columns ["x", "y"] (correct, from GetResultSetMetadata)
	correctFields := []arrow.Field{
		{Name: "x", Type: arrow.PrimitiveTypes.Int32},
		{Name: "y", Type: arrow.BinaryTypes.String},
	}
	correctSchema := arrow.NewSchema(correctFields, nil)
	var schemaBuf bytes.Buffer
	schemaWriter := ipc.NewWriter(&schemaBuf, ipc.WithSchema(correctSchema))
	if err := schemaWriter.Close(); err != nil {
		t.Fatal(err)
	}
	correctSchemaBytes := schemaBuf.Bytes()

	// Serve stale IPC data via mock HTTP
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, err := w.Write(staleIPCBytes)
		if err != nil {
			panic(err)
		}
	}))
	defer server.Close()

	t.Run("should override stale column names with authoritative schema", func(t *testing.T) {
		links := []*cli_service.TSparkArrowResultLink{
			{
				FileLink:       server.URL,
				ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
				StartRowOffset: 0,
				RowCount:       3,
			},
		}

		cfg := config.WithDefaults()
		cfg.UseLz4Compression = false
		cfg.MaxDownloadThreads = 1

		bi, err := NewCloudBatchIterator(
			context.Background(),
			links,
			0,
			correctSchemaBytes,
			cfg,
		)
		assert.Nil(t, err)

		batch, batchErr := bi.Next()
		assert.Nil(t, batchErr)
		assert.NotNil(t, batch)

		rec, recErr := batch.Next()
		assert.Nil(t, recErr)
		assert.NotNil(t, rec)

		// The record schema must use the authoritative names, not the stale ones
		assert.Equal(t, "x", rec.Schema().Field(0).Name)
		assert.Equal(t, "y", rec.Schema().Field(1).Name)

		// Data must be preserved
		assert.Equal(t, int64(3), rec.NumRows())
		assert.Equal(t, 2, len(rec.Schema().Fields()))

		rec.Release()
	})

	t.Run("should pass through unchanged when no override schema provided", func(t *testing.T) {
		links := []*cli_service.TSparkArrowResultLink{
			{
				FileLink:       server.URL,
				ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
				StartRowOffset: 0,
				RowCount:       3,
			},
		}

		cfg := config.WithDefaults()
		cfg.UseLz4Compression = false
		cfg.MaxDownloadThreads = 1

		bi, err := NewCloudBatchIterator(
			context.Background(),
			links,
			0,
			nil,
			cfg,
		)
		assert.Nil(t, err)

		batch, batchErr := bi.Next()
		assert.Nil(t, batchErr)

		rec, recErr := batch.Next()
		assert.Nil(t, recErr)

		// Without override, the original (stale) column names are preserved
		assert.Equal(t, "id", rec.Schema().Field(0).Name)
		assert.Equal(t, "name", rec.Schema().Field(1).Name)

		rec.Release()
	})
}

func generateArrowRecord() arrow.Record {
	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())

	fields := []arrow.Field{
		{Name: "id", Type: arrow.PrimitiveTypes.Int32},
		{Name: "name", Type: arrow.BinaryTypes.String},
	}
	schema := arrow.NewSchema(fields, nil)

	builder := array.NewRecordBuilder(mem, schema)
	defer builder.Release()

	builder.Field(0).(*array.Int32Builder).AppendValues([]int32{1, 2, 3}, nil)
	builder.Field(1).(*array.StringBuilder).AppendValues([]string{"one", "two", "three"}, nil)

	record := builder.NewRecord()

	return record
}

func generateMockArrowBytes(record arrow.Record) []byte {

	defer record.Release()

	var buf bytes.Buffer
	w := ipc.NewWriter(&buf, ipc.WithSchema(record.Schema()))
	if err := w.Write(record); err != nil {
		return nil
	}
	if err := w.Write(record); err != nil {
		return nil
	}
	if err := w.Close(); err != nil {
		return nil
	}
	return buf.Bytes()
}