Skip to content

Commit

Permalink
Allow role-based authentication for Athena
Browse files Browse the repository at this point in the history
You can now specify role_arn to assume a role for querying Athena,
rather than using the access key & secret directly.

STS requires a region, which might not be the same region as the
database's region. To try & clean up the distinction I've moved all the
Athena credential settings to a new `credentials` sub-hash.
  • Loading branch information
jdelStrother committed Sep 30, 2024
1 parent 16b1b62 commit efde307
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 8 deletions.
9 changes: 7 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -601,9 +601,14 @@ data_sources:
# optional settings
output_location: s3://some-bucket/
workgroup: primary
access_key_id: ...
secret_access_key: ...
region: ...
credentials:
access_key_id: ...
secret_access_key: ...
# optional credential-settings, for role-based authentication:
role_arn: ...
region: ...
```

Here’s an example IAM policy:
Expand Down
34 changes: 28 additions & 6 deletions lib/blazer/adapters/athena_adapter.rb
Original file line number Diff line number Diff line change
Expand Up @@ -168,13 +168,35 @@ def glue
end

def client_options
@client_options ||= begin
options = {}
if settings["access_key_id"] || settings["secret_access_key"]
options[:credentials] = Aws::Credentials.new(settings["access_key_id"], settings["secret_access_key"])
options = {}
if credentials = client_credentials
options[:credentials] = credentials
end
options[:region] = settings["region"] if settings["region"]
options
end

def client_credentials
@client_credentials ||= begin
# Loading the access key & secret from the top-level settings is supported for backwards compatibility,
# but prefer loading them from the 'credentials' sub-hash.
creds = (settings["credentials"] || {}).with_defaults(settings.slice("access_key_id", "secret_access_key", "region"))
access_key_id = creds["access_key_id"]
secret_access_key = creds["secret_access_key"]
role_arn = creds["role_arn"]
if role_arn
region = creds["region"]
role_session_name = creds["role_session_name"] || "blazer"
Aws::AssumeRoleCredentials.new(
access_key_id: access_key_id,
secret_access_key: secret_access_key,
region: region,
role_arn: role_arn,
role_session_name: role_session_name,
)
elsif access_key_id && secret_access_key
Aws::Credentials.new(access_key_id, secret_access_key)
end
options[:region] = settings["region"] if settings["region"]
options
end
end
end
Expand Down

0 comments on commit efde307

Please sign in to comment.