Repository: snowflakedb/gosnowflake
Pull request #945 (2023-10-25 06:55:08)

SNOW-949229: support overwrite option in PUT
### Description
Supports overwrite in PUT commands. Before, overwrite retrieved from Snowflake was not taken into account.
https://github.com/snowflakedb/snowflake-sdks-drivers-issues-teamwork/issues/712

### 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


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/945?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
> Merging [#945](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/945?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) (7e7d9b4) into [master](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/99921d481830252887ec89396a6ebcddb5856ec4?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) (99921d4) will **decrease** coverage by `1.09%`.
> The diff coverage is `66.66%`.

```diff
@@            Coverage Diff             @@
##           master     #945      +/-   ##
==========================================
- Coverage   83.24%   82.15%   -1.09%     
==========================================
  Files          49       48       -1     
  Lines        9340     9342       +2     
==========================================
- Hits         7775     7675     -100     
- Misses       1565     1667     +102     
```


| [Files](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/945?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) | Coverage Δ | |
|---|---|---|
| [file\_transfer\_agent.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/945?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-ZmlsZV90cmFuc2Zlcl9hZ2VudC5nbw==) | `73.24% <100.00%> (-1.87%)` | :arrow_down: |
| [storage\_client.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/945?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-c3RvcmFnZV9jbGllbnQuZ28=) | `77.32% <60.00%> (+1.73%)` | :arrow_up: |

... and [6 files with indirect coverage changes](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/945/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb)


Patch:
diff --git a/doc.go b/doc.go
index e234263..df4a610 100644
--- a/doc.go
+++ b/doc.go
@@ -960,5 +960,10 @@ Remember, to encode slashes.
 Example:
 
 	u:p@a.r.c.snowflakecomputing.com/db/s?account=a.r.c&tmpDirPath=%2Fother%2Ftmp
+
+## Using custom configuration for PUT/GET
+
+If you want to override some default configuration options, you can use `WithFileTransferOptions` context.
+There are multiple config parameters including progress bars or compression.
 */
 package gosnowflake
diff --git a/file_transfer_agent.go b/file_transfer_agent.go
index 9e3a646..ae63f9a 100644
--- a/file_transfer_agent.go
+++ b/file_transfer_agent.go
@@ -92,7 +92,6 @@ type SnowflakeFileTransferOptions struct {
 	compressSourceFromStream bool
 
 	/* PUT */
-	DisablePutOverwrite     bool
 	putCallback             *snowflakeProgressPercentage
 	putAzureCallback        *snowflakeProgressPercentage
 	putCallbackOutputStream *io.Writer
@@ -268,7 +267,7 @@ func (sfa *snowflakeFileTransferAgent) parseCommand() error {
 	if sfa.data.Parallel != 0 {
 		sfa.parallel = sfa.data.Parallel
 	}
-	sfa.overwrite = !sfa.options.DisablePutOverwrite
+	sfa.overwrite = sfa.data.Overwrite
 	sfa.stageLocationType = cloudType(strings.ToUpper(sfa.data.StageInfo.LocationType))
 	sfa.stageInfo = &sfa.data.StageInfo
 	sfa.presignedURLs = make([]string, 0)
diff --git a/gcs_storage_client.go b/gcs_storage_client.go
index 5d9dac9..9affaa7 100644
--- a/gcs_storage_client.go
+++ b/gcs_storage_client.go
@@ -34,7 +34,7 @@ func (util *snowflakeGcsClient) createClient(info *execResponseStageInfo, _ bool
 		logger.Debug("Using GCS downscoped token")
 		return info.Creds.GcsAccessToken, nil
 	}
-	logger.Debug("No access token received from GS, using presigned url")
+	logger.Debugf("No access token received from GS, using presigned url: %s", info.PresignedURL)
 	return "", nil
 }
 
diff --git a/query.go b/query.go
index 300233d..162b452 100644
--- a/query.go
+++ b/query.go
@@ -127,6 +127,7 @@ type execResponseData struct {
 	Parallel                int64                 `json:"parallel,omitempty"`
 	Threshold               int64                 `json:"threshold,omitempty"`
 	AutoCompress            bool                  `json:"autoCompress,omitempty"`
+	Overwrite               bool                  `json:"overwrite,omitempty"`
 	SourceCompression       string                `json:"sourceCompression,omitempty"`
 	ShowEncryptionParameter bool                  `json:"clientShowEncryptionParameter,omitempty"`
 	EncryptionMaterial      encryptionWrapper     `json:"encryptionMaterial,omitempty"`
diff --git a/s3_storage_client.go b/s3_storage_client.go
index ed3bca5..82638e7 100644
--- a/s3_storage_client.go
+++ b/s3_storage_client.go
@@ -80,11 +80,7 @@ func (util *snowflakeS3Client) getFileHeader(meta *fileMetadata, filename string
 		if errors.As(err, &ae) {
 			if ae.ErrorCode() == notFound {
 				meta.resStatus = notFoundFile
-				return &fileHeader{
-					digest:             "",
-					contentLength:      0,
-					encryptionMetadata: nil,
-				}, nil
+				return nil, fmt.Errorf("could not find file")
 			} else if ae.ErrorCode() == expiredToken {
 				meta.resStatus = renewToken
 				return nil, fmt.Errorf("received expired token. renewing")
diff --git a/storage_client.go b/storage_client.go
index a7385c2..ee746a6 100644
--- a/storage_client.go
+++ b/storage_client.go
@@ -88,7 +88,12 @@ func (rsu *remoteStorageUtil) uploadOneFile(meta *fileMetadata) error {
 	for retry := 0; retry < maxRetry; retry++ {
 		if !meta.overwrite {
 			header, err := utilClass.getFileHeader(meta, meta.dstFileName)
-			if err != nil {
+			if meta.resStatus == notFoundFile {
+				err := utilClass.uploadFile(dataFile, meta, encryptMeta, maxConcurrency, meta.options.MultiPartThreshold)
+				if err != nil {
+					logger.Warnf("Error uploading %v. err: %v", dataFile, err)
+				}
+			} else if err != nil {
 				return err
 			}
 			if header != nil && meta.resStatus == uploaded {


Test patch:
diff --git a/aaa_test.go b/aaa_test.go
index de86ae1..534a3b1 100644
--- a/aaa_test.go
+++ b/aaa_test.go
@@ -1,6 +1,8 @@
 package gosnowflake
 
-import "testing"
+import (
+	"testing"
+)
 
 func TestShowServerVersion(t *testing.T) {
 	runDBTest(t, func(dbt *DBTest) {
diff --git a/assert_test.go b/assert_test.go
index f886062..185e254 100644
--- a/assert_test.go
+++ b/assert_test.go
@@ -26,6 +26,10 @@ func assertEqualF(t *testing.T, actual any, expected any, descriptions ...string
 	fatalOnNonEmpty(t, validateEqual(actual, expected, descriptions...))
 }
 
+func assertTrueF(t *testing.T, actual bool, descriptions ...string) {
+	fatalOnNonEmpty(t, validateEqual(actual, true, descriptions...))
+}
+
 func assertStringContainsE(t *testing.T, actual string, expectedToContain string, descriptions ...string) {
 	errorOnNonEmpty(t, validateStringContains(actual, expectedToContain, descriptions...))
 }
diff --git a/file_transfer_agent_test.go b/file_transfer_agent_test.go
index 163fe3c..f4ec224 100644
--- a/file_transfer_agent_test.go
+++ b/file_transfer_agent_test.go
@@ -237,9 +237,6 @@ func TestParseCommandInvalidStorageClientException(t *testing.T) {
 					EncryptionMaterials:     []snowflakeFileEncryption{mockEncMaterial1},
 				},
 			},
-			options: &SnowflakeFileTransferOptions{
-				DisablePutOverwrite: false,
-			},
 		}
 
 		err = sfa.parseCommand()
diff --git a/put_get_test.go b/put_get_test.go
index fee0516..e44c0d8 100644
--- a/put_get_test.go
+++ b/put_get_test.go
@@ -298,16 +298,13 @@ func TestPutOverwrite(t *testing.T) {
 	f.Close()
 
 	runDBTest(t, func(dbt *DBTest) {
-		if _, err = dbt.exec("use role sysadmin"); err != nil {
-			t.Skip("snowflake admin account not accessible")
-		}
 		dbt.mustExec("rm @~/test_put_overwrite")
 
 		f, _ = os.Open(testData)
 		rows := dbt.mustQueryContext(
 			WithFileStream(context.Background(), f),
 			fmt.Sprintf("put 'file://%v' @~/test_put_overwrite",
-				strings.ReplaceAll(testData, "\\", "\\\\")))
+				strings.ReplaceAll(testData, "\\", "/")))
 		defer rows.Close()
 		f.Close()
 		defer dbt.mustExec("rm @~/test_put_overwrite")
@@ -321,36 +318,60 @@ func TestPutOverwrite(t *testing.T) {
 			t.Fatalf("expected UPLOADED, got %v", s6)
 		}
 
+		rows = dbt.mustQuery("ls @~/test_put_overwrite")
+		defer rows.Close()
+		assertTrueF(t, rows.Next(), "expected new rows")
+		if err = rows.Scan(&s0, &s1, &s2, &s3); err != nil {
+			t.Fatal(err)
+		}
+		md5Column := s2
+
 		f, _ = os.Open(testData)
-		ctx := WithFileTransferOptions(context.Background(),
-			&SnowflakeFileTransferOptions{
-				DisablePutOverwrite: true,
-			})
 		rows = dbt.mustQueryContext(
-			WithFileStream(ctx, f),
+			WithFileStream(context.Background(), f),
 			fmt.Sprintf("put 'file://%v' @~/test_put_overwrite",
-				strings.ReplaceAll(testData, "\\", "\\\\")))
+				strings.ReplaceAll(testData, "\\", "/")))
 		defer rows.Close()
 		f.Close()
-		if rows.Next() {
-			if err = rows.Scan(&s0, &s1, &s2, &s3, &s4, &s5, &s6, &s7); err != nil {
-				t.Fatal(err)
-			}
+		assertTrueF(t, rows.Next(), "expected new rows")
+		if err = rows.Scan(&s0, &s1, &s2, &s3, &s4, &s5, &s6, &s7); err != nil {
+			t.Fatal(err)
 		}
-		if s6 != skipped.String() {
+		if runningOnGCP() && s6 != uploaded.String() {
+			// when this condition fails it means, that presgined URLs are replaced with downscoped tokens
+			// when it happens, all clouds should not overwrite by default, so all clouds should pass the `s6 == skipped` test
+			t.Fatalf("expected UPLOADED as long as presigned URLs are used, got %v", s6)
+		} else if !runningOnGCP() && s6 != skipped.String() {
 			t.Fatalf("expected SKIPPED, got %v", s6)
 		}
 
+		rows = dbt.mustQuery("ls @~/test_put_overwrite")
+		defer rows.Close()
+		assertTrueF(t, rows.Next(), "expected new rows")
+
+		if err = rows.Scan(&s0, &s1, &s2, &s3); err != nil {
+			t.Fatal(err)
+		}
+		if runningOnGCP() {
+			if s2 == md5Column {
+				// when this condition fails it means, that presgined URLs are replaced with downscoped tokens
+				// when it happens, all clouds should not overwrite by default, so all clouds should pass the `s2 == md5Column` check
+				t.Fatal("For GCP and presigned URLs (current on Github Actions) it should be overwritten by default")
+			}
+		} else if s2 != md5Column {
+			t.Fatal("The MD5 column should have stayed the same")
+		}
+
 		f, _ = os.Open(testData)
 		rows = dbt.mustQueryContext(
 			WithFileStream(context.Background(), f),
 			fmt.Sprintf("put 'file://%v' @~/test_put_overwrite overwrite=true",
-				strings.ReplaceAll(testData, "\\", "\\\\")))
+				strings.ReplaceAll(testData, "\\", "/")))
+		defer rows.Close()
 		f.Close()
-		if rows.Next() {
-			if err = rows.Scan(&s0, &s1, &s2, &s3, &s4, &s5, &s6, &s7); err != nil {
-				t.Fatal(err)
-			}
+		assertTrueF(t, rows.Next(), "expected new rows")
+		if err = rows.Scan(&s0, &s1, &s2, &s3, &s4, &s5, &s6, &s7); err != nil {
+			t.Fatal(err)
 		}
 		if s6 != uploaded.String() {
 			t.Fatalf("expected UPLOADED, got %v", s6)
@@ -358,14 +379,16 @@ func TestPutOverwrite(t *testing.T) {
 
 		rows = dbt.mustQuery("ls @~/test_put_overwrite")
 		defer rows.Close()
-		if rows.Next() {
-			if err = rows.Scan(&s0, &s1, &s2, &s3); err != nil {
-				t.Fatal(err)
-			}
+		assertTrueF(t, rows.Next(), "expected new rows")
+		if err = rows.Scan(&s0, &s1, &s2, &s3); err != nil {
+			t.Fatal(err)
 		}
 		if s0 != fmt.Sprintf("test_put_overwrite/%v.gz", baseName(testData)) {
 			t.Fatalf("expected test_put_overwrite/%v.gz, got %v", baseName(testData), s0)
 		}
+		if s2 == md5Column {
+			t.Fatalf("file should have been overwritten.")
+		}
 	})
 }
 
@@ -417,10 +440,9 @@ func testPutGet(t *testing.T, isStream bool) {
 		defer rows.Close()
 
 		var s0, s1, s2, s3, s4, s5, s6, s7 string
-		if rows.Next() {
-			if err = rows.Scan(&s0, &s1, &s2, &s3, &s4, &s5, &s6, &s7); err != nil {
-				t.Fatal(err)
-			}
+		assertTrueF(t, rows.Next(), "expected new rows")
+		if err = rows.Scan(&s0, &s1, &s2, &s3, &s4, &s5, &s6, &s7); err != nil {
+			t.Fatal(err)
 		}
 		if s6 != uploaded.String() {
 			t.Fatalf("expected %v, got: %v", uploaded, s6)
diff --git a/s3_storage_client_test.go b/s3_storage_client_test.go
index 8843dcd..1de79c2 100644
--- a/s3_storage_client_test.go
+++ b/s3_storage_client_test.go
@@ -296,7 +296,7 @@ func TestGetHeaderNotFoundError(t *testing.T) {
 	}
 
 	_, err := new(snowflakeS3Client).getFileHeader(&meta, "file.txt")
-	if err != nil {
+	if err != nil && err.Error() != "could not find file" {
 		t.Error(err)
 	}