nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: snowflakedb/gosnowflake
Pull request #955 (2023-11-05 18:27:20)

Support DescribeOnly for QueryArrowStream
### Description
Modify snowflakeConn.QueryArrowStream to take the "describe only" flag on the context into account.

### Checklist
- [x] Code compiles correctly
- [x] Run ``make fmt`` to fix inconsistent formats
- [x] Run ``make lint`` to get lint errors and fix all of them
- [x] Created tests which fail without the change (if possible)
- [x] All tests passing
- [ ] Extended the README / documentation, if necessary

References #954 


Patch:
diff --git a/connection.go b/connection.go
index 5b39d14..ef4abcf 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 {


Test patch:
diff --git a/chunk_test.go b/chunk_test.go
index 56b2435..989673a 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))
+		}
+	})
+}
Repository: snowflakedb/gosnowflake
Pull request #1049 (2024-02-07 01:13:10)

fix a tz bug when validate ErrTooHighTimestampPrecision
### Description
We got error `268003 (22007): Cannot convert timestamp 2021-01-01 01:24:54 +0000 UTC in column Date to Arrow.Timestamp data type due to too high precision. Please use context with WithOriginalTimestamp.` However this should be within the range that nanosecond supports. 

Notice that the date is the first day of the year, that points to where the problem is. The right side is always in UTC while the left side is from `arrowSnowflakeTimestampToTime()` and can have a timezone. 

### Checklist
- [ ] Code compiles correctly
- [ ] Run ``make fmt`` to fix inconsistent formats
- [ ] Run ``make lint`` to get lint errors and fix all of them
- [ ] Created tests which fail without the change (if possible)
- [ ] All tests passing
- [ ] Extended the README / documentation, if necessary


Yifeng-Sigma: @sfc-gh-pfus could you please re-trigger the tests? I think they might be flaky?
sfc-gh-dheyman-1: Hi @Yifeng-Sigma, I rerun the tests and it seems that everything is passing. 
Yifeng-Sigma: > Hi @Yifeng-Sigma, I rerun the tests and it seems that everything is passing.

@sfc-gh-dheyman-1 Thank you! can we merge this?

