Repository: snowflakedb/gosnowflake
Pull request #571 (2022-04-13 21:03:14)

SNOW-563221 Fix Missing Locale
### Description
Fix crash when location cannot be parsed for whatever reason. Add another safeguard when parsing timestamp_ltz as well.

### 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



Patch:
diff --git a/arrow_chunk.go b/arrow_chunk.go
index 2ef2170..cce02a0 100644
--- a/arrow_chunk.go
+++ b/arrow_chunk.go
@@ -6,6 +6,7 @@ import (
 	"bytes"
 	"encoding/base64"
 	"io"
+	"time"
 
 	"github.com/apache/arrow/go/arrow/array"
 	"github.com/apache/arrow/go/arrow/ipc"
@@ -17,6 +18,7 @@ type arrowResultChunk struct {
 	rowCount         int
 	uncompressedSize int
 	allocator        memory.Allocator
+	loc              *time.Location
 }
 
 func (arc *arrowResultChunk) decodeArrowChunk(rowType []execResponseRowType, highPrec bool) ([]chunkRowType, error) {
@@ -37,7 +39,7 @@ func (arc *arrowResultChunk) decodeArrowChunk(rowType []execResponseRowType, hig
 
 		for colIdx, col := range columns {
 			destcol := make([]snowflakeValue, numRows)
-			if err = arrowToValue(&destcol, rowType[colIdx], col, highPrec); err != nil {
+			if err = arrowToValue(&destcol, rowType[colIdx], col, arc.loc, highPrec); err != nil {
 				return nil, err
 			}
 
@@ -63,7 +65,7 @@ func (arc *arrowResultChunk) decodeArrowBatch(scd *snowflakeChunkDownloader) (*[
 		} else if err != nil {
 			return nil, err
 		}
-		record, err := arrowToRecord(rawRecord, scd.RowSet.RowType)
+		record, err := arrowToRecord(rawRecord, scd.RowSet.RowType, arc.loc)
 		rawRecord.Release()
 		if err != nil {
 			return nil, err
@@ -75,7 +77,7 @@ func (arc *arrowResultChunk) decodeArrowBatch(scd *snowflakeChunkDownloader) (*[
 }
 
 // Build arrow chunk based on RowSet of base64
-func buildFirstArrowChunk(rowsetBase64 string) arrowResultChunk {
+func buildFirstArrowChunk(rowsetBase64 string, loc *time.Location) arrowResultChunk {
 	rowSetBytes, err := base64.StdEncoding.DecodeString(rowsetBase64)
 	if err != nil {
 		return arrowResultChunk{}
@@ -85,5 +87,5 @@ func buildFirstArrowChunk(rowsetBase64 string) arrowResultChunk {
 		return arrowResultChunk{}
 	}
 
-	return arrowResultChunk{*rr, 0, 0, memory.NewGoAllocator()}
+	return arrowResultChunk{*rr, 0, 0, memory.NewGoAllocator(), loc}
 }
diff --git a/chunk_downloader.go b/chunk_downloader.go
index 10cf577..0247cce 100644
--- a/chunk_downloader.go
+++ b/chunk_downloader.go
@@ -105,7 +105,11 @@ func (scd *snowflakeChunkDownloader) start() error {
 	if scd.getQueryResultFormat() == arrowFormat && scd.RowSet.RowSetBase64 != "" {
 		// if the rowsetbase64 retrieved from the server is empty, move on to downloading chunks
 		var err error
-		firstArrowChunk := buildFirstArrowChunk(scd.RowSet.RowSetBase64)
+		var loc *time.Location
+		if scd.sc != nil {
+			loc = getCurrentLocation(scd.sc.cfg.Params)
+		}
+		firstArrowChunk := buildFirstArrowChunk(scd.RowSet.RowSetBase64, loc)
 		higherPrecision := higherPrecisionEnabled(scd.ctx)
 		scd.CurrentChunk, err = firstArrowChunk.decodeArrowChunk(scd.RowSet.RowType, higherPrecision)
 		scd.CurrentChunkSize = firstArrowChunk.rowCount
@@ -261,7 +265,11 @@ func getChunk(
 func (scd *snowflakeChunkDownloader) startArrowBatches() error {
 	var err error
 	chunkMetaLen := len(scd.ChunkMetas)
-	firstArrowChunk := buildFirstArrowChunk(scd.RowSet.RowSetBase64)
+	var loc *time.Location
+	if scd.sc != nil {
+		loc = getCurrentLocation(scd.sc.cfg.Params)
+	}
+	firstArrowChunk := buildFirstArrowChunk(scd.RowSet.RowSetBase64, loc)
 	scd.FirstBatch = &ArrowBatch{
 		idx:                0,
 		scd:                scd,
@@ -414,11 +422,16 @@ func decodeChunk(scd *snowflakeChunkDownloader, idx int, bufStream *bufio.Reader
 		if err != nil {
 			return err
 		}
+		var loc *time.Location
+		if scd.sc != nil {
+			loc = getCurrentLocation(scd.sc.cfg.Params)
+		}
 		arc := arrowResultChunk{
 			*ipcReader,
 			0,
 			int(scd.totalUncompressedSize()),
 			memory.NewGoAllocator(),
+			loc,
 		}
 		if usesArrowBatches(scd.ctx) {
 			if scd.ArrowBatches[idx].rec, err = arc.decodeArrowBatch(scd); err != nil {
diff --git a/connection_util.go b/connection_util.go
index 14b25d4..10801c8 100644
--- a/connection_util.go
+++ b/connection_util.go
@@ -152,11 +152,6 @@ func (sc *snowflakeConn) populateSessionParameters(parameters []nameValueParamet
 		logger.Debugf("parameter. name: %v, value: %v", param.Name, v)
 		sc.cfg.Params[strings.ToLower(param.Name)] = &v
 	}
-	if tz, ok := sc.cfg.Params["timezone"]; ok {
-		localLocation, _ = time.LoadLocation(*tz)
-	} else {
-		localLocation = time.Now().Location()
-	}
 }
 
 func isAsyncMode(ctx context.Context) bool {
diff --git a/converter.go b/converter.go
index c77c8b0..94716a8 100644
--- a/converter.go
+++ b/converter.go
@@ -184,9 +184,13 @@ func extractTimestamp(srcValue *string) (sec int64, nsec int64, err error) {
 	return sec, nsec, nil
 }
 
-// stringToValue converts a pointer of string data to an arbitrary golang variable. This is mainly used in fetching
-// data.
-func stringToValue(dest *driver.Value, srcColumnMeta execResponseRowType, srcValue *string) error {
+// stringToValue converts a pointer of string data to an arbitrary golang variable
+// This is mainly used in fetching data.
+func stringToValue(
+	dest *driver.Value,
+	srcColumnMeta execResponseRowType,
+	srcValue *string,
+	loc *time.Location) error {
 	if srcValue == nil {
 		logger.Debugf("snowflake data type: %v, raw value: nil", srcColumnMeta.Type)
 		*dest = nil
@@ -224,7 +228,10 @@ func stringToValue(dest *driver.Value, srcColumnMeta execResponseRowType, srcVal
 		if err != nil {
 			return err
 		}
-		*dest = time.Unix(sec, nsec).In(localLocation)
+		if loc == nil {
+			loc = time.Now().Location()
+		}
+		*dest = time.Unix(sec, nsec).In(loc)
 		return nil
 	case "timestamp_tz":
 		logger.Debugf("tz: %v", *srcValue)
@@ -295,6 +302,7 @@ func arrowToValue(
 	destcol *[]snowflakeValue,
 	srcColumnMeta execResponseRowType,
 	srcValue array.Interface,
+	loc *time.Location,
 	higherPrecision bool) error {
 	data := srcValue.Data()
 	var err error
@@ -490,7 +498,7 @@ func arrowToValue(
 			fraction := array.NewInt32Data(structData.Field(1).Data()).Int32Values()
 			for i := range *destcol {
 				if !srcValue.IsNull(i) {
-					(*destcol)[i] = time.Unix(epoch[i], int64(fraction[i])).In(localLocation)
+					(*destcol)[i] = time.Unix(epoch[i], int64(fraction[i])).In(loc)
 				}
 			}
 		} else {
@@ -498,7 +506,7 @@ func arrowToValue(
 				if !srcValue.IsNull(i) {
 					q := t / int64(math.Pow10(int(srcColumnMeta.Scale)))
 					r := t % int64(math.Pow10(int(srcColumnMeta.Scale)))
-					(*destcol)[i] = time.Unix(q, r).In(localLocation)
+					(*destcol)[i] = time.Unix(q, r).In(loc)
 				}
 			}
 		}
@@ -754,8 +762,8 @@ func higherPrecisionEnabled(ctx context.Context) bool {
 	return ok && d
 }
 
-func arrowToRecord(record array.Record, rowType []execResponseRowType) (array.Record, error) {
-	s, err := recordToSchema(record.Schema(), rowType)
+func arrowToRecord(record array.Record, rowType []execResponseRowType, loc *time.Location) (array.Record, error) {
+	s, err := recordToSchema(record.Schema(), rowType, loc)
 	if err != nil {
 		return nil, err
 	}
@@ -902,7 +910,7 @@ func arrowToRecord(record array.Record, rowType []execResponseRowType) (array.Re
 			}
 			newCol = tb.NewArray()
 		case timestampLtzType:
-			tb := array.NewTimestampBuilder(pool, &arrow.TimestampType{TimeZone: "UTC"})
+			tb := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: loc.String()})
 			if col.DataType().ID() == arrow.STRUCT {
 				structData := array.NewStructData(data)
 				epoch := array.NewInt64Data(structData.Field(0).Data()).Int64Values()
@@ -966,7 +974,7 @@ func arrowToRecord(record array.Record, rowType []execResponseRowType) (array.Re
 	return array.NewRecord(s, cols, numRows), nil
 }
 
-func recordToSchema(sc *arrow.Schema, rowType []execResponseRowType) (*arrow.Schema, error) {
+func recordToSchema(sc *arrow.Schema, rowType []execResponseRowType, loc *time.Location) (*arrow.Schema, error) {
 	var fields []arrow.Field
 	for i := 0; i < len(sc.Fields()); i++ {
 		f := sc.Field(i)
@@ -995,7 +1003,7 @@ func recordToSchema(sc *arrow.Schema, rowType []execResponseRowType) (*arrow.Sch
 		case timestampNtzType, timestampTzType:
 			t = &arrow.TimestampType{}
 		case timestampLtzType:
-			t = &arrow.TimestampType{TimeZone: localLocation.String()}
+			t = &arrow.TimestampType{TimeZone: loc.String()}
 		default:
 			converted = false
 		}
diff --git a/location.go b/location.go
index a1116f1..c96528f 100644
--- a/location.go
+++ b/location.go
@@ -10,7 +10,6 @@ import (
 )
 
 var (
-	localLocation       *time.Location
 	timezones           map[int]*time.Location
 	updateTimezoneMutex *sync.Mutex
 )
@@ -90,3 +89,16 @@ func init() {
 		timezones[i] = genTimezone(i)
 	}
 }
+
+// retrieve current location based on connection
+func getCurrentLocation(params map[string]*string) *time.Location {
+	loc := time.Now().Location()
+	var err error
+	if tz, ok := params["timezone"]; ok {
+		loc, err = time.LoadLocation(*tz)
+		if err != nil {
+			loc = time.Now().Location()
+		}
+	}
+	return loc
+}
diff --git a/rows.go b/rows.go
index 0663c97..98713f7 100644
--- a/rows.go
+++ b/rows.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"reflect"
 	"strings"
+	"time"
 )
 
 const (
@@ -170,7 +171,11 @@ func (rows *snowflakeRows) Next(dest []driver.Value) (err error) {
 		for i, n := 0, len(row.RowSet); i < n; i++ {
 			// could move to chunk downloader so that each go routine
 			// can convert data
-			err = stringToValue(&dest[i], rows.ChunkDownloader.getRowType()[i], row.RowSet[i])
+			var loc *time.Location
+			if rows.sc != nil {
+				loc = getCurrentLocation(rows.sc.cfg.Params)
+			}
+			err = stringToValue(&dest[i], rows.ChunkDownloader.getRowType()[i], row.RowSet[i], loc)
 			if err != nil {
 				return err
 			}


Test patch:
diff --git a/converter_test.go b/converter_test.go
index f481018..728bfd2 100644
--- a/converter_test.go
+++ b/converter_test.go
@@ -3,6 +3,7 @@
 package gosnowflake
 
 import (
+	"context"
 	"database/sql/driver"
 	"fmt"
 	"math/big"
@@ -177,8 +178,7 @@ func TestStringToValue(t *testing.T) {
 		rowType = &execResponseRowType{
 			Type: tt,
 		}
-		err = stringToValue(&dest, *rowType, &source)
-		if err == nil {
+		if err = stringToValue(&dest, *rowType, &source, nil); err == nil {
 			t.Errorf("should raise error. type: %v, value:%v", tt, source)
 		}
 	}
@@ -197,15 +197,14 @@ func TestStringToValue(t *testing.T) {
 			rowType = &execResponseRowType{
 				Type: tt,
 			}
-			err = stringToValue(&dest, *rowType, &ss)
-			if err == nil {
+			if err = stringToValue(&dest, *rowType, &ss, nil); err == nil {
 				t.Errorf("should raise error. type: %v, value:%v", tt, source)
 			}
 		}
 	}
 
 	src := "1549491451.123456789"
-	if err = stringToValue(&dest, execResponseRowType{Type: "timestamp_ltz"}, &src); err != nil {
+	if err = stringToValue(&dest, execResponseRowType{Type: "timestamp_ltz"}, &src, nil); err != nil {
 		t.Errorf("unexpected error: %v", err)
 	} else if ts, ok := dest.(time.Time); !ok {
 		t.Errorf("expected type: 'time.Time', got '%v'", reflect.TypeOf(dest))
@@ -561,8 +560,7 @@ func TestArrowToValue(t *testing.T) {
 			meta := tc.rowType
 			meta.Type = tc.logical
 
-			err := arrowToValue(&dest, meta, arr, true)
-			if err != nil {
+			if err := arrowToValue(&dest, meta, arr, localTime.Location(), true); err != nil {
 				t.Fatalf("error: %s", err)
 			}
 
@@ -972,7 +970,7 @@ func TestArrowToRecord(t *testing.T) {
 			meta := tc.rowType
 			meta.Type = tc.logical
 
-			transformedRec, err := arrowToRecord(rawRec, []execResponseRowType{meta})
+			transformedRec, err := arrowToRecord(rawRec, []execResponseRowType{meta}, localTime.Location())
 			if err != nil {
 				t.Fatalf("error: %s", err)
 			}
@@ -991,6 +989,45 @@ func TestArrowToRecord(t *testing.T) {
 				}
 			}
 		})
+	}
+}
 
+func TestTimestampLTZLocation(t *testing.T) {
+	config, err := ParseDSN(dsn)
+	if err != nil {
+		t.Error(err)
+	}
+	ctx := context.Background()
+	sc, err := buildSnowflakeConn(ctx, *config)
+	if err != nil {
+		t.Error(err)
+	}
+	if err = authenticateWithConfig(sc); err != nil {
+		t.Error(err)
+	}
+
+	src := "1549491451.123456789"
+	var dest driver.Value
+	loc, _ := time.LoadLocation(PSTLocation)
+	if err = stringToValue(&dest, execResponseRowType{Type: "timestamp_ltz"}, &src, loc); err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+	ts, ok := dest.(time.Time)
+	if !ok {
+		t.Errorf("expected type: 'time.Time', got '%v'", reflect.TypeOf(dest))
+	}
+	if ts.Location() != loc {
+		t.Errorf("expected location to be %v, got '%v'", loc, ts.Location())
+	}
+
+	if err = stringToValue(&dest, execResponseRowType{Type: "timestamp_ltz"}, &src, nil); err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+	ts, ok = dest.(time.Time)
+	if !ok {
+		t.Errorf("expected type: 'time.Time', got '%v'", reflect.TypeOf(dest))
+	}
+	if ts.Location() != time.Local {
+		t.Errorf("expected location to be local, got '%v'", ts.Location())
 	}
 }