Add assumerole with creds from metadata

fix sts_client

remove commented code

add docs

Author: zvonand

docs/en/sql-reference/table-functions/s3.md

@@ -283,19 +283,21 @@ SELECT count() FROM s3('https://datasets-documentation.s3.eu-west-3.amazonaws.co
 
 ## 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.
+ClickHouse supports assuming an AWS IAM role using a set of AWS credentials (`access_key_id`, `secret_access_key`, `session_token`) or EC2 metadata (only when running on EC2 instance).
+This allows ClickHouse to obtain temporary credentials for accessing an S3 bucket, even if the original credentials or instance 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:
+To use this mechanism, pass parameters via the `extra_credentials` argument to 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.
+* `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 try to access the bucket as-is.**
+* `role_session_name` (optional) — Custom session name to include in the AssumeRole request. If not specified, a random UUID will be assigned.
 * `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.
 
+If explicit `access_key_id` and `secret_access_key` are provided as parameters to `s3(...)` function, then they will be used for retrieving temporary credentials from STS:
+
 ```sql
 SELECT count() FROM s3(
     '<s3_bucket_uri>/*.csv',
@@ -309,6 +311,20 @@ SELECT count() FROM s3(
     )
 )
 ```
+
+Otherwise, ClickHouse will attempt to extract credentials from EC2 metadata:
+
+```sql
+SELECT count() FROM s3(
+    '<s3_bucket_uri>/*.csv',
+    'CSVWithNames',
+    extra_credentials(
+        role_arn = 'arn:aws:iam::111111111111:role/BucketAccessRole-001',
+        role_session_name = 'ClickHouseSession'
+    )
+)
+```
+
 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}

src/IO/S3/Credentials.cpp

@@ -563,6 +563,86 @@ void AwsAuthSTSAssumeRoleWebIdentityCredentialsProvider::refreshIfExpired()
     Reload();
 }
 
+AWSInstanceMetadataAssumeRoleCredentialsProvider::AWSInstanceMetadataAssumeRoleCredentialsProvider(
+    const Aws::String & role_arn_,
+    const Aws::String & session_name_,
+    DB::S3::PocoHTTPClientConfiguration aws_client_configuration,
+    uint64_t expiration_window_seconds_)
+    : role_arn(role_arn_)
+    , session_name(session_name_.empty() ? Aws::String(Aws::Utils::UUID::RandomUUID()) : session_name_)
+    , logger(getLogger("AWSInstanceMetadataAssumeRoleCredentialsProvider"))
+    , expiration_window_seconds(expiration_window_seconds_)
+{
+    // Create metadata credentials provider
+    auto ec2_metadata_client = createEC2MetadataClient(aws_client_configuration);
+    auto config_loader = std::make_shared<AWSEC2InstanceProfileConfigLoader>(ec2_metadata_client, true);
+    metadata_provider = std::make_shared<AWSInstanceProfileCredentialsProvider>(config_loader);
+
+    aws_client_configuration.scheme = Aws::Http::Scheme::HTTPS;
+
+    std::vector<Aws::String> retryable_errors;
+    retryable_errors.push_back("IDPCommunicationError");
+    retryable_errors.push_back("InvalidIdentityToken");
+    aws_client_configuration.retryStrategy = std::make_shared<Aws::Client::SpecifiedRetryableErrorsRetryStrategy>(
+        retryable_errors, /*maxRetries=*/ 3);
+
+    sts_client = std::make_unique<Aws::STS::STSClient>(metadata_provider, aws_client_configuration);
+
+    LOG_INFO(logger, "Created STS AssumeRole provider using EC2 instance metadata");
+}
+
+Aws::Auth::AWSCredentials AWSInstanceMetadataAssumeRoleCredentialsProvider::GetAWSCredentials()
+{
+    refreshIfExpired();
+    Aws::Utils::Threading::ReaderLockGuard guard(m_reloadLock);
+    return credentials;
+}
+
+void AWSInstanceMetadataAssumeRoleCredentialsProvider::Reload()
+{
+    // Get fresh metadata credentials
+    auto metadata_creds = metadata_provider->GetAWSCredentials();
+    if (metadata_creds.IsEmpty())
+    {
+        LOG_ERROR(logger, "Failed to obtain instance metadata credentials");
+        return;
+    }
+
+    // Perform AssumeRole
+    Aws::STS::Model::AssumeRoleRequest request;
+    request.SetRoleArn(role_arn);
+    request.SetRoleSessionName(session_name);
+
+    auto outcome = sts_client->AssumeRole(request);
+    if (!outcome.IsSuccess())
+    {
+        LOG_ERROR(logger, "Failed to assume role: {}", outcome.GetError().GetMessage());
+        return;
+    }
+
+    const auto & result = outcome.GetResult().GetCredentials();
+    credentials = Aws::Auth::AWSCredentials(
+        result.GetAccessKeyId(),
+        result.GetSecretAccessKey(),
+        result.GetSessionToken(),
+        result.GetExpiration().SecondsWithMSPrecision());
+
+    LOG_TRACE(logger, "Successfully assumed role using metadata credentials");
+}
+
+void AWSInstanceMetadataAssumeRoleCredentialsProvider::refreshIfExpired()
+{
+    Aws::Utils::Threading::ReaderLockGuard guard(m_reloadLock);
+    if (!areCredentialsEmptyOrExpired(credentials, expiration_window_seconds))
+        return;
+
+    guard.UpgradeToWriterLock();
+    if (!areCredentialsEmptyOrExpired(credentials, expiration_window_seconds))
+        return;
+
+    Reload();
+}
+
 
 SSOCredentialsProvider::SSOCredentialsProvider(DB::S3::PocoHTTPClientConfiguration aws_client_configuration_, uint64_t expiration_window_seconds_)
     : profile_to_use(Aws::Auth::GetConfigProfileName())
