chunk_test.go

package gosnowflake

import (
	"bytes"
	"context"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strings"
	"sync"
	"testing"

	"github.com/apache/arrow-go/v18/arrow"
	"github.com/apache/arrow-go/v18/arrow/array"
	"github.com/apache/arrow-go/v18/arrow/memory"
	ia "github.com/snowflakedb/gosnowflake/v2/internal/arrow"
	sferrors "github.com/snowflakedb/gosnowflake/v2/internal/errors"
)

func TestBadChunkData(t *testing.T) {
	testDecodeErr(t, "")
	testDecodeErr(t, "null")
	testDecodeErr(t, "42")
	testDecodeErr(t, "\"null\"")
	testDecodeErr(t, "{}")

	testDecodeErr(t, "[[]")
	testDecodeErr(t, "[null]")
	testDecodeErr(t, `[[hello world]]`)

	testDecodeErr(t, `[[""hello world""]]`)
	testDecodeErr(t, `[["\"hello world""]]`)
	testDecodeErr(t, `[[""hello world\""]]`)
	testDecodeErr(t, `[["hello world`)
	testDecodeErr(t, `[["hello world"`)
	testDecodeErr(t, `[["hello world"]`)

	testDecodeErr(t, `[["\uQQQQ"]]`)

	for b := range byte(' ') {
		testDecodeErr(t, string([]byte{
			'[', '[', '"', b, '"', ']', ']',
		}))
	}
}

func TestValidChunkData(t *testing.T) {
	testDecodeOk(t, "[]")
	testDecodeOk(t, "[  ]")
	testDecodeOk(t, "[[]]")
	testDecodeOk(t, "[ [  ]   ]")
	testDecodeOk(t, "[[],[],[],[]]")
	testDecodeOk(t, "[[] , []  , [], []  ]")

	testDecodeOk(t, "[[null]]")
	testDecodeOk(t, "[[\n\t\r null]]")
	testDecodeOk(t, "[[null,null]]")
	testDecodeOk(t, "[[ null , null ]]")
	testDecodeOk(t, "[[null],[null],[null]]")
	testDecodeOk(t, "[[null],[ null  ] ,  [null]]")

	testDecodeOk(t, `[[""]]`)
	testDecodeOk(t, `[["false"]]`)
	testDecodeOk(t, `[["true"]]`)
	testDecodeOk(t, `[["42"]]`)

	testDecodeOk(t, `[[""]]`)
	testDecodeOk(t, `[["hello"]]`)
	testDecodeOk(t, `[["hello world"]]`)

	testDecodeOk(t, `[["/ ' \\ \b \t \n \f \r \""]]`)
	testDecodeOk(t, `[["❄"]]`)
	testDecodeOk(t, `[["\u2744"]]`)
	testDecodeOk(t, `[["\uFfFc"]]`)       // consume replacement chars
	testDecodeOk(t, `[["\ufffd"]]`)       // consume replacement chars
	testDecodeOk(t, `[["\u0000"]]`)       // yes, this is valid
	testDecodeOk(t, `[["\uD834\uDD1E"]]`) // surrogate pair
	testDecodeOk(t, `[["\uD834\u0000"]]`) // corrupt surrogate pair

	testDecodeOk(t, `[["$"]]`)      // "$"
	testDecodeOk(t, `[["\u0024"]]`) // "$"

	testDecodeOk(t, `[["\uC2A2"]]`) // "¢"
	testDecodeOk(t, `[["¢"]]`)      // "¢"

	testDecodeOk(t, `[["\u00E2\u82AC"]]`) // "€"
	testDecodeOk(t, `[["€"]]`)            // "€"

	testDecodeOk(t, `[["\uF090\u8D88"]]`) // "𐍈"
	testDecodeOk(t, `[["𐍈"]]`)            // "𐍈"
}

func TestSmallBufferChunkData(t *testing.T) {
	r := strings.NewReader(`[
	  [null,"hello world"],
	  ["foo bar", null],
	  [null, null] ,
	  ["foo bar",   "hello world" ]
	]`)

	lcd := largeChunkDecoder{
		r, 0, 0,
		0, 0,
		make([]byte, 1),
		bytes.NewBuffer(make([]byte, defaultStringBufferSize)),
		nil,
	}

	if _, err := lcd.decode(); err != nil {
		t.Fatalf("failed with small buffer: %s", err)
	}
}

