support overwrite option in PUT

snowflakedb · Oct 25, 2023 · f5a0ec6 · f5a0ec6
1 parent 8ed5c10
commit f5a0ec6
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 16 deletions.
diff --git 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/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
@@ -321,13 +321,18 @@ func TestPutOverwrite(t *testing.T) {
 			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)
+			}
+		}
+		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()
@@ -341,11 +346,23 @@ func TestPutOverwrite(t *testing.T) {
 			t.Fatalf("expected SKIPPED, 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)
+			}
+		}
+		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 {
@@ -366,6 +383,9 @@ func TestPutOverwrite(t *testing.T) {
 		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.")
+		}
 	})
 }
 

diff --git 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
@@ -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
@@ -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.Debugf("Error uploading %v. err: %v", dataFile, err)
+				}
+			} else if err != nil {
 				return err
 			}
 			if header != nil && meta.resStatus == uploaded {