From fe256e8301be1e257b552fca5851a69563a22399 Mon Sep 17 00:00:00 2001
From: Curt Hagenlocher <curt@hagenlocher.org>
Date: Sun, 5 Nov 2023 10:15:55 -0800
Subject: [PATCH 1/2] Support DescribeOnly for QueryArrowStream

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

diff --git a/chunk_test.go b/chunk_test.go
index 56b24354b..f6b3b20eb 100644
--- a/chunk_test.go
+++ b/chunk_test.go
@@ -611,3 +611,27 @@ 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)
+		if err != nil {
+			t.Error(err)
+		}
+
+		if loader.TotalRows() != 0 {
+			t.Errorf("total numrows did not match expected, wanted 0, got %v", loader.TotalRows())
+		}
+
+		batches, err := loader.GetBatches()
+		if err != nil {
+			t.Error(err)
+		}
+		if len(batches) != 0 {
+			t.Errorf("batches length did not match expected, wanted 0, got %v", len(batches))
+		}
+	})
+}
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 {

From 857b610fb8fa33e0d43e00394b16b25820229497 Mon Sep 17 00:00:00 2001
From: Curt Hagenlocher <curt@hagenlocher.org>
Date: Mon, 6 Nov 2023 05:09:50 -0800
Subject: [PATCH 2/2] PR feedback

---
 chunk_test.go | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/chunk_test.go b/chunk_test.go
index f6b3b20eb..989673ac8 100644
--- a/chunk_test.go
+++ b/chunk_test.go
@@ -618,20 +618,21 @@ func TestQueryArrowStreamDescribeOnly(t *testing.T) {
 
 		query := fmt.Sprintf(selectRandomGenerator, numrows)
 		loader, err := sct.sc.QueryArrowStream(WithDescribeOnly(sct.sc.ctx), query)
-		if err != nil {
-			t.Error(err)
-		}
+		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()
-		if err != nil {
-			t.Error(err)
-		}
+		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))
+		}
 	})
 }