func TestEnsureBytes(t *testing.T) {
	// the content here doesn't matter
	r := strings.NewReader("0123456789")

	lcd := largeChunkDecoder{
		r, 0, 0,
		3, 8189,
		make([]byte, 8192),
		bytes.NewBuffer(make([]byte, defaultStringBufferSize)),
		nil,
	}

	lcd.ensureBytes(4)

	// we expect the new remainder to be 3 + 10 (length of r)
	if lcd.rem != 13 {
		t.Fatalf("buffer was not refilled correctly")
	}
}

func testDecodeOk(t *testing.T, s string) {
	var rows [][]*string
	if err := json.Unmarshal([]byte(s), &rows); err != nil {
		t.Fatalf("test case is not valid json / [][]*string: %s", s)
	}

	// NOTE we parse and stringify the expected result to
	// remove superficial differences, like whitespace
	expect, err := json.Marshal(rows)
	if err != nil {
		t.Fatalf("unreachable: %s", err)
	}

	rows, err = decodeLargeChunk(strings.NewReader(s), 0, 0)
	if err != nil {
		t.Fatalf("expected decode to succeed: %s", err)
	}

	actual, err := json.Marshal(rows)
	if err != nil {
		t.Fatalf("json marshal failed: %s", err)
	}
	if string(actual) != string(expect) {
		t.Fatalf(`
		result did not match expected result
		  expect=%s
		   bytes=(%v)

		  acutal=%s
		   bytes=(%v)`,
			string(expect), expect,
			string(actual), actual,
		)
	}
}

func testDecodeErr(t *testing.T, s string) {
	if _, err := decodeLargeChunk(strings.NewReader(s), 0, 0); err == nil {
		t.Fatalf("expected decode to fail for input: %s", s)
	}
}

func TestEnableArrowBatches(t *testing.T) {
	runSnowflakeConnTest(t, func(sct *SCTest) {
		ctx := ia.EnableArrowBatches(sct.sc.ctx)
		numrows := 3000 // approximately 6 ArrowBatch objects

		pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
		defer pool.AssertSize(t, 0)
		ctx = WithArrowAllocator(ctx, pool)

		query := fmt.Sprintf(selectRandomGenerator, numrows)
		rows := sct.mustQueryContext(ctx, query, []driver.NamedValue{})
		defer rows.Close()

		// getting result batches via raw bridge
		info, err := rows.(*snowflakeRows).GetArrowBatches()
		if err != nil {
			t.Error(err)
		}
		batches := info.Batches
		numBatches := len(batches)
		maxWorkers := 10 // enough for 3000 rows
		type count struct {
			m       sync.Mutex
			recVal  int
			metaVal int
		}
		cnt := count{recVal: 0}
		var wg sync.WaitGroup
		chunks := make(chan int, numBatches)

		for w := 1; w <= maxWorkers; w++ {
			wg.Go(func() {
				for i := range chunks {
					batch := batches[i]
					var recs *[]arrow.Record
					if batch.Records != nil {
						recs = batch.Records
					} else if batch.Download != nil {
						var downloadErr error
						recs, _, downloadErr = batch.Download(context.Background())
						if downloadErr != nil {
							t.Error(downloadErr)
						}
					}
					if recs != nil {
						for _, r := range *recs {
							cnt.m.Lock()
							cnt.recVal += int(r.NumRows())
							cnt.m.Unlock()
							r.Release()
						}
					}
					cnt.m.Lock()
					cnt.metaVal += batch.RowCount
					cnt.m.Unlock()
				}
			})
		}
		for j := range numBatches {
			chunks <- j
		}
		close(chunks)

		wg.Wait()
		if cnt.recVal != numrows {
			t.Errorf("number of rows from records didn't match. expected: %v, got: %v", numrows, cnt.recVal)
		}
		if cnt.metaVal != numrows {
			t.Errorf("number of rows from arrow batch metadata didn't match. expected: %v, got: %v", numrows, cnt.metaVal)
		}
	})
}