Patch:
diff --git a/converter.go b/converter.go
index d4b3a83..7d7b20e 100644
--- a/converter.go
+++ b/converter.go
@@ -1084,7 +1084,7 @@ func arrowToRecord(ctx context.Context, record arrow.Record, pool memory.Allocat
 							ar = arrow.Timestamp(ts.UnixNano())
 							// in case of overflow in arrow timestamp return error
 							// this could only happen for nanosecond case
-							if ts.Year() != ar.ToTime(arrow.Nanosecond).Year() {
+							if ts.UTC().Year() != ar.ToTime(arrow.Nanosecond).Year() {
 								return nil, &SnowflakeError{
 									Number:   ErrTooHighTimestampPrecision,
 									SQLState: SQLStateInvalidDataTimeFormat,


Test patch:
diff --git a/converter_test.go b/converter_test.go
index d596bd0..96d33b7 100644
--- a/converter_test.go
+++ b/converter_test.go
@@ -892,7 +892,7 @@ func TestArrowToRecord(t *testing.T) {
 	defer pool.AssertSize(t, 0) // ensure no arrow memory leaks
 	var valids []bool           // AppendValues() with an empty valid array adds every value by default
 
-	localTime := time.Date(2019, 2, 6, 14, 17, 31, 123456789, time.FixedZone("-08:00", -8*3600))
+	localTime := time.Date(2019, 1, 1, 1, 17, 31, 123456789, time.FixedZone("-08:00", -8*3600))
 	localTimeFarIntoFuture := time.Date(9000, 2, 6, 14, 17, 31, 123456789, time.FixedZone("-08:00", -8*3600))
 
 	epochField := arrow.Field{Name: "epoch", Type: &arrow.Int64Type{}}
Repository: snowflakedb/gosnowflake
Pull request #892 (2023-08-23 10:16:55)

SNOW-726742 Remove queryID from snowflakeConn
### Description
QueryID was persisted in snowflakeConn and read by some functions. It could cause race conditions if the same connection is reused among threads. Now, the queryId is read from the responses.

### Checklist
- [x] Code compiles correctly
- [x] Run ``make fmt`` to fix inconsistent formats
- [x] Run ``make lint`` to get lint errors and fix all of them
- [x] Created tests which fail without the change (if possible)
- [x] All tests passing
- [x] Extended the README / documentation, if necessary



Patch:
diff --git a/connection.go b/connection.go
index aa26e85..cf129a1 100644
--- a/connection.go
+++ b/connection.go
@@ -64,8 +64,6 @@ type snowflakeConn struct {
 	cfg             *Config
 	rest            *snowflakeRestful
 	SequenceCounter uint64
-	QueryID         string
-	SQLState        string
 	telemetry       *snowflakeTelemetry
 	internal        InternalClient
 }
@@ -158,8 +156,6 @@ func (sc *snowflakeConn) exec(
 	sc.cfg.Schema = data.Data.FinalSchemaName
 	sc.cfg.Role = data.Data.FinalRoleName
 	sc.cfg.Warehouse = data.Data.FinalWarehouseName
-	sc.QueryID = data.Data.QueryID
-	sc.SQLState = data.Data.SQLState
 	sc.populateSessionParameters(data.Data.Parameters)
 	return data, err
 }
@@ -282,7 +278,7 @@ func (sc *snowflakeConn) ExecContext(
 		return &snowflakeResult{
 			affectedRows: updatedRows,
 			insertID:     -1,
-			queryID:      sc.QueryID,
+			queryID:      data.Data.QueryID,
 		}, nil // last insert id is not supported by Snowflake
 	} else if isMultiStmt(&data.Data) {
 		return sc.handleMultiExec(ctx, data.Data)
@@ -353,7 +349,7 @@ func (sc *snowflakeConn) queryContextInternal(
 
 	rows := new(snowflakeRows)
 	rows.sc = sc
-	rows.queryID = sc.QueryID
+	rows.queryID = data.Data.QueryID
 
 	if isMultiStmt(&data.Data) {
 		// handleMultiQuery is responsible to fill rows with childResults
diff --git a/multistatement.go b/multistatement.go
index c8b13e2..ce9d991 100644
--- a/multistatement.go
+++ b/multistatement.go
@@ -78,7 +78,7 @@ func (sc *snowflakeConn) handleMultiExec(
 	return &snowflakeResult{
 		affectedRows: updatedRows,
 		insertID:     -1,
-		queryID:      sc.QueryID,
+		queryID:      data.QueryID,
 	}, nil
 }
 


Test patch:
diff --git a/rows_test.go b/rows_test.go
index bf414f8..0d8fbc1 100644
--- a/rows_test.go
+++ b/rows_test.go
@@ -460,6 +460,7 @@ func TestWithArrowBatchesNotImplementedForResult(t *testing.T) {
 	if err != nil {
 		t.Error(err)
 	}
+	defer sc.Close()
 	if err = authenticateWithConfig(sc); err != nil {
 		t.Error(err)
 	}
@@ -474,12 +475,6 @@ func TestWithArrowBatchesNotImplementedForResult(t *testing.T) {
 		t.Error(err)
 	}
 
-	result.(*snowflakeResult).GetStatus()
-	queryID := result.(*snowflakeResult).GetQueryID()
-	if queryID != sc.QueryID {
-		t.Fatalf("failed to get query ID")
-	}
-
 	_, err = result.(*snowflakeResult).GetArrowBatches()
 	if err == nil {
 		t.Fatal("should have raised an error")
Repository: snowflakedb/gosnowflake
Pull request #1247 (2024-11-18 15:19:14)

SNOW-1806123 Propagate context to snowflakeConn
### Description

SNOW-1806123 It is a bit of reversion of https://github.com/snowflakedb/gosnowflake/pull/1196 . Firstly, the issue was that we persisted a auth context in `snowflakeConn` and reused it in closing. But if the context was cancelled or timed out, it ended up in problems with closing connection. After 1196 context was not propagated to auth calls so auth calls couldn't have been cancelled or timed out. Now the context is propagated to auth calls again, but it is not used during session closing.

### Checklist
- [x] Created tests which fail without the change (if possible)
- [x] Extended the README / documentation, if necessary



Patch:
diff --git a/connection.go b/connection.go
index 6dd4ae0..26fa8f0 100644
--- a/connection.go
+++ b/connection.go
@@ -284,7 +284,8 @@ func (sc *snowflakeConn) Close() (err error) {
 	defer sc.cleanup()
 
 	if sc.cfg != nil && !sc.cfg.KeepSessionAlive {
-		if err = sc.rest.FuncCloseSession(sc.ctx, sc.rest, sc.rest.RequestTimeout); err != nil {
+		// we have to replace context with background, otherwise we can use a one that is cancelled or timed out
+		if err = sc.rest.FuncCloseSession(context.Background(), sc.rest, sc.rest.RequestTimeout); err != nil {
 			logger.WithContext(sc.ctx).Error(err)
 		}
 	}
@@ -775,7 +776,7 @@ func (scd *snowflakeArrowStreamChunkDownloader) GetBatches() (out []ArrowStreamB
 func buildSnowflakeConn(ctx context.Context, config Config) (*snowflakeConn, error) {
 	sc := &snowflakeConn{
 		SequenceCounter:     0,
-		ctx:                 context.Background(),
+		ctx:                 ctx,
 		cfg:                 &config,
 		queryContextCache:   (&queryContextCache{}).init(),
 		currentTimeProvider: defaultTimeProvider,


Test patch:
diff --git a/auth_test.go b/auth_test.go
index e05a96d..ccc0f2c 100644
--- a/auth_test.go
+++ b/auth_test.go
@@ -977,3 +977,30 @@ func TestOktaRetryWithNewToken(t *testing.T) {
 	assertEqualF(t, authResponse.MfaToken, expectedMfaToken)
 	assertEqualF(t, authResponse.SessionInfo.DatabaseName, expectedDatabaseName)
 }
+
+func TestContextPropagatedToAuthWhenUsingOpen(t *testing.T) {
+	t.Skip("dsnConnector loses context when opening using sql.Open")
+	db, err := sql.Open("snowflake", dsn)
+	assertNilF(t, err)
+	defer db.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	_, err = db.QueryContext(ctx, "SELECT 1")
+	assertNotNilF(t, err)
+	assertStringContainsE(t, err.Error(), "context deadline exceeded")
+	cancel()
+}
+
+func TestContextPropagatedToAuthWhenUsingOpenDB(t *testing.T) {
+	cfg, err := ParseDSN(dsn)
+	assertNilF(t, err)
+	connector := NewConnector(&SnowflakeDriver{}, *cfg)
+	db := sql.OpenDB(connector)
+	defer db.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	_, err = db.QueryContext(ctx, "SELECT 1")
+	assertNotNilF(t, err)
+	assertStringContainsE(t, err.Error(), "context deadline exceeded")
+	cancel()
+}