SNOW-895536: Limit query context cache

connection.go

@@ -141,7 +141,7 @@ func (sc *snowflakeConn) exec(
 		return nil, err
 	}
 
-	sc.queryContextCache.add(data.Data.QueryContext.Entries...)
+	sc.queryContextCache.add(sc, data.Data.QueryContext.Entries...)
 
 	// handle PUT/GET commands
 	if isFileTransfer(query) {

htap.go

@@ -1,6 +1,14 @@
 package gosnowflake
 
-import "sync"
+import (
+	"strconv"
+	"sync"
+)
+
+const (
+	queryContextCacheSizeParamName = "QUERY_CONTEXT_CACHE_SIZE"
+	defaultQueryContextCacheSize   = 5
+)
 
 type queryContextEntry struct {
 	ID        int   `json:"id"`
@@ -19,8 +27,27 @@ func (qcc *queryContextCache) init() *queryContextCache {
 	return qcc
 }
 
-func (qcc *queryContextCache) add(qces ...queryContextEntry) {
+func (qcc *queryContextCache) add(sc *snowflakeConn, qces ...queryContextEntry) {
 	qcc.mutex.Lock()
 	defer qcc.mutex.Unlock()
 	qcc.entries = append(qcc.entries, qces...)
+	qcc.prune(qcc.getQueryContextCacheSize(sc))
+}
+
+func (qcc *queryContextCache) prune(size int) {
+	if len(qcc.entries) > size {
+		qcc.entries = qcc.entries[0:size]
+	}
+}
+
+func (qcc *queryContextCache) getQueryContextCacheSize(sc *snowflakeConn) int {
+	if sizeStr, ok := sc.cfg.Params[queryContextCacheSizeParamName]; ok {
+		size, err := strconv.Atoi(*sizeStr)
+		if err != nil {
+			logger.Warnf("cannot parse %v as int as query context cache size: %v", sizeStr, err)
+		} else {
+			return size
+		}
+	}
+	return defaultQueryContextCacheSize
 }

htap_test.go

@@ -2,6 +2,7 @@ package gosnowflake
 
 import (
 	"encoding/json"
+	"fmt"
 	"reflect"
 	"strings"
 	"testing"
@@ -147,3 +148,56 @@ func containsNewEntries(entriesAfter []queryContextEntry, entriesBefore []queryC
 
 	return false
 }
+
+func TestPrune(t *testing.T) {
+	size := "3"
+	qce1 := queryContextEntry{1, 1, 1, nil}
+	qce2 := queryContextEntry{2, 2, 2, nil}
+	qce3 := queryContextEntry{3, 3, 3, nil}
+
+	testcases := []struct {
+		size     int
+		expected []queryContextEntry
+	}{
+		{
+			size:     1,
+			expected: []queryContextEntry{qce1},
+		},
+		{
+			size:     2,
+			expected: []queryContextEntry{qce1, qce2},
+		},
+		{
+			size:     3,
+			expected: []queryContextEntry{qce1, qce2, qce3},
+		},
+		{
+			size:     4,
+			expected: []queryContextEntry{qce1, qce2, qce3},
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(fmt.Sprintf("%v", tc.size), func(t *testing.T) {
+			sc := &snowflakeConn{
+				cfg: &Config{
+					Params: map[string]*string{
+						queryContextCacheSizeParamName: &size,
+					},
+				},
+			}
+
+			qcc := (&queryContextCache{}).init()
+
+			qcc.add(sc, qce1)
+			qcc.add(sc, qce2)
+			qcc.add(sc, qce3)
+
+			qcc.prune(tc.size)
+
+			if !reflect.DeepEqual(qcc.entries, tc.expected) {
+				t.Errorf("unexpected cache entries. expected: %v, got: %v", tc.expected, qcc.entries)
+			}
+		})
+	}
+}