func TestWithArrowBatchesAsync(t *testing.T) {
	runSnowflakeConnTest(t, func(sct *SCTest) {
		ctx := WithAsyncMode(sct.sc.ctx)
		ctx = ia.EnableArrowBatches(ctx)
		numrows := 50000

		pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
		defer pool.AssertSize(t, 0)
		ctx = WithArrowAllocator(ctx, pool)

		query := fmt.Sprintf(selectRandomGenerator, numrows)
		rows := sct.mustQueryContext(ctx, query, []driver.NamedValue{})
		defer rows.Close()

		info, err := rows.(*snowflakeRows).GetArrowBatches()
		if err != nil {
			t.Error(err)
		}
		batches := info.Batches
		numBatches := len(batches)
		maxWorkers := 10
		type count struct {
			m       sync.Mutex
			recVal  int
			metaVal int
		}
		cnt := count{recVal: 0}
		var wg sync.WaitGroup
		chunks := make(chan int, numBatches)

		for w := 1; w <= maxWorkers; w++ {
			wg.Go(func() {
				for i := range chunks {
					batch := batches[i]
					var recs *[]arrow.Record
					if batch.Records != nil {
						recs = batch.Records
					} else if batch.Download != nil {
						var downloadErr error
						recs, _, downloadErr = batch.Download(context.Background())
						if downloadErr != nil {
							t.Error(downloadErr)
						}
					}
					if recs != nil {
						for _, r := range *recs {
							cnt.m.Lock()
							cnt.recVal += int(r.NumRows())
							cnt.m.Unlock()
							r.Release()
						}
					}
					cnt.m.Lock()
					cnt.metaVal += batch.RowCount
					cnt.m.Unlock()
				}
			})
		}
		for j := range numBatches {
			chunks <- j
		}
		close(chunks)

		wg.Wait()
		if cnt.recVal != numrows {
			t.Errorf("number of rows from records didn't match. expected: %v, got: %v", numrows, cnt.recVal)
		}
		if cnt.metaVal != numrows {
			t.Errorf("number of rows from arrow batch metadata didn't match. expected: %v, got: %v", numrows, cnt.metaVal)
		}
	})
}

func TestWithArrowBatchesButReturningJSON(t *testing.T) {
	testWithArrowBatchesButReturningJSON(t, false)
}

func TestWithArrowBatchesButReturningJSONAsync(t *testing.T) {
	testWithArrowBatchesButReturningJSON(t, true)
}

func testWithArrowBatchesButReturningJSON(t *testing.T, async bool) {
	runSnowflakeConnTest(t, func(sct *SCTest) {
		requestID := NewUUID()
		pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
		defer pool.AssertSize(t, 0)
		ctx := WithArrowAllocator(context.Background(), pool)
		ctx = ia.EnableArrowBatches(ctx)
		ctx = WithRequestID(ctx, requestID)
		if async {
			ctx = WithAsyncMode(ctx)
		}

		sct.mustExec(forceJSON, nil)
		rows := sct.mustQueryContext(ctx, "SELECT 'hello'", nil)
		defer rows.Close()
		_, err := rows.(ia.BatchDataProvider).GetArrowBatches()
		assertNotNilF(t, err)
		var se *SnowflakeError
		assertTrueE(t, errors.As(err, &se))
		assertEqualE(t, se.Message, sferrors.ErrMsgNonArrowResponseInArrowBatches)
		assertEqualE(t, se.Number, ErrNonArrowResponseInArrowBatches)

		v := make([]driver.Value, 1)
		assertNilE(t, rows.Next(v))
		assertEqualE(t, v[0], "hello")
	})
}

func TestWithArrowBatchesMultistatement(t *testing.T) {
	testWithArrowBatchesMultistatement(t, false)
}

func TestWithArrowBatchesMultistatementAsync(t *testing.T) {
	testWithArrowBatchesMultistatement(t, true)
}

func testWithArrowBatchesMultistatement(t *testing.T, async bool) {
	runSnowflakeConnTest(t, func(sct *SCTest) {
		sct.mustExec("ALTER SESSION SET ENABLE_FIX_1758055_ADD_ARROW_SUPPORT_FOR_MULTI_STMTS = true", nil)
		pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
		defer pool.AssertSize(t, 0)
		ctx := WithMultiStatement(ia.EnableArrowBatches(WithArrowAllocator(context.Background(), pool)), 2)
		if async {
			ctx = WithAsyncMode(ctx)
		}
		driverRows := sct.mustQueryContext(ctx, "SELECT 'abc' UNION SELECT 'def' ORDER BY 1; SELECT 'ghi' UNION SELECT 'jkl' ORDER BY 1", nil)
		defer driverRows.Close()
		sfRows := driverRows.(SnowflakeRows)
		expectedResults := [][]string{{"abc", "def"}, {"ghi", "jkl"}}
		resultSetIdx := 0
		for hasNextResultSet := true; hasNextResultSet; hasNextResultSet = sfRows.NextResultSet() != io.EOF {
			info, err := driverRows.(ia.BatchDataProvider).GetArrowBatches()
			assertNilF(t, err)
			assertEqualF(t, len(info.Batches), 1)
			batch := info.Batches[0]
			assertNotNilF(t, batch.Records)
			records := *batch.Records
			assertEqualF(t, len(records), 1)
			record := records[0]
			defer record.Release()
			assertEqualF(t, record.Column(0).(*array.String).Value(0), expectedResults[resultSetIdx][0])
			assertEqualF(t, record.Column(0).(*array.String).Value(1), expectedResults[resultSetIdx][1])
			resultSetIdx++
		}
		assertEqualF(t, resultSetIdx, len(expectedResults))
		err := sfRows.NextResultSet()
		assertErrIsE(t, err, io.EOF)
	})
}

