This is an automated email from the ASF dual-hosted git repository.
pankaj pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new a3f0d83679 Add assume_role_kwargs in hashicorp backend config (#39279)
a3f0d83679 is described below
commit a3f0d836796c34d3643842f1674f1d5b04bb34bd
Author: Pankaj Singh <[email protected]>
AuthorDate: Sat Apr 27 11:09:11 2024 +0530
Add assume_role_kwargs in hashicorp backend config (#39279)
Add assume_role_kwargs in hashicorp backend config
```
AIRFLOW__SECRETS__BACKEND_KWARGS='{"kv_engine_version": 1, "mount_point":
"kv", "variables_path": "airflow", "url": "http://127.0.0.0:8200/";,
"auth_type": "aws_iam", "assume_role_kwargs": {"RoleArn":
"arn:aws:iam::1234567890000:role/hashicorp-aws-iam", "RoleSessionName":
"airflow", "DurationSeconds": 900}}'
```
---
airflow/providers/hashicorp/_internal_client/vault_client.py | 12 +++++++-----
airflow/providers/hashicorp/secrets/vault.py | 8 +++++---
2 files changed, 12 insertions(+), 8 deletions(-)
diff --git a/airflow/providers/hashicorp/_internal_client/vault_client.py
b/airflow/providers/hashicorp/_internal_client/vault_client.py
index ffc338217a..a90188f996 100644
--- a/airflow/providers/hashicorp/_internal_client/vault_client.py
+++ b/airflow/providers/hashicorp/_internal_client/vault_client.py
@@ -74,7 +74,9 @@ class _VaultClient(LoggingMixin):
:param key_id: Key ID for Authentication (for ``aws_iam`` and ''azure``
auth_type).
:param secret_id: Secret ID for Authentication (for ``approle``,
``aws_iam`` and ``azure`` auth_types).
:param role_id: Role ID for Authentication (for ``approle``, ``aws_iam``
auth_types).
- :param role_arn: AWS arn role (for ``aws_iam`` auth_type)
+ :param assume_role_kwargs: AWS assume role param.
+ See AWS STS Docs:
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role.html
:param kubernetes_role: Role for Authentication (for ``kubernetes``
auth_type).
:param kubernetes_jwt_path: Path for kubernetes jwt token (for
``kubernetes`` auth_type, default:
``/var/run/secrets/kubernetes.io/serviceaccount/token``).
@@ -104,7 +106,7 @@ class _VaultClient(LoggingMixin):
password: str | None = None,
key_id: str | None = None,
secret_id: str | None = None,
- role_arn: str | None = None,
+ assume_role_kwargs: dict | None = None,
role_id: str | None = None,
kubernetes_role: str | None = None,
kubernetes_jwt_path: str | None =
"/var/run/secrets/kubernetes.io/serviceaccount/token",
@@ -163,7 +165,7 @@ class _VaultClient(LoggingMixin):
self.key_id = key_id
self.secret_id = secret_id
self.role_id = role_id
- self.role_arn = role_arn
+ self.assume_role_kwargs = assume_role_kwargs
self.kubernetes_role = kubernetes_role
self.kubernetes_jwt_path = kubernetes_jwt_path
self.gcp_key_path = gcp_key_path
@@ -330,9 +332,9 @@ class _VaultClient(LoggingMixin):
else:
import boto3
- if self.role_arn:
+ if self.assume_role_kwargs:
sts_client = boto3.client("sts")
- credentials = sts_client.assume_role(RoleArn=self.role_arn,
RoleSessionName="airflow")
+ credentials = sts_client.assume_role(**self.assume_role_kwargs)
auth_args = {
"access_key": credentials["Credentials"]["AccessKeyId"],
"secret_key":
credentials["Credentials"]["SecretAccessKey"],
diff --git a/airflow/providers/hashicorp/secrets/vault.py
b/airflow/providers/hashicorp/secrets/vault.py
index b29ae77461..2591c77652 100644
--- a/airflow/providers/hashicorp/secrets/vault.py
+++ b/airflow/providers/hashicorp/secrets/vault.py
@@ -74,7 +74,9 @@ class VaultBackend(BaseSecretsBackend, LoggingMixin):
:param key_id: Key ID for Authentication (for ``aws_iam`` and ''azure``
auth_type).
:param secret_id: Secret ID for Authentication (for ``approle``,
``aws_iam`` and ``azure`` auth_types).
:param role_id: Role ID for Authentication (for ``approle``, ``aws_iam``
auth_types).
- :param role_arn: AWS arn role,
+ :param assume_role_kwargs: AWS assume role param.
+ See AWS STS Docs:
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role.html
:param kubernetes_role: Role for Authentication (for ``kubernetes``
auth_type).
:param kubernetes_jwt_path: Path for kubernetes jwt token (for
``kubernetes`` auth_type, default:
``/var/run/secrets/kubernetes.io/serviceaccount/token``).
@@ -108,7 +110,7 @@ class VaultBackend(BaseSecretsBackend, LoggingMixin):
key_id: str | None = None,
secret_id: str | None = None,
role_id: str | None = None,
- role_arn: str | None = None,
+ assume_role_kwargs: dict | None = None,
kubernetes_role: str | None = None,
kubernetes_jwt_path: str =
"/var/run/secrets/kubernetes.io/serviceaccount/token",
gcp_key_path: str | None = None,
@@ -149,7 +151,7 @@ class VaultBackend(BaseSecretsBackend, LoggingMixin):
key_id=key_id,
secret_id=secret_id,
role_id=role_id,
- role_arn=role_arn,
+ assume_role_kwargs=assume_role_kwargs,
kubernetes_role=kubernetes_role,
kubernetes_jwt_path=kubernetes_jwt_path,
gcp_key_path=gcp_key_path,
