From b8c6a93404527e64833fe0cf30dd320a9aefb86a Mon Sep 17 00:00:00 2001
From: Joyce Ling <115662568+sfc-gh-ext-simba-jl@users.noreply.github.com>
Date: Thu, 7 Sep 2023 14:40:36 -0700
Subject: [PATCH] add mutext lock on params map

---
 chunk_downloader.go | 6 ++++++
 connection.go       | 4 ++++
 location.go         | 2 --
 rows.go             | 2 ++
 4 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/chunk_downloader.go b/chunk_downloader.go
index f715992c0..45f5440c5 100644
--- a/chunk_downloader.go
+++ b/chunk_downloader.go
@@ -106,9 +106,11 @@ func (scd *snowflakeChunkDownloader) start() error {
 		// if the rowsetbase64 retrieved from the server is empty, move on to downloading chunks
 		var err error
 		var loc *time.Location
+		paramsMutex.Lock()
 		if scd.sc != nil && scd.sc.cfg != nil {
 			loc = getCurrentLocation(scd.sc.cfg.Params)
 		}
+		paramsMutex.Unlock()
 		firstArrowChunk := buildFirstArrowChunk(scd.RowSet.RowSetBase64, loc, scd.pool)
 		higherPrecision := higherPrecisionEnabled(scd.ctx)
 		scd.CurrentChunk, err = firstArrowChunk.decodeArrowChunk(scd.RowSet.RowType, higherPrecision)
@@ -271,9 +273,11 @@ func (scd *snowflakeChunkDownloader) startArrowBatches() error {
 	var err error
 	chunkMetaLen := len(scd.ChunkMetas)
 	var loc *time.Location
+	paramsMutex.Lock()
 	if scd.sc != nil && scd.sc.cfg != nil {
 		loc = getCurrentLocation(scd.sc.cfg.Params)
 	}
+	paramsMutex.Unlock()
 	firstArrowChunk := buildFirstArrowChunk(scd.RowSet.RowSetBase64, loc, scd.pool)
 	scd.FirstBatch = &ArrowBatch{
 		idx:                0,
@@ -430,9 +434,11 @@ func decodeChunk(scd *snowflakeChunkDownloader, idx int, bufStream *bufio.Reader
 			return err
 		}
 		var loc *time.Location
+		paramsMutex.Lock()
 		if scd.sc != nil && scd.sc.cfg != nil {
 			loc = getCurrentLocation(scd.sc.cfg.Params)
 		}
+		paramsMutex.Unlock()
 		arc := arrowResultChunk{
 			ipcReader,
 			0,
diff --git a/connection.go b/connection.go
index a48af3600..69611300f 100644
--- a/connection.go
+++ b/connection.go
@@ -200,7 +200,9 @@ func (sc *snowflakeConn) cleanup() {
 		sc.rest.Client.CloseIdleConnections()
 	}
 	sc.rest = nil
+	paramsMutex.Lock()
 	sc.cfg = nil
+	paramsMutex.Unlock()
 }
 
 func (sc *snowflakeConn) Close() (err error) {
@@ -599,9 +601,11 @@ type snowflakeArrowStreamChunkDownloader struct {
 }
 
 func (scd *snowflakeArrowStreamChunkDownloader) Location() *time.Location {
+	paramsMutex.Lock()
 	if scd.sc != nil {
 		return getCurrentLocation(scd.sc.cfg.Params)
 	}
+	paramsMutex.Unlock()
 	return nil
 }
 func (scd *snowflakeArrowStreamChunkDownloader) TotalRows() int64 { return scd.Total }
diff --git a/location.go b/location.go
index 6209f3181..16fdbb163 100644
--- a/location.go
+++ b/location.go
@@ -94,13 +94,11 @@ func init() {
 func getCurrentLocation(params map[string]*string) *time.Location {
 	loc := time.Now().Location()
 	var err error
-	paramsMutex.Lock()
 	if tz, ok := params["timezone"]; ok && tz != nil {
 		loc, err = time.LoadLocation(*tz)
 		if err != nil {
 			loc = time.Now().Location()
 		}
 	}
-	paramsMutex.Unlock()
 	return loc
 }
diff --git a/rows.go b/rows.go
index 83f49ba94..f63d2f747 100644
--- a/rows.go
+++ b/rows.go
@@ -47,9 +47,11 @@ type snowflakeRows struct {
 }
 
 func (rows *snowflakeRows) getLocation() *time.Location {
+	paramsMutex.Lock()
 	if rows.location == nil && rows.sc != nil && rows.sc.cfg != nil {
 		rows.location = getCurrentLocation(rows.sc.cfg.Params)
 	}
+	paramsMutex.Unlock()
 	return rows.location
 }