func TestWithArrowBatchesMultistatementWithJSONResponse(t *testing.T) {
	runSnowflakeConnTest(t, func(sct *SCTest) {
		sct.mustExec(forceJSON, nil)
		pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
		defer pool.AssertSize(t, 0)
		ctx := WithMultiStatement(ia.EnableArrowBatches(WithArrowAllocator(context.Background(), pool)), 2)
		driverRows := sct.mustQueryContext(ctx, "SELECT 'abc' UNION SELECT 'def' ORDER BY 1; SELECT 'ghi' UNION SELECT 'jkl' ORDER BY 1", nil)
		defer driverRows.Close()
		sfRows := driverRows.(SnowflakeRows)
		resultSetIdx := 0
		for hasNextResultSet := true; hasNextResultSet; hasNextResultSet = sfRows.NextResultSet() != io.EOF {
			_, err := driverRows.(ia.BatchDataProvider).GetArrowBatches()
			assertNotNilF(t, err)
			var se *SnowflakeError
			assertTrueF(t, errors.As(err, &se))
			assertEqualE(t, se.Number, ErrNonArrowResponseInArrowBatches)
			assertEqualE(t, se.Message, sferrors.ErrMsgNonArrowResponseInArrowBatches)
			resultSetIdx++
		}
		assertEqualF(t, resultSetIdx, 2)
		err := sfRows.NextResultSet()
		assertErrIsE(t, err, io.EOF)
	})
}

func TestWithArrowBatchesMultistatementWithLargeResultSet(t *testing.T) {
	runSnowflakeConnTest(t, func(sct *SCTest) {
		sct.mustExec("ALTER SESSION SET ENABLE_FIX_1758055_ADD_ARROW_SUPPORT_FOR_MULTI_STMTS = true", nil)
		pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
		defer pool.AssertSize(t, 0)
		ctx := WithMultiStatement(ia.EnableArrowBatches(WithArrowAllocator(context.Background(), pool)), 2)
		driverRows := sct.mustQueryContext(ctx, "SELECT 'abc' FROM TABLE(GENERATOR(ROWCOUNT => 1000000)); SELECT 'abc' FROM TABLE(GENERATOR(ROWCOUNT => 1000000))", nil)
		defer driverRows.Close()
		sfRows := driverRows.(SnowflakeRows)
		rowCount := 0
		for hasNextResultSet := true; hasNextResultSet; hasNextResultSet = sfRows.NextResultSet() != io.EOF {
			info, err := driverRows.(ia.BatchDataProvider).GetArrowBatches()
			assertNilF(t, err)
			assertTrueF(t, len(info.Batches) > 1)
			for _, batch := range info.Batches {
				var recs *[]arrow.Record
				if batch.Records != nil {
					recs = batch.Records
				} else if batch.Download != nil {
					recs, _, err = batch.Download(context.Background())
					assertNilF(t, err)
				}
				if recs != nil {
					for _, record := range *recs {
						defer record.Release()
						for i := 0; i < int(record.NumRows()); i++ {
							assertEqualF(t, record.Column(0).(*array.String).Value(i), "abc")
							rowCount++
						}
					}
				}
			}
		}
		err := sfRows.NextResultSet()
		assertErrIsE(t, err, io.EOF)
	})
}

func TestQueryArrowStream(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 loader.TotalRows() != int64(numrows) {
			t.Errorf("total numrows did not match expected, wanted %v, got %v", numrows, loader.TotalRows())
		}

		batches, err := loader.GetBatches()
		assertNilF(t, err)
		assertTrueF(t, len(batches) > 0, "should have at least one batch")
		assertTrueF(t, len(loader.RowTypes()) > 0, "should have row types")
	})
}

