Merge pull request duckdb#128 from carlopi/one_request_upload

Single PUT when uploading single buffer-files to S3 (version 2)

Author: carlopi

extension/httpfs/include/s3fs.hpp

@@ -127,6 +127,7 @@ class S3FileHandle : public HTTPFileHandle {
 
 	S3AuthParams auth_params;
 	const S3ConfigParams config_params;
+	bool initialized_multipart_upload {false};
 
 public:
 	void Close() override;
@@ -214,6 +215,8 @@ class S3FileSystem : public HTTPFileSystem {
 	// Uploads the contents of write_buffer to S3.
 	// Note: caller is responsible to not call this method twice on the same buffer
 	static void UploadBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer);
+	static void UploadSingleBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer);
+	static void UploadBufferImplementation(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer, string query_param, bool direct_throw);
 
 	vector<OpenFileInfo> Glob(const string &glob_pattern, FileOpener *opener = nullptr) override;
 	bool ListFiles(const string &directory, const std::function<void(const string &, bool)> &callback,

extension/httpfs/s3fs.cpp

@@ -336,6 +336,8 @@ string S3FileSystem::InitializeMultipartUpload(S3FileHandle &file_handle) {
 
 	open_tag_pos += 10; // Skip open tag
 
+	file_handle.initialized_multipart_upload = true;
+
 	return result.substr(open_tag_pos, close_tag_pos - open_tag_pos);
 }
 
@@ -353,10 +355,21 @@ void S3FileSystem::NotifyUploadsInProgress(S3FileHandle &file_handle) {
 }
 
 void S3FileSystem::UploadBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer) {
-	auto &s3fs = (S3FileSystem &)file_handle.file_system;
-
 	string query_param = "partNumber=" + to_string(write_buffer->part_no + 1) + "&" +
 	                     "uploadId=" + S3FileSystem::UrlEncode(file_handle.multipart_upload_id, true);
+
+	UploadBufferImplementation(file_handle, write_buffer, query_param, false);
+
+	NotifyUploadsInProgress(file_handle);
+}
+
+void S3FileSystem::UploadSingleBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer) {
+	UploadBufferImplementation(file_handle, write_buffer, "", true);
+}
+
+void S3FileSystem::UploadBufferImplementation(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer, string query_param, bool single_upload) {
+	auto &s3fs = (S3FileSystem &)file_handle.file_system;
+
 	unique_ptr<HTTPResponse> res;
 	string etag;
 
@@ -374,6 +387,9 @@ void S3FileSystem::UploadBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuf
 		}
 		etag = res->headers.GetHeaderValue("ETag");
 	} catch (std::exception &ex) {
+		if (single_upload) {
+			throw;
+		}
 		ErrorData error(ex);
 		if (error.Type() != ExceptionType::IO && error.Type() != ExceptionType::HTTP) {
 			throw;
@@ -385,6 +401,7 @@ void S3FileSystem::UploadBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuf
 			file_handle.upload_exception = std::current_exception();
 		}
 
+		D_ASSERT(!single_upload); // If we are here we are in the multi-buffer situation
 		NotifyUploadsInProgress(file_handle);
 		return;
 	}
@@ -399,8 +416,6 @@ void S3FileSystem::UploadBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuf
 
 	// Free up space for another thread to acquire an S3WriteBuffer
 	write_buffer.reset();
-
-	NotifyUploadsInProgress(file_handle);
 }
 
 void S3FileSystem::FlushBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer) {
@@ -437,6 +452,9 @@ void S3FileSystem::FlushBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuff
 #endif
 		file_handle.uploads_in_progress++;
 	}
+	if (file_handle.initialized_multipart_upload == false) {
+		file_handle.multipart_upload_id = InitializeMultipartUpload(file_handle);
+	}
 
 #ifdef SAME_THREAD_UPLOAD
 	UploadBuffer(file_handle, write_buffer);
@@ -459,6 +477,16 @@ void S3FileSystem::FlushAllBuffers(S3FileHandle &file_handle) {
 	}
 	file_handle.write_buffers_lock.unlock();
 
+	if (file_handle.initialized_multipart_upload == false) {
+		// TODO (carlo): unclear how to handle kms_key_id, but given currently they are custom, leave the multiupload codepath in that case
+		if (to_flush.size() == 1 && file_handle.auth_params.kms_key_id.empty()) {
+			UploadSingleBuffer(file_handle, to_flush[0]);
+			file_handle.upload_finalized= true;
+			return;
+		} else {
+			file_handle.multipart_upload_id = InitializeMultipartUpload(file_handle);
+		}
+	}
 	// Flush all buffers that aren't already uploading
 	for (auto &write_buffer : to_flush) {
 		if (!write_buffer->uploading) {
@@ -475,6 +503,10 @@ void S3FileSystem::FlushAllBuffers(S3FileHandle &file_handle) {
 
 void S3FileSystem::FinalizeMultipartUpload(S3FileHandle &file_handle) {
 	auto &s3fs = (S3FileSystem &)file_handle.file_system;
+	if (file_handle.upload_finalized) {
+		return;
+	}
+
 	file_handle.upload_finalized = true;
 
 	std::stringstream ss;
@@ -889,8 +921,6 @@ void S3FileHandle::Initialize(optional_ptr<FileOpener> opener) {
 		part_size = ((minimum_part_size + Storage::DEFAULT_BLOCK_SIZE - 1) / Storage::DEFAULT_BLOCK_SIZE) *
 		            Storage::DEFAULT_BLOCK_SIZE;
 		D_ASSERT(part_size * max_part_count >= config_params.max_file_size);
-
-		multipart_upload_id = s3fs.InitializeMultipartUpload(*this);
 	}
 }
 

test/configs/duckdb-tests.json

@@ -11,6 +11,12 @@
         "test/sql/secrets/secret_autoloading_errors.test",
         "test/sql/copy/csv/zstd_crash.test"
       ]
+    },
+    {
+      "reason": "Improved from 1 PUT + 2 POST to 1 PUT",
+      "paths": [
+        "test/sql/copy/s3/metadata_cache.test"
+      ]
     }
   ]
 }

test/sql/copy/no_head_on_write.test

@@ -48,19 +48,18 @@ copy (select 1 as a) to 's3://test-bucket/test-file.parquet'
 query I
 select request.type FROM duckdb_logs_parsed('HTTP')
 ----
-POST
 PUT
-POST
 
 statement ok
 CALL truncate_duckdb_logs();
 
 statement ok
-copy (select 1 as a) to 's3://test-bucket/test-file.csv'
+copy (select random() as a FROM range(8000000)) to 's3://test-bucket/test-file2.csv'
 
 query I
 select request.type FROM duckdb_logs_parsed('HTTP')
 ----
 POST
 PUT
+PUT
 POST