Merge branch 'SNOW-1029631-optimize-put-mem' of https://github.com/sn…

…owflakedb/gosnowflake into SNOW-1029631-optimize-put-mem

.github/workflows/build-test.yml

@@ -1,22 +1,22 @@
 name: Build and Test
 
 on:
-    push:
-        braches:
-            - master
-        tags:
-            - v*
-    pull_request:
-        branches:
-            - master
-            - SNOW-*
-    schedule:
-      - cron:  '7 3 * * *'
-    workflow_dispatch:
-        inputs:
-          goTestParams:
-            default:
-            description: "Parameters passed to go test"
+  push:
+    branches:
+      - master
+    tags:
+      - v*
+  pull_request:
+    branches:
+      - master
+      - SNOW-*
+  schedule:
+    - cron: '7 3 * * *'
+  workflow_dispatch:
+    inputs:
+      goTestParams:
+        default:
+        description: 'Parameters passed to go test'
 
 concurrency:
   # older builds for the same pull request numer or branch should be cancelled
@@ -34,7 +34,7 @@ jobs:
           - name: Setup go
             uses: actions/setup-go@v5
             with:
-              go-version: '1.21'
+              go-version-file: './go.mod'
           - name: golangci-lint
             uses: golangci/golangci-lint-action@v6
             with:

.github/workflows/jira_issue.yml

@@ -37,7 +37,7 @@ jobs:
           summary: '${{ github.event.issue.title }}'
           description: |
             ${{ github.event.issue.body }} \\ \\ _Created from GitHub Action_ for ${{ github.event.issue.html_url }}
-          fields: '{ "customfield_11401": {"id": "14723"}, "assignee": {"id": "712020:3c0352b5-63f7-4e26-9afe-38f6f9f0f4c5"}, "components":[{"id":"19286"}] }'
+          fields: '{ "customfield_11401": {"id": "14723"}, "assignee": {"id": "712020:e1f41916-da57-4fe8-b317-116d5229aa51"}, "components":[{"id":"19286"}], "labels": ["oss"], "priority": {"id": "10001"} }'
 
       - name: Update GitHub Issue
         uses: ./jira/gajira-issue-update

.gitignore

@@ -1,4 +1,5 @@
 .idea/
+.vscode/
 parameters*.json
 parameters*.bat
 *.p8
@@ -11,6 +12,8 @@ wss-golang-agent.config
 wss-unified-agent.jar
 whitesource/
 *.swp
+cp.out
+__debug_bin*
 
 # exclude vendor
 vendor

Jenkinsfile