func TestQueryArrowStreamDescribeOnly(t *testing.T) {
	runSnowflakeConnTest(t, func(sct *SCTest) {
		numrows := 50000

		query := fmt.Sprintf(selectRandomGenerator, numrows)
		loader, err := sct.sc.QueryArrowStream(WithDescribeOnly(sct.sc.ctx), query)
		assertNilF(t, err, "failed to run query")

		if loader.TotalRows() != 0 {
			t.Errorf("total numrows did not match expected, wanted 0, got %v", loader.TotalRows())
		}

		if len(loader.RowTypes()) != 2 {
			t.Errorf("rowTypes length did not match expected, wanted 2, got %v", len(loader.RowTypes()))
		}
	})
}

func TestRetainChunkWOHighPrecision(t *testing.T) {
	runDBTest(t, func(dbt *DBTest) {
		var rows driver.Rows
		var err error

		err = dbt.conn.Raw(func(connection any) error {
			rows, err = connection.(driver.QueryerContext).QueryContext(ia.EnableArrowBatches(context.Background()), "select 0", nil)
			return err
		})
		assertNilF(t, err, "error running select 0 query")

		info, err := rows.(ia.BatchDataProvider).GetArrowBatches()
		assertNilF(t, err, "error getting arrow batch data")
		assertEqualF(t, len(info.Batches), 1, "should have one batch")

		records := info.Batches[0].Records
		assertNotNilF(t, records, "records should not be nil")

		numRecords := len(*records)
		assertEqualF(t, numRecords, 1, "should have exactly one record")

		record := (*records)[0]
		assertEqualF(t, len(record.Columns()), 1, "should have exactly one column")

		column := record.Column(0).(*array.Int8)
		row := column.Len()
		assertEqualF(t, row, 1, "should have exactly one row")

		int8Val := column.Value(0)
		assertEqualF(t, int8Val, int8(0), "value of cell should be 0")
	})
}

func TestQueryArrowStreamMultiStatement(t *testing.T) {
	runSnowflakeConnTest(t, func(sct *SCTest) {
		sct.mustExec("ALTER SESSION SET ENABLE_FIX_1758055_ADD_ARROW_SUPPORT_FOR_MULTI_STMTS = true", nil)
		ctx := WithMultiStatement(ia.EnableArrowBatches(sct.sc.ctx), 2)
		loader, err := sct.sc.QueryArrowStream(ctx, "SELECT 'abc'; SELECT 'abc' UNION SELECT 'def' ORDER BY 1")
		assertNilF(t, err)
		assertTrueF(t, len(loader.RowTypes()) > 0, "should have row types")
		assertTrueF(t, loader.TotalRows() > 0, "should have total rows")
	})
}

func TestQueryArrowStreamMultiStatementForJSONData(t *testing.T) {
	runSnowflakeConnTest(t, func(sct *SCTest) {
		ctx := WithMultiStatement(ia.EnableArrowBatches(sct.sc.ctx), 2)
		loader, err := sct.sc.QueryArrowStream(ctx, "SELECT 'abc'; SELECT 'abc'")
		assertNilF(t, err)
		assertTrueF(t, loader.TotalRows() > 0, "should return data")
	})
}

func TestQueryArrowStreamStoredProcedureLargeJSONResponse(t *testing.T) {
	runSnowflakeConnTest(t, func(sct *SCTest) {
		sct.mustExec(`CREATE OR REPLACE PROCEDURE test_arrow_stream_large_sp()
RETURNS TABLE ()
LANGUAGE SQL
AS
$$
DECLARE res RESULTSET;
BEGIN
  res := (
    SELECT SEQ4() AS id, 'item_' || SEQ4()::VARCHAR AS name
    FROM TABLE(GENERATOR(ROWCOUNT => 10000))
  );
  RETURN TABLE(res);
END;
$$`, nil)
		defer sct.mustExec("DROP PROCEDURE IF EXISTS test_arrow_stream_large_sp()", nil)

		loader, err := sct.sc.QueryArrowStream(sct.sc.ctx, "CALL test_arrow_stream_large_sp()")
		assertNilF(t, err)
		fp, ok := loader.(QueryResultFormatProvider)
		assertTrueF(t, ok, "loader should implement QueryResultFormatProvider")
		// currently SQL and JS stored procedure, which RETURNS TABLE, has "json" result format
		// with the inline rowset (JSONData) being empty, but the resultset returned in the chunks as JSON
		assertEqualE(t, fp.QueryResultFormat(), "json")
		assertTrueF(t, loader.TotalRows() == 10000, "should report 10000 total rows")
		assertEqualE(t, len(loader.JSONData()), 0)

		batches, err := loader.GetBatches()
		assertNilF(t, err)
		assertTrueF(t, len(batches) > 0, "large JSON result should have chunk batches")

		stream, err := batches[0].GetStream(sct.sc.ctx)
		assertNilF(t, err)
		defer stream.Close()

		body, err := io.ReadAll(stream)
		assertNilF(t, err)
		assertTrueF(t, len(body) > 0, "batch stream should contain data")

		wrapped := append([]byte("["), body...)
		wrapped = append(wrapped, ']')
		var rows [][]*string
		assertNilF(t, json.Unmarshal(wrapped, &rows))
		assertTrueF(t, len(rows) > 0, "batch should contain JSON rows")
		assertEqualE(t, len(rows[0]), 2)
	})
}