@@ -707,13 +787,17 @@ S3CredentialsProviderChain::S3CredentialsProviderChain(
     if (credentials_configuration.no_sign_request)
         return;
 
-    /// add explicit credentials to the front of the chain
-    /// because it's manually defined by the user
-    if (!credentials.IsEmpty())
+    if (credentials_configuration.role_arn.empty())
     {
-        if (credentials_configuration.role_arn.empty())
+        if (!credentials.IsEmpty())
+        {
             AddProvider(std::make_shared<Aws::Auth::SimpleAWSCredentialsProvider>(credentials));
-        else
+            return;
+        }
+    }
+    else
+    {
+        if (!credentials.IsEmpty())
         {
             auto sts_client_config = Aws::STS::STSClientConfiguration();
 
@@ -740,6 +824,26 @@ S3CredentialsProviderChain::S3CredentialsProviderChain(
                 )
             );
         }
+        else
+        {
+            DB::S3::PocoHTTPClientConfiguration aws_client_configuration = DB::S3::ClientFactory::instance().createClientConfiguration(
+                configuration.region,
+                configuration.remote_host_filter,
+                configuration.s3_max_redirects,
+                configuration.s3_retry_attempts,
+                configuration.enable_s3_requests_logging,
+                configuration.for_disk_s3,
+                configuration.get_request_throttler,
+                configuration.put_request_throttler);
+
+            // Use metadata credentials for AssumeRole
+            AddProvider(std::make_shared<AWSInstanceMetadataAssumeRoleCredentialsProvider>(
+                Aws::String(credentials_configuration.role_arn),
+                Aws::String(credentials_configuration.role_session_name),
+                std::move(aws_client_configuration),
+                credentials_configuration.expiration_window_seconds)
+                );
+        }
         return;
     }
 

src/IO/S3/Credentials.h

@@ -11,6 +11,8 @@
 #    include <aws/core/config/AWSProfileConfigLoader.h>
 #    include <aws/core/auth/AWSCredentialsProviderChain.h>
 #    include <aws/core/auth/bearer-token-provider/SSOBearerTokenProvider.h>
+#    include <aws/sts/STSClient.h>
+#    include <aws/sts/model/AssumeRoleRequest.h>
 
 #    include <IO/S3/PocoHTTPClient.h>
 #    include <IO/S3Defines.h>
@@ -108,6 +110,32 @@ class AWSInstanceProfileCredentialsProvider : public Aws::Auth::AWSCredentialsPr
     LoggerPtr logger;
 };
 
+class AWSInstanceMetadataAssumeRoleCredentialsProvider : public Aws::Auth::AWSCredentialsProvider
+{
+public:
+    explicit AWSInstanceMetadataAssumeRoleCredentialsProvider(
+        const Aws::String & role_arn_,
+        const Aws::String & session_name_,
+        DB::S3::PocoHTTPClientConfiguration aws_client_configuration,
+        uint64_t expiration_window_seconds_);
+
+    Aws::Auth::AWSCredentials GetAWSCredentials() override;
+
+protected:
+    void Reload() override;
+
+private:
+    void refreshIfExpired();
+
+    std::shared_ptr<AWSInstanceProfileCredentialsProvider> metadata_provider;
+    std::unique_ptr<Aws::STS::STSClient> sts_client;
+    Aws::Auth::AWSCredentials credentials;
+    Aws::String role_arn;
+    Aws::String session_name;
+    LoggerPtr logger;
+    uint64_t expiration_window_seconds;
+};
+
 class AwsAuthSTSAssumeRoleWebIdentityCredentialsProvider : public Aws::Auth::AWSCredentialsProvider
 {
     /// See STSAssumeRoleWebIdentityCredentialsProvider.

src/Storages/ObjectStorage/S3/Configuration.cpp

@@ -233,7 +233,6 @@ void StorageS3Configuration::extractExtraCreds(ASTs & args, ContextPtr context)
             }
 
             extra_creds_it = arg_it;
-            continue;
         }
     }