@@ -10,7 +10,7 @@ timestamps {
       env.GIT_COMMIT = scmInfo.GIT_COMMIT
     }
     params = [
-      string(name: 'svn_revision', value: 'main'),
+      string(name: 'svn_revision', value: 'bptp-built'),
       string(name: 'branch', value: 'main'),
       string(name: 'client_git_commit', value: scmInfo.GIT_COMMIT),
       string(name: 'client_git_branch', value: scmInfo.GIT_BRANCH),

README.md

@@ -56,13 +56,13 @@ This driver currently does not support GCP regional endpoints. Please ensure tha
 Snowflake provides a set of sample programs to test with. Set the environment variable ``$GOPATH`` to the top directory of your workspace, e.g., ``~/go`` and make certain to
 include ``$GOPATH/bin`` in the environment variable ``$PATH``. Run the ``make`` command to build all sample programs.
 
-```
+```sh
 make install
 ```
 
 In the following example, the program ``select1.go`` is built and installed in ``$GOPATH/bin`` and can be run from the command line:
 
-```
+```sh
 SNOWFLAKE_TEST_ACCOUNT=<your_account> \
 SNOWFLAKE_TEST_USER=<your_user> \
 SNOWFLAKE_TEST_PASSWORD=<your_password> \
@@ -79,7 +79,7 @@ The developer notes are hosted with the source code on [GitHub](https://github.c
 
 Set the Snowflake connection info in ``parameters.json``:
 
-```
+```json
 {
     "testconnection": {
         "SNOWFLAKE_TEST_USER":      "<your_user>",
@@ -88,21 +88,25 @@ Set the Snowflake connection info in ``parameters.json``:
         "SNOWFLAKE_TEST_WAREHOUSE": "<your_warehouse>",
         "SNOWFLAKE_TEST_DATABASE":  "<your_database>",
         "SNOWFLAKE_TEST_SCHEMA":    "<your_schema>",
-        "SNOWFLAKE_TEST_ROLE":      "<your_role>"
+        "SNOWFLAKE_TEST_ROLE":      "<your_role>",
+        "SNOWFLAKE_TEST_DEBUG":     "false"
     }
 }
 ```
 
 Install [jq](https://stedolan.github.io/jq) so that the parameters can get parsed correctly, and run ``make test`` in your Go development environment:
 
-```
+```sh
 make test
 ```
 
+### Setting debug mode during tests
+This is for debugging Large SQL statements (greater than 300 characters). If you want to enable debug mode, set `SNOWFLAKE_TEST_DEBUG` to `true` in `parameters.json`, or export it in your shell instance.
+
 ## customizing Logging Tags
 
 If you would like to ensure that certain tags are always present in the logs, `RegisterClientLogContextHook` can be used in your init function. See example below.
-```
+```go
 import "github.com/snowflakedb/gosnowflake"
 
 func init() {
@@ -116,7 +120,7 @@ func init() {
 
 ## Setting Log Level
 If you want to change the log level, `SetLogLevel` can be used in your init function like this:
-```
+```go
 import "github.com/snowflakedb/gosnowflake"
 
 func init() {
@@ -138,15 +142,15 @@ The following is a list of options you can pass in to set the level from least t
 
 Configure your testing environment as described above and run ``make cov``. The coverage percentage will be printed on the console when the testing completes.
 
-```
+```sh
 make cov
 ```
 
 For more detailed analysis, results are printed to ``coverage.txt`` in the project directory.
 
 To read the coverage report, run:
 
-```
+```sh
 go tool cover -html=coverage.txt
 ```
 

arrow_chunk.go

@@ -21,6 +21,7 @@ type arrowResultChunk struct {
 }
 
 func (arc *arrowResultChunk) decodeArrowChunk(ctx context.Context, rowType []execResponseRowType, highPrec bool, params map[string]*string) ([]chunkRowType, error) {
+	defer arc.reader.Release()
 	logger.Debug("Arrow Decoder")
 	var chunkRows []chunkRowType
 

arrow_test.go

@@ -6,6 +6,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"github.com/apache/arrow/go/v16/arrow/memory"
 	"math/big"
 	"reflect"
 	"strings"
@@ -577,3 +578,26 @@ func TestArrowVariousTypes(t *testing.T) {
 		}*/
 	})
 }
+
+func TestArrowMemoryCleanedUp(t *testing.T) {
+	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+	defer mem.AssertSize(t, 0)
+
+	runDBTest(t, func(dbt *DBTest) {
+		ctx := WithArrowAllocator(
+			context.Background(),
+			mem,
+		)
+
+		rows := dbt.mustQueryContext(ctx, "select 1 UNION select 2")
+		defer rows.Close()
+		var v int
+		assertTrueF(t, rows.Next())
+		assertNilF(t, rows.Scan(&v))
+		assertEqualE(t, v, 1)
+		assertTrueF(t, rows.Next())
+		assertNilF(t, rows.Scan(&v))
+		assertEqualE(t, v, 2)
+		assertFalseE(t, rows.Next())
+	})
+}

assert_test.go

@@ -68,6 +68,10 @@ func assertFalseF(t *testing.T, actual bool, descriptions ...string) {
 	fatalOnNonEmpty(t, validateEqual(actual, false, descriptions...))
 }
 
+func assertFalseE(t *testing.T, actual bool, descriptions ...string) {
+	errorOnNonEmpty(t, validateEqual(actual, false, descriptions...))
+}
+
 func assertStringContainsE(t *testing.T, actual string, expectedToContain string, descriptions ...string) {
 	errorOnNonEmpty(t, validateStringContains(actual, expectedToContain, descriptions...))
 }

async.go

@@ -119,6 +119,7 @@ func (sr *snowflakeRestful) getAsync(
 				rows.errChannel <- err
 				return err
 			}
+			rows.format = respd.Data.QueryResultFormat
 			rows.errChannel <- nil // mark query status complete
 		}
 	} else {

azure_storage_client.go

@@ -23,6 +23,7 @@ import (
 )
 
 type snowflakeAzureClient struct {
+	cfg *Config
 }
 
 type azureLocation struct {
@@ -85,9 +86,11 @@ func (util *snowflakeAzureClient) getFileHeader(meta *fileMetadata, filename str
 	if meta.mockAzureClient != nil {
 		blobClient = meta.mockAzureClient
 	}
-	resp, err := blobClient.GetProperties(context.Background(), &blob.GetPropertiesOptions{
-		AccessConditions: &blob.AccessConditions{},
-		CPKInfo:          &blob.CPKInfo{},
+	resp, err := withCloudStorageTimeout(util.cfg, func(ctx context.Context) (blob.GetPropertiesResponse, error) {
+		return blobClient.GetProperties(ctx, &blob.GetPropertiesOptions{
+			AccessConditions: &blob.AccessConditions{},
+			CPKInfo:          &blob.CPKInfo{},
+		})
 	})
 	if err != nil {
 		var se *azcore.ResponseError
@@ -203,9 +206,11 @@ func (util *snowflakeAzureClient) uploadFile(
 		if meta.realSrcStream != nil {
 			uploadSrc = meta.realSrcStream
 		}
-		_, err = blobClient.UploadStream(context.Background(), uploadSrc, &azblob.UploadStreamOptions{
-			BlockSize: int64(uploadSrc.Len()),
-			Metadata:  azureMeta,
+		_, err = withCloudStorageTimeout(util.cfg, func(ctx context.Context) (azblob.UploadStreamResponse, error) {
+			return blobClient.UploadStream(ctx, uploadSrc, &azblob.UploadStreamOptions{
+				BlockSize: int64(uploadSrc.Len()),
+				Metadata:  azureMeta,
+			})
 		})
 	} else {
 		var f *os.File
@@ -228,7 +233,9 @@ func (util *snowflakeAzureClient) uploadFile(
 		if meta.options.putAzureCallback != nil {
 			blobOptions.Progress = meta.options.putAzureCallback.call
 		}
-		_, err = blobClient.UploadFile(context.Background(), f, blobOptions)
+		_, err = withCloudStorageTimeout(util.cfg, func(ctx context.Context) (azblob.UploadFileResponse, error) {
+			return blobClient.UploadFile(ctx, f, blobOptions)
+		})
 	}
 	if err != nil {
 		var se *azcore.ResponseError
@@ -279,7 +286,9 @@ func (util *snowflakeAzureClient) nativeDownloadFile(
 		blobClient = meta.mockAzureClient
 	}
 	if meta.options.GetFileToStream {
-		blobDownloadResponse, err := blobClient.DownloadStream(context.Background(), &azblob.DownloadStreamOptions{})
+		blobDownloadResponse, err := withCloudStorageTimeout(util.cfg, func(ctx context.Context) (azblob.DownloadStreamResponse, error) {
+			return blobClient.DownloadStream(ctx, &azblob.DownloadStreamOptions{})
+		})
 		if err != nil {
 			return err
 		}
@@ -295,9 +304,11 @@ func (util *snowflakeAzureClient) nativeDownloadFile(
 			return err
 		}
 		defer f.Close()
-		_, err = blobClient.DownloadFile(
-			context.Background(), f, &azblob.DownloadFileOptions{
-				Concurrency: uint16(maxConcurrency)})
+		_, err = withCloudStorageTimeout(util.cfg, func(ctx context.Context) (any, error) {
+			return blobClient.DownloadFile(
+				ctx, f, &azblob.DownloadFileOptions{
+					Concurrency: uint16(maxConcurrency)})
+		})
 		if err != nil {
 			return err
 		}