func TestQueryArrowStreamMultiStatementMixedFormats(t *testing.T) {
	runSnowflakeConnTest(t, func(sct *SCTest) {
		sct.mustExec("ALTER SESSION SET ENABLE_FIX_1758055_ADD_ARROW_SUPPORT_FOR_MULTI_STMTS = true", nil)
		ctx := WithMultiStatement(ia.EnableArrowBatches(sct.sc.ctx), 3)
		loader, err := sct.sc.QueryArrowStream(ctx, "SELECT 1; SHOW PARAMETERS; SELECT 2")
		assertNilF(t, err)

		fp, ok := loader.(QueryResultFormatProvider)
		assertTrueF(t, ok, "loader should implement QueryResultFormatProvider")

		expectedFormats := []string{"arrow", "json", "arrow"}
		for i, expectedFmt := range expectedFormats {
			format := fp.QueryResultFormat()
			assertEqualE(t, format, expectedFmt, fmt.Sprintf("result set %d: expected format %q, got %q", i, expectedFmt, format))

			if format == "arrow" {
				batches, batchErr := loader.GetBatches()
				assertNilF(t, batchErr)
				assertEqualE(t, len(batches), 1, fmt.Sprintf("result set %d: expected 1 arrow batch", i))
			} else {
				assertTrueF(t, len(loader.JSONData()) > 0 || len(loader.RowTypes()) > 0,
					fmt.Sprintf("result set %d: json result should have inline data or row types", i))
			}

			if i < len(expectedFormats)-1 {
				assertNilF(t, loader.NextResultSet(context.Background()))
			}
		}

		assertEqualE(t, loader.NextResultSet(context.Background()), io.EOF)
	})
}

func TestQueryArrowStreamArrowResponseExposesFormat(t *testing.T) {
	runSnowflakeConnTest(t, func(sct *SCTest) {
		loader, err := sct.sc.QueryArrowStream(sct.sc.ctx, "SELECT 1")
		assertNilF(t, err)
		assertTrueF(t, loader.TotalRows() > 0, "should have total rows")
		fp, ok := loader.(QueryResultFormatProvider)
		assertTrueF(t, ok, "loader should implement QueryResultFormatProvider")
		assertEqualE(t, fp.QueryResultFormat(), "arrow")
		batches, err := loader.GetBatches()
		assertNilF(t, err)
		assertTrueF(t, len(batches) > 0, "should have Arrow batches")
	})
}

func TestQueryArrowStreamResetAndRereadBatch(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, "QueryArrowStream should succeed")

		batches, err := loader.GetBatches()
		assertNilF(t, err, "GetBatches should succeed")
		assertTrueF(t, len(batches) > 0, "should have at least one batch")

		// Pick a batch that requires a download (skip index 0 which may be inline).
		idx := 0
		if len(batches) > 1 {
			idx = 1
		}

		// First read.
		stream1, err := batches[idx].GetStream(sct.sc.ctx)
		assertNilF(t, err, "first GetStream should succeed")
		defer stream1.Close()
		data1, err := io.ReadAll(stream1)
		assertNilF(t, err, "reading first stream should succeed")
		assertTrueF(t, len(data1) > 0, "first read should return data")

		// Reset the batch.
		assertNilF(t, batches[idx].Reset(), "Reset should succeed")

		// Re-read the same batch.
		stream2, err := batches[idx].GetStream(sct.sc.ctx)
		assertNilF(t, err, "second GetStream should succeed")
		defer stream2.Close()
		data2, err := io.ReadAll(stream2)
		assertNilF(t, err, "reading second stream should succeed")
		assertTrueF(t, len(data2) > 0, "re-read after Reset should return data")

		// Both reads should return the same content.
		assertBytesEqualE(t, data1, data2, "data should match after Reset + re-download")
	})
}