From 8445dcaf46d9d92119f0a2179359e7b0256aecab Mon Sep 17 00:00:00 2001
From: Curt Hagenlocher <curt@hagenlocher.org>
Date: Mon, 6 Nov 2023 06:24:08 -0800
Subject: [PATCH] Support DescribeOnly for QueryArrowStream (#955)

---
 chunk_test.go | 25 +++++++++++++++++++++++++
 connection.go |  3 ++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/chunk_test.go b/chunk_test.go
index 56b24354b..989673ac8 100644
--- a/chunk_test.go
+++ b/chunk_test.go
@@ -611,3 +611,28 @@ func TestQueryArrowStream(t *testing.T) {
 		}
 	})
 }
+
+func TestQueryArrowStreamDescribeOnly(t *testing.T) {
+	runSnowflakeConnTest(t, func(sct *SCTest) {
+		numrows := 50000 // approximately 10 ArrowBatch objects
+
+		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())
+		}
+
+		batches, err := loader.GetBatches()
+		assertNilF(t, err, "failed to get result")
+		if len(batches) != 0 {
+			t.Errorf("batches length did not match expected, wanted 0, got %v", len(batches))
+		}
+
+		rowtypes := loader.RowTypes()
+		if len(rowtypes) != 2 {
+			t.Errorf("rowTypes length did not match expected, wanted 2, got %v", len(rowtypes))
+		}
+	})
+}
diff --git a/connection.go b/connection.go
index 5b39d1460..ef4abcf31 100644
--- a/connection.go
+++ b/connection.go
@@ -482,7 +482,8 @@ func (sc *snowflakeConn) GetQueryStatus(
 func (sc *snowflakeConn) QueryArrowStream(ctx context.Context, query string, bindings ...driver.NamedValue) (ArrowStreamLoader, error) {
 	ctx = WithArrowBatches(context.WithValue(ctx, asyncMode, false))
 	ctx = setResultType(ctx, queryResultType)
-	data, err := sc.exec(ctx, query, false, false /* isinternal */, false, bindings)
+	isDesc := isDescribeOnly(ctx)
+	data, err := sc.exec(ctx, query, false, false /* isinternal */, isDesc, bindings)
 	if err != nil {
 		logger.WithContext(ctx).Errorf("error: %v", err)
 		if data != nil {