Merge pull request #688 from Altinity/s3-roles

Enmk · zvonand · commit 0c1e13fc3455 · 2025-07-16T18:47:12.000Z
Role-based S3 access
diff --git a/contrib/aws-cmake/CMakeLists.txt b/contrib/aws-cmake/CMakeLists.txt
@@ -57,6 +57,10 @@ SET(AWS_SDK_S3_DIR "${AWS_SDK_DIR}/generated/src/aws-cpp-sdk-s3")
 SET(AWS_SDK_KMS_DIR "${AWS_SDK_DIR}/generated/src/aws-cpp-sdk-kms")
 SET(AWS_SDK_GLUE_DIR "${AWS_SDK_DIR}/generated/src/aws-cpp-sdk-glue")
 
+SET(AWS_SDK_IDENTITY_MANAGEMENT_DIR "${AWS_SDK_DIR}/src/aws-cpp-sdk-identity-management")
+SET(AWS_SDK_STS_DIR "${AWS_SDK_DIR}/generated/src/aws-cpp-sdk-sts")
+SET(AWS_SDK_COGNITO_IDENTITY_DIR "${AWS_SDK_DIR}/generated/src/aws-cpp-sdk-cognito-identity")
+
 SET(AWS_AUTH_DIR "${ClickHouse_SOURCE_DIR}/contrib/aws-c-auth")
 SET(AWS_CAL_DIR "${ClickHouse_SOURCE_DIR}/contrib/aws-c-cal")
 SET(AWS_CHECKSUMS_DIR "${ClickHouse_SOURCE_DIR}/contrib/aws-checksums")
@@ -145,6 +149,34 @@ list(APPEND AWS_SOURCES ${AWS_SDK_S3_SRC})
 list(APPEND AWS_PUBLIC_INCLUDES "${AWS_SDK_S3_DIR}/include/")
 
 
+# aws-cpp-sdk-identity-management
+file(GLOB AWS_SDK_IDENTITY_MANAGEMENT_SRC
+    "${AWS_SDK_IDENTITY_MANAGEMENT_DIR}/source/auth/*.cpp"
+)
+
+list(APPEND AWS_SOURCES ${AWS_SDK_IDENTITY_MANAGEMENT_SRC})
+list(APPEND AWS_PUBLIC_INCLUDES "${AWS_SDK_IDENTITY_MANAGEMENT_DIR}/include/")
+
+
+# aws-cpp-sdk-sts
+file(GLOB AWS_SDK_STS_SRC
+        "${AWS_SDK_STS_DIR}/source/*.cpp"
+        "${AWS_SDK_STS_DIR}/source/model/*.cpp"
+)
+
+list(APPEND AWS_SOURCES ${AWS_SDK_STS_SRC})
+list(APPEND AWS_PUBLIC_INCLUDES "${AWS_SDK_STS_DIR}/include/")
+
+# aws-cpp-sdk-cognito-identity
+file(GLOB AWS_SDK_COGNITO_IDENTITY_SRC
+        "${AWS_SDK_COGNITO_IDENTITY_DIR}/source/*.cpp"
+        "${AWS_SDK_COGNITO_IDENTITY_DIR}/source/model/*.cpp"
+)
+
+list(APPEND AWS_SOURCES ${AWS_SDK_COGNITO_IDENTITY_SRC})
+list(APPEND AWS_PRIVATE_INCLUDES "${AWS_SDK_COGNITO_IDENTITY_DIR}/include/")
+
+
 if(CLICKHOUSE_CLOUD)
     # aws-cpp-sdk-kms
     file(GLOB AWS_SDK_KMS_SRC
diff --git a/contrib/sparse-checkout/update-aws.sh b/contrib/sparse-checkout/update-aws.sh
@@ -6,10 +6,13 @@ FILES_TO_CHECKOUT=$(git rev-parse --git-dir)/info/sparse-checkout
 echo '/*' > $FILES_TO_CHECKOUT
 echo '!/*/*' >> $FILES_TO_CHECKOUT
 echo '/src/aws-cpp-sdk-core/*' >> $FILES_TO_CHECKOUT
+echo '/src/aws-cpp-sdk-identity-management/*' >> $FILES_TO_CHECKOUT
 echo '/generated/src/aws-cpp-sdk-s3/*' >> $FILES_TO_CHECKOUT
 echo '/generated/src/aws-cpp-sdk-aws/*' >> $FILES_TO_CHECKOUT
 echo '/generated/src/aws-cpp-sdk-glue/*' >> $FILES_TO_CHECKOUT
 echo '/generated/src/aws-cpp-sdk-kms/*' >> $FILES_TO_CHECKOUT
+echo '/generated/src/aws-cpp-sdk-sts/*' >> $FILES_TO_CHECKOUT
+echo '/generated/src/aws-cpp-sdk-cognito-identity/*' >> $FILES_TO_CHECKOUT
 
 git config core.sparsecheckout true
 git checkout $1
diff --git a/docs/en/sql-reference/table-functions/s3.md b/docs/en/sql-reference/table-functions/s3.md
@@ -22,8 +22,8 @@ When using the `s3 table function` with [`INSERT INTO...SELECT`](../../sql-refer
 
 ## Syntax {#syntax}
 
-```sql
-s3(url [, NOSIGN | access_key_id, secret_access_key, [session_token]] [,format] [,structure] [,compression_method],[,headers])
+``` sql
+s3(url [, NOSIGN | access_key_id, secret_access_key, [session_token]] [,format] [,structure] [,compression_method] [,headers] [,extra_credentials])
 s3(named_collection[, option=value [,..]])
 ```
 
@@ -47,6 +47,7 @@ For GCS, substitute your HMAC key and HMAC secret where you see `access_key_id`
 | `structure`                             | Structure of the table. Format `'column1_name column1_type, column2_name column2_type, ...'`.                                                                                                          |
 | `compression_method`                    | Parameter is optional. Supported values: `none`, `gzip` or `gz`, `brotli` or `br`, `xz` or `LZMA`, `zstd` or `zst`. By default, it will autodetect compression method by file extension.                 |
 | `headers`                               | Parameter is optional. Allows headers to be passed in the S3 request. Pass in the format `headers(key=value)` e.g. `headers('x-amz-request-payer' = 'requester')`.                                    |
+| `extra_credentials`                     | Parameter is optional. Allows to specify role ARN and role session name for AssumeRole (see below). Pass in the format `extra_credentials(key=value)`.                                    |
 
 :::note GCS
 The GCS url is in this format as the endpoint for the Google XML API is different than the JSON API:
@@ -280,7 +281,35 @@ Once configured, a `roleARN` can be passed to the s3 function via an `extra_cred
 SELECT count() FROM s3('https://datasets-documentation.s3.eu-west-3.amazonaws.com/mta/*.tsv','CSVWithNames',extra_credentials(role_arn = 'arn:aws:iam::111111111111:role/ClickHouseAccessRole-001'))
 ```
 
-Further examples can be found [here](/cloud/security/secure-s3#access-your-s3-bucket-with-the-clickhouseaccess-role)
+## Role Assumption
+
+ClickHouse supports assuming an AWS IAM role using a set of AWS credentials (`access_key_id`, `secret_access_key`, `session_token`).
+This allows ClickHouse to obtain temporary credentials for accessing an S3 bucket, even if the original credentials do not have direct access.
+
+For example, if the provided credentials have permission to assume a role but lack direct access to the S3 bucket, ClickHouse will first request temporary credentials from AWS STS and then use those credentials to access S3.
+
+For more details on role assumption, read [AWS AssumeRole documentation](https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRole.html).
+
+To enable role assumption, pass parameters via the extra_credentials argument in the s3 function. The following keys are supported:
+
+* `role_arn` (required) — ARN of the IAM role to assume. **If this key is not provided, ClickHouse will not attempt to assume a role and will use the original credentials as-is.**
+* `role_session_name` (optional) — Custom session name to include in the AssumeRole request.
+* `sts_endpoint_override` (optional) — Overrides the default AWS STS endpoint (https://sts.amazonaws.com). Useful for testing with a mock or when using another STS-compatible service.
+
+```sql
+SELECT count() FROM s3(
+    '<s3_bucket_uri>/*.csv',
+    access_key_id,
+    secret_access_key,
+    'CSVWithNames',
+    extra_credentials(
+        role_arn = 'arn:aws:iam::111111111111:role/BucketAccessRole-001',
+        role_session_name = 'ClickHouseSession',
+        sts_endpoint_override = 'http://mock-sts:8080'
+    )
+)
+```
+Further examples can be found [here](/docs/cloud/security/secure-s3#access-your-s3-bucket-with-the-clickhouseaccess-role)
 
 ## Working with archives {#working-with-archives}
 
diff --git a/src/Disks/ObjectStorages/S3/diskSettings.cpp b/src/Disks/ObjectStorages/S3/diskSettings.cpp
@@ -50,6 +50,9 @@ namespace S3AuthSetting
     extern const S3AuthSettingsString secret_access_key;
     extern const S3AuthSettingsString server_side_encryption_customer_key_base64;
     extern const S3AuthSettingsString session_token;
+    extern const S3AuthSettingsString role_arn;
+    extern const S3AuthSettingsString role_session_name;
+    extern const S3AuthSettingsString sts_endpoint_override;
     extern const S3AuthSettingsBool use_adaptive_timeouts;
     extern const S3AuthSettingsBool use_environment_credentials;
     extern const S3AuthSettingsBool use_insecure_imds_request;
@@ -175,6 +178,9 @@ std::unique_ptr<S3::Client> getClient(
         auth_settings[S3AuthSetting::use_insecure_imds_request],
         auth_settings[S3AuthSetting::expiration_window_seconds],
         auth_settings[S3AuthSetting::no_sign_request],
+        auth_settings[S3AuthSetting::role_arn],
+        auth_settings[S3AuthSetting::role_session_name],
+        auth_settings[S3AuthSetting::sts_endpoint_override]
     };
 
     return S3::ClientFactory::instance().create(
diff --git a/src/IO/S3/Credentials.cpp b/src/IO/S3/Credentials.cpp
@@ -39,6 +39,9 @@ namespace S3
 #    include <aws/core/utils/UUID.h>
 #    include <aws/core/http/HttpClientFactory.h>
 
+#    include <aws/sts/STSClient.h>
+#    include <aws/identity-management/auth/STSAssumeRoleCredentialsProvider.h>
+
 #    include <aws/core/utils/HashingUtils.h>
 #    include <aws/core/platform/FileSystem.h>
 
@@ -708,7 +711,35 @@ S3CredentialsProviderChain::S3CredentialsProviderChain(
     /// because it's manually defined by the user
     if (!credentials.IsEmpty())
     {
-        AddProvider(std::make_shared<Aws::Auth::SimpleAWSCredentialsProvider>(credentials));
+        if (credentials_configuration.role_arn.empty())
+            AddProvider(std::make_shared<Aws::Auth::SimpleAWSCredentialsProvider>(credentials));
+        else
+        {
+            auto sts_client_config = Aws::STS::STSClientConfiguration();
+
+            if (!credentials_configuration.sts_endpoint_override.empty())
+            {
+                auto endpoint_uri = Poco::URI(credentials_configuration.sts_endpoint_override);
+
+                String url_without_scheme = endpoint_uri.getHost();
+                if (endpoint_uri.getPort() != 0)
+                    url_without_scheme += ":" + std::to_string(endpoint_uri.getPort());
+
+                sts_client_config.endpointOverride = url_without_scheme;
+                sts_client_config.scheme = endpoint_uri.getScheme() == "https" ? Aws::Http::Scheme::HTTPS : Aws::Http::Scheme::HTTP;
+            }
+
+            AddProvider(std::make_shared<Aws::Auth::STSAssumeRoleCredentialsProvider>(
+                credentials_configuration.role_arn,
+                /* sessionName */ credentials_configuration.role_session_name,
+                /* externalId */ Aws::String(),
+                /* loadFrequency */ Aws::Auth::DEFAULT_CREDS_LOAD_FREQ_SECONDS,
+                std::make_shared<Aws::STS::STSClient>(credentials,
+                                                      /* endpointProvider */ Aws::MakeShared<Aws::STS::STSEndpointProvider>(Aws::STS::STSClient::ALLOCATION_TAG),
+                                                      /* clientConfiguration */ sts_client_config)
+                )
+            );
+        }
         return;
     }
 
diff --git a/src/IO/S3/Credentials.h b/src/IO/S3/Credentials.h
@@ -174,6 +174,9 @@ struct CredentialsConfiguration
     bool use_insecure_imds_request = false;
     uint64_t expiration_window_seconds = DEFAULT_EXPIRATION_WINDOW_SECONDS;
     bool no_sign_request = false;
+    String role_arn = ""; // NOLINT(*-redundant-string-init)
+    String role_session_name = ""; // NOLINT(*-redundant-string-init)
+    String sts_endpoint_override = ""; // NOLINT(*-redundant-string-init)
 };
 
 class S3CredentialsProviderChain : public Aws::Auth::AWSCredentialsProviderChain
diff --git a/src/IO/S3AuthSettings.cpp b/src/IO/S3AuthSettings.cpp
@@ -36,7 +36,10 @@ namespace ErrorCodes
     DECLARE(String, secret_access_key, "", "", 0) \
     DECLARE(String, session_token, "", "", 0) \
     DECLARE(String, region, "", "", 0) \
-    DECLARE(String, server_side_encryption_customer_key_base64, "", "", 0)
+    DECLARE(String, server_side_encryption_customer_key_base64, "", "", 0) \
+    DECLARE(String, role_arn, "", "", 0) \
+    DECLARE(String, role_session_name, "", "", 0) \
+    DECLARE(String, sts_endpoint_override, "", "", 0)
 
 #define CLIENT_SETTINGS_LIST(M, ALIAS) \
     CLIENT_SETTINGS(M, ALIAS) \
diff --git a/src/Storages/ObjectStorage/S3/Configuration.cpp b/src/Storages/ObjectStorage/S3/Configuration.cpp
@@ -14,6 +14,8 @@
 #include <Disks/ObjectStorages/S3/S3ObjectStorage.h>
 #include <Disks/ObjectStorages/S3/diskSettings.h>
 
+#include <Interpreters/evaluateConstantExpression.h>
+
 #include <Parsers/ASTFunction.h>
 #include <Parsers/ASTIdentifier.h>
 #include <Parsers/ASTLiteral.h>
@@ -46,6 +48,9 @@ namespace S3AuthSetting
     extern const S3AuthSettingsString secret_access_key;
     extern const S3AuthSettingsString session_token;
     extern const S3AuthSettingsBool use_environment_credentials;
+    extern const S3AuthSettingsString role_arn;
+    extern const S3AuthSettingsString role_session_name;
+    extern const S3AuthSettingsString sts_endpoint_override;
 }
 
 namespace ErrorCodes
@@ -182,8 +187,66 @@ void StorageS3Configuration::fromNamedCollection(const NamedCollection & collect
     keys = {url.key};
 }
 
+void StorageS3Configuration::extractExtraCreds(ASTs & args, ContextPtr context)
+{
+    ASTs::iterator extra_creds_it = args.end();
+
+    for (auto * arg_it = args.begin(); arg_it != args.end(); ++arg_it)
+    {
+        const auto * extra_creds_ast_function = (*arg_it)->as<ASTFunction>();
+        if (extra_creds_ast_function && extra_creds_ast_function->name == "extra_credentials")
+        {
+            if (extra_creds_it != args.end())
+                throw Exception(
+                    ErrorCodes::BAD_ARGUMENTS,
+                    "S3 table function can have only one extra_credentials argument");
+
+            const auto * extra_creds_function_args_expr = assert_cast<const ASTExpressionList *>(extra_creds_ast_function->arguments.get());
+            auto extra_creds_function_args = extra_creds_function_args_expr->children;
+
+            for (auto & extra_cred_arg : extra_creds_function_args)
+            {
+                const auto * extra_cred_ast = extra_cred_arg->as<ASTFunction>();
+                if (!extra_cred_ast || extra_cred_ast->name != "equals")
+                    throw Exception(ErrorCodes::BAD_ARGUMENTS, "extra_credentials argument is incorrect: shall be key=value");
+
+                const auto * extra_cred_args_expr = assert_cast<const ASTExpressionList *>(extra_cred_ast->arguments.get());
+                auto extra_cred_args = extra_cred_args_expr->children;
+                if (extra_cred_args.size() != 2)
+                    throw Exception(
+                        ErrorCodes::BAD_ARGUMENTS,
+                        "extra_credentials argument is incorrect: expected 2 arguments, got {}",
+                        extra_cred_args.size());
+
+                auto ast_literal = evaluateConstantExpressionOrIdentifierAsLiteral(extra_cred_args[0], context);
+                auto arg_name_value = ast_literal->as<ASTLiteral>()->value;
+                if (arg_name_value.getType() != Field::Types::Which::String)
+                    throw Exception(ErrorCodes::BAD_ARGUMENTS, "Expected string as extra_credentials name");
+                auto arg_name = arg_name_value.safeGet<String>();
+
+                ast_literal = evaluateConstantExpressionOrIdentifierAsLiteral(extra_cred_args[1], context);
+                auto arg_value = ast_literal->as<ASTLiteral>()->value;
+                if (arg_value.getType() != Field::Types::Which::String)
+                    throw Exception(ErrorCodes::BAD_ARGUMENTS, "Expected string as extra_credentials value");
+
+                extra_credentials_from_ast.emplace_back(arg_name, arg_value.safeGet<String>());
+            }
+
+            extra_creds_it = arg_it;
+            continue;
+        }
+    }
+
+    /// To avoid making unnecessary changes and avoid potential conflicts in future,
+    /// simply remove the "extra" argument after processing if it exists.
+    if (extra_creds_it != args.end())
+        args.erase(extra_creds_it);
+}
+
 void StorageS3Configuration::fromAST(ASTs & args, ContextPtr context, bool with_structure)
 {
+    extractExtraCreds(args, context);
+
     size_t count = StorageURL::evalArgsAndCollectHeaders(args, headers_from_ast, context);
 
     if (count == 0 || count > getMaxNumberOfArguments(with_structure))
@@ -381,6 +444,23 @@ void StorageS3Configuration::fromAST(ASTs & args, ContextPtr context, bool with_
     if (no_sign_request)
         auth_settings[S3AuthSetting::no_sign_request] = no_sign_request;
 
+    if (!extra_credentials_from_ast.empty())
+    {
+        auto extract_extra_cred_value = [&extra_creds = this->extra_credentials_from_ast](const String & cred_name) -> String
+        {
+            auto role_arn_it = std::find_if(extra_creds.begin(), extra_creds.end(),
+                                            [&cred_name](const HTTPHeaderEntry & entry) { return entry.name == cred_name; });
+            if (role_arn_it != extra_creds.end())
+                return role_arn_it->value;
+
+            return {};
+        };
+
+        auth_settings[S3AuthSetting::role_arn] = extract_extra_cred_value("role_arn");
+        auth_settings[S3AuthSetting::role_session_name] = extract_extra_cred_value("role_session_name");
+        auth_settings[S3AuthSetting::sts_endpoint_override] = extract_extra_cred_value("sts_endpoint_override");
+    }
+
     static_configuration = !auth_settings[S3AuthSetting::access_key_id].value.empty() || auth_settings[S3AuthSetting::no_sign_request].changed;
     auth_settings[S3AuthSetting::no_sign_request] = no_sign_request;
 
diff --git a/src/Storages/ObjectStorage/S3/Configuration.h b/src/Storages/ObjectStorage/S3/Configuration.h
@@ -37,7 +37,7 @@ class StorageS3Configuration : public StorageObjectStorage::Configuration
         " - url, access_key_id, secret_access_key, session_token, format, structure\n"
         " - url, access_key_id, secret_access_key, format, structure, compression_method\n"
         " - url, access_key_id, secret_access_key, session_token, format, structure, compression_method\n"
-        "All signatures supports optional headers (specified as `headers('name'='value', 'name2'='value2')`)";
+        "All signatures supports optional headers (specified as `headers('name'='value', 'name2'='value2')`) and extra credentials for role assumption (`extra_credentials(role_arn=value, role_session_name=value)`)";
 
     /// All possible signatures for S3 storage without structure argument (for example for S3 table engine).
     static constexpr auto max_number_of_arguments_without_structure = 6;
@@ -54,7 +54,7 @@ class StorageS3Configuration : public StorageObjectStorage::Configuration
         " - url, access_key_id, secret_access_key, session_token, format\n"
         " - url, access_key_id, secret_access_key, format, compression_method\n"
         " - url, access_key_id, secret_access_key, session_token, format, compression_method\n"
-        "All signatures supports optional headers (specified as `headers('name'='value', 'name2'='value2')`)";
+        "All signatures supports optional headers (specified as `headers('name'='value', 'name2'='value2')`) and extra credentials for role assumption (`extra_credentials(role_arn=value, role_session_name=value)`)";
 
     StorageS3Configuration() = default;
 
@@ -96,6 +96,7 @@ class StorageS3Configuration : public StorageObjectStorage::Configuration
         bool with_structure) override;
 
 private:
+    void extractExtraCreds(ASTs & args, ContextPtr context);
     void fromNamedCollection(const NamedCollection & collection, ContextPtr context) override;
     void fromAST(ASTs & args, ContextPtr context, bool with_structure) override;
 
@@ -105,6 +106,8 @@ class StorageS3Configuration : public StorageObjectStorage::Configuration
     S3::S3AuthSettings auth_settings;
     S3::S3RequestSettings request_settings;
     HTTPHeaderEntries headers_from_ast; /// Headers from ast is a part of static configuration.
+    HTTPHeaderEntries extra_credentials_from_ast; /// Avoid duplicated entities: HTTPHeaderEntry structure matches our needs here, use it.
+
     /// If s3 configuration was passed from ast, then it is static.
     /// If from config - it can be changed with config reload.
     bool static_configuration = true;
