From 692d7f5b3de29ea2a87d1726de5b675258f7901e Mon Sep 17 00:00:00 2001 From: Pankaj Date: Fri, 26 Apr 2024 20:29:07 +0530 Subject: [PATCH] Add assume_role_kwargs in hashicorp backend config --- .../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 ffc338217a3a9..a90188f996e77 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 @@ def __init__( 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 @@ def __init__( 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 @@ def _auth_aws_iam(self, _client: hvac.Client) -> None: 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 b29ae774612af..2591c77652e26 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 @@ def __init__( 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 @@ def __init__( 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,