support overwrite option in PUT

assert_test.go

@@ -25,6 +25,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...))
 }

doc.go

@@ -960,5 +960,10 @@ Remember, to encode slashes.
 Example:
 
 	u:[email protected]/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

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)

file_transfer_agent_test.go

@@ -237,9 +237,6 @@ func TestParseCommandInvalidStorageClientException(t *testing.T) {
 					EncryptionMaterials:     []snowflakeFileEncryption{mockEncMaterial1},
 				},
 			},
-			options: &SnowflakeFileTransferOptions{
-				DisablePutOverwrite: false,
-			},
 		}
 
 		err = sfa.parseCommand()

put_get_test.go

@@ -288,6 +288,9 @@ func TestPutWithAutoCompressFalse(t *testing.T) {
 }
 
 func TestPutOverwrite(t *testing.T) {
+	if runningOnGCP() {
+		t.Skip("Overwriting is default as long as presigned URLs are enabled")
+	}
 	tmpDir := t.TempDir()
 	testData := filepath.Join(tmpDir, "data.txt")
 	f, err := os.Create(testData)
@@ -298,9 +301,6 @@ 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)
@@ -321,51 +321,67 @@ 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)
+		}
+		uploadTime := s3
+
 		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, "\\", "\\\\")))
 		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() {
 			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 s3 != uploadTime {
+			t.Fatalf("upload time should have stayed the same, expected: %v, got: %v", uploadTime, s3)
+		}
+
 		f, _ = os.Open(testData)
 		rows = dbt.mustQueryContext(
 			WithFileStream(context.Background(), f),
 			fmt.Sprintf("put 'file://%v' @~/test_put_overwrite overwrite=true",
 				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)
 		}
 
 		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 s3 == uploadTime {
+			t.Fatalf("file should have been overwritten.")
+		}
 	})
 }
 
@@ -417,10 +433,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)

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"`

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, errors.New("could not find file")
 			} else if ae.ErrorCode() == expiredToken {
 				meta.resStatus = renewToken
 				return nil, fmt.Errorf("received expired token. renewing")

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)
 	}
 

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 {