Merge pull request #85 from gizmodata/main

Adding "requester pays" mode for S3 buckets

Author: samansmink

extension/httpfs/create_secret_functions.cpp

@@ -108,6 +108,12 @@ unique_ptr<BaseSecret> CreateS3SecretFunctions::CreateSecretFunctionInternal(Cli
 			}
 			refresh = true;
 			secret->secret_map["refresh_info"] = MapToStruct(named_param.second);
+		} else if (lower_name == "requester_pays") {
+			if (named_param.second.type() != LogicalType::BOOLEAN) {
+				throw InvalidInputException("Invalid type past to secret option: '%s', found '%s', expected: 'BOOLEAN'",
+											lower_name, named_param.second.type().ToString());
+			}
+			secret->secret_map["requester_pays"] = Value::BOOLEAN(named_param.second.GetValue<bool>());
 		} else {
 			throw InvalidInputException("Unknown named parameter passed to CreateSecretFunctionInternal: " +
 			                            lower_name);
@@ -185,6 +191,7 @@ void CreateS3SecretFunctions::SetBaseNamedParams(CreateSecretFunction &function,
 	function.named_parameters["use_ssl"] = LogicalType::BOOLEAN;
 	function.named_parameters["kms_key_id"] = LogicalType::VARCHAR;
 	function.named_parameters["url_compatibility_mode"] = LogicalType::BOOLEAN;
+    function.named_parameters["requester_pays"] = LogicalType::BOOLEAN;
 
 	// Whether a secret refresh attempt should be made when the secret appears to be incorrect
 	function.named_parameters["refresh"] = LogicalType::VARCHAR;

extension/httpfs/httpfs_extension.cpp

@@ -70,6 +70,7 @@ static void LoadInternal(DatabaseInstance &instance) {
 	config.AddExtensionOption("s3_kms_key_id", "S3 KMS Key ID", LogicalType::VARCHAR);
 	config.AddExtensionOption("s3_url_compatibility_mode", "Disable Globs and Query Parameters on S3 URLs",
 	                          LogicalType::BOOLEAN, Value(false));
+	config.AddExtensionOption("s3_requester_pays", "S3 use requester pays mode", LogicalType::BOOLEAN, Value(false));
 
 	// S3 Uploader config
 	config.AddExtensionOption("s3_uploader_max_filesize", "S3 Uploader max filesize (between 50GB and 5TB)",

extension/httpfs/include/s3fs.hpp

@@ -30,6 +30,7 @@ struct S3AuthParams {
 	string url_style;
 	bool use_ssl = true;
 	bool s3_url_compatibility_mode = false;
+    bool requester_pays = false;
 
 	static S3AuthParams ReadFrom(optional_ptr<FileOpener> opener, FileOpenerInfo &info);
 };
@@ -43,6 +44,8 @@ struct AWSEnvironmentCredentialsProvider {
 	static constexpr const char *DUCKDB_ENDPOINT_ENV_VAR = "DUCKDB_S3_ENDPOINT";
 	static constexpr const char *DUCKDB_USE_SSL_ENV_VAR = "DUCKDB_S3_USE_SSL";
 	static constexpr const char *DUCKDB_KMS_KEY_ID_ENV_VAR = "DUCKDB_S3_KMS_KEY_ID";
+	static constexpr const char *DUCKDB_REQUESTER_PAYS_ENV_VAR = "DUCKDB_S3_REQUESTER_PAYS";
+
 
 	explicit AWSEnvironmentCredentialsProvider(DBConfig &config) : config(config) {};
 

extension/httpfs/s3fs.cpp

@@ -62,7 +62,12 @@ static HTTPHeaders create_s3_header(string url, string query, string host, strin
 		res["x-amz-server-side-encryption-aws-kms-key-id"] = auth_params.kms_key_id;
 	}
 
-	string signed_headers = "";
+	bool use_requester_pays = auth_params.requester_pays;
+	if (use_requester_pays) {
+		res["x-amz-request-payer"] = "requester";
+	}
+
+    string signed_headers = "";
 	hash_bytes canonical_request_hash;
 	hash_str canonical_request_hash_str;
 	if (content_type.length() > 0) {
@@ -75,7 +80,10 @@ static HTTPHeaders create_s3_header(string url, string query, string host, strin
 	if (use_sse_kms) {
 		signed_headers += ";x-amz-server-side-encryption;x-amz-server-side-encryption-aws-kms-key-id";
 	}
-	auto canonical_request = method + "\n" + S3FileSystem::UrlEncode(url) + "\n" + query;
+	if (use_requester_pays) {
+		signed_headers += ";x-amz-request-payer";
+	}
+    auto canonical_request = method + "\n" + S3FileSystem::UrlEncode(url) + "\n" + query;
 	if (content_type.length() > 0) {
 		canonical_request += "\ncontent-type:" + content_type;
 	}
@@ -87,6 +95,9 @@ static HTTPHeaders create_s3_header(string url, string query, string host, strin
 		canonical_request += "\nx-amz-server-side-encryption:aws:kms";
 		canonical_request += "\nx-amz-server-side-encryption-aws-kms-key-id:" + auth_params.kms_key_id;
 	}
+	if (use_requester_pays) {
+		canonical_request += "\nx-amz-request-payer:requester";
+	}
 
 	canonical_request += "\n\n" + signed_headers + "\n" + payload_hash;
 	sha256(canonical_request.c_str(), canonical_request.length(), canonical_request_hash);
@@ -143,6 +154,7 @@ void AWSEnvironmentCredentialsProvider::SetAll() {
 	this->SetExtensionOptionValue("s3_endpoint", DUCKDB_ENDPOINT_ENV_VAR);
 	this->SetExtensionOptionValue("s3_use_ssl", DUCKDB_USE_SSL_ENV_VAR);
 	this->SetExtensionOptionValue("s3_kms_key_id", DUCKDB_KMS_KEY_ID_ENV_VAR);
+	this->SetExtensionOptionValue("s3_requester_pays", DUCKDB_REQUESTER_PAYS_ENV_VAR);
 }
 
 S3AuthParams AWSEnvironmentCredentialsProvider::CreateParams() {
@@ -156,6 +168,7 @@ S3AuthParams AWSEnvironmentCredentialsProvider::CreateParams() {
 	params.endpoint = DUCKDB_ENDPOINT_ENV_VAR;
 	params.kms_key_id = DUCKDB_KMS_KEY_ID_ENV_VAR;
 	params.use_ssl = DUCKDB_USE_SSL_ENV_VAR;
+    params.requester_pays = DUCKDB_REQUESTER_PAYS_ENV_VAR;
 
 	return params;
 }
@@ -181,6 +194,8 @@ S3AuthParams S3AuthParams::ReadFrom(optional_ptr<FileOpener> opener, FileOpenerI
 	secret_reader.TryGetSecretKeyOrSetting("kms_key_id", "s3_kms_key_id", result.kms_key_id);
 	secret_reader.TryGetSecretKeyOrSetting("s3_url_compatibility_mode", "s3_url_compatibility_mode",
 	                                       result.s3_url_compatibility_mode);
+	secret_reader.TryGetSecretKeyOrSetting("requester_pays", "s3_requester_pays",
+										   result.requester_pays);
 
 	// Endpoint and url style are slightly more complex and require special handling for gcs and r2
 	auto endpoint_result = secret_reader.TryGetSecretKeyOrSetting("endpoint", "s3_endpoint", result.endpoint);
@@ -219,6 +234,7 @@ unique_ptr<KeyValueSecret> CreateSecret(vector<string> &prefix_paths_p, string &
 	return_value->secret_map["use_ssl"] = params.use_ssl;
 	return_value->secret_map["kms_key_id"] = params.kms_key_id;
 	return_value->secret_map["s3_url_compatibility_mode"] = params.s3_url_compatibility_mode;
+	return_value->secret_map["requester_pays"] = params.requester_pays;
 
 	//! Set redact keys
 	return_value->redact_keys = {"secret", "session_token"};
@@ -531,9 +547,21 @@ void S3FileSystem::ReadQueryParams(const string &url_query_param, S3AuthParams &
 		}
 		query_params.erase(found_param);
 	}
+	auto found_requester_pays_param = query_params.find("s3_requester_pays");
+	if (found_requester_pays_param != query_params.end()) {
+		if (found_requester_pays_param->second == "true") {
+			params.requester_pays = true;
+		} else if (found_requester_pays_param->second == "false") {
+			params.requester_pays = false;
+		} else {
+			throw IOException("Incorrect setting found for s3_requester_pays, allowed values are: 'true' or 'false'");
+		}
+		query_params.erase(found_requester_pays_param);
+	}
 	if (!query_params.empty()) {
 		throw IOException("Invalid query parameters found. Supported parameters are:\n's3_region', 's3_access_key_id', "
-		                  "'s3_secret_access_key', 's3_session_token',\n's3_endpoint', 's3_url_style', 's3_use_ssl'");
+		                  "'s3_secret_access_key', 's3_session_token',\n's3_endpoint', 's3_url_style', 's3_use_ssl', "
+                          "'s3_requester_pays'");
 	}
 }
 

test/sql/secret/secret_s3_requester_pays.test

@@ -0,0 +1,48 @@
+# name: test/sql/secret/secret_aws_requester_pays.test
+# description: Tests secret refreshing with AWS requester pays mode
+# group: [secrets]
+
+require-env S3_TEST_SERVER_AVAILABLE 1
+
+require-env AWS_DEFAULT_REGION
+
+require-env AWS_ACCESS_KEY_ID
+
+require-env AWS_SECRET_ACCESS_KEY
+
+require-env DUCKDB_S3_ENDPOINT
+
+require-env DUCKDB_S3_USE_SSL
+
+require httpfs
+
+require parquet
+
+statement ok
+SET enable_logging=true
+
+statement ok
+set s3_use_ssl='${DUCKDB_S3_USE_SSL}'
+
+statement ok
+set s3_endpoint='${DUCKDB_S3_ENDPOINT}'
+
+statement ok
+set s3_region='${AWS_DEFAULT_REGION}'
+
+# Create some test data
+statement ok
+CREATE SECRET s1 (
+    TYPE S3,
+    KEY_ID '${AWS_ACCESS_KEY_ID}',
+    SECRET '${AWS_SECRET_ACCESS_KEY}',
+    REQUESTER_PAYS true
+)
+
+statement ok
+copy (select 1 as a) to 's3://test-bucket/test-file.parquet'
+
+query I
+FROM "s3://test-bucket/test-file.parquet"
+----
+1