Habilitar a federação de identidades de carga de trabalho para cargas de trabalho de IAM na AWS

As cargas de trabalho do AWS podem se autenticar em Azure Databricks sem segredos de longo prazo trocando um token de identidade OIDC assinado pela AWS. O caminho recomendado chama a API STS GetWebIdentityToken da AWS, que funciona em qualquer lugar em que a carga de trabalho tenha credenciais de AWS.

Casos de uso

  • Funções Lambda que chamam APIs do Azure Databricks (acionando trabalhos, consultando armazéns SQL)
  • Pipelines de ETL baseados no EC2/ECS que se autenticam no Azure Databricks sem segredos
  • Cargas de trabalho de ML baseadas no EKS que acessam pontos de extremidade do Serviço de Modelo
  • Padrões entre contas em que cargas de trabalho em uma conta do AWS são federadas em uma conta Azure Databricks gerenciada por outra equipe
  • Aprimoramento da postura de segurança pela eliminação de PATs de longa duração do Azure Databricks ou segredos do AWS Secrets Manager

Pré-requisitos do AWS

As etapas a seguir devem ser concluídas para habilitar a federação de identidade para cargas de trabalho do AWS IAM.

Etapa 1: Habilitar a federação de identidade de saída do AWS IAM

Habilite a federação de identidade de saída em sua conta do AWS:

import boto3
boto3.client('iam').enable_outbound_web_identity_federation()

Você também pode habilitar isso no Console do IAM em Configurações de conta>Habilitar "Federação de identidade web de saída".

Etapa 2: Conceder sts:GetWebIdentityToken permissão

Conceda à função de IAM da carga de trabalho a permissão sts:GetWebIdentityToken:

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": ["sts:GetWebIdentityToken"],
      "Resource": "*",
      "Condition": {
        "ForAllValues:StringEquals": {
          "sts:IdentityTokenAudience": "databricks"
        },
        "NumericLessThanEquals": {
          "sts:DurationSeconds": 300
        }
      }
    }
  ]
}

Observação

A condição de público-alvo garante que a função só possa solicitar tokens direcionados ao Azure Databricks. A condição de duração limita o tempo de vida do token a 300 segundos. Você pode ajustar a duração até 3600 segundos com base nas suas necessidades de carga de trabalho, mas tempos de vida mais curtos são recomendados.

Etapa 3: Observe a URL do emissor específico da conta

Recupere a URL do emissor específica da conta:

import boto3
info = boto3.client('iam').get_outbound_web_identity_federation_info()
print(info['IssuerUrl'])  # https://<uuid>.tokens.sts.global.api.aws

Criar uma política de federação

Importante

Azure Databricks políticas de federação são criadas no nível da conta (não no nível do workspace). O host da CLI Azure Databricks deve ser definido https://accounts.azuredatabricks.net e o usuário deve ser um administrador de conta.

Crie uma política de federação de identidade de carga de trabalho usando a CLI do Databricks. Defina o emissor como a URL do emissor específico da sua conta na Etapa 3. Para obter instruções detalhadas, consulte Configurar uma política de federação de entidade de serviço.

databricks account service-principal-federation-policy create ${SP_ID} --json '{
  "oidc_policy": {
    "issuer": "https://<uuid>.tokens.sts.global.api.aws",
    "audiences": ["databricks"],
    "subject": "arn:aws:iam::<account-id>:role/<workload-role-name>"
  }
}'

Autentique-se no Databricks

Depois de criar a política de federação, use o SDK do Databricks para autenticar suas cargas de trabalho do AWS. O exemplo a seguir usa o padrão do IdTokenSource SDK para recuperar um token STS do AWS e trocá-lo por um token OAuth Azure Databricks.

import boto3
from databricks.sdk import WorkspaceClient
from databricks.sdk import oidc
from databricks.sdk.core import Config, credentials_strategy, oidc_credentials_provider


class AwsStsTokenSource(oidc.IdTokenSource):
    def __init__(self, audience="databricks", region="us-east-1"):
        self._audience = audience
        self._region = region

    def id_token(self) -> oidc.IdToken:
        sts = boto3.client("sts", region_name=self._region)
        resp = sts.get_web_identity_token(
            Audience=[self._audience],
            SigningAlgorithm="RS256",
            DurationSeconds=300,
        )
        return oidc.IdToken(jwt=resp["WebIdentityToken"])


@credentials_strategy("aws-sts-wif", [])
def aws_sts_wif_strategy(cfg: Config):
    return oidc_credentials_provider(cfg, AwsStsTokenSource())


w = WorkspaceClient(
    host="https://my-workspace.cloud.databricks.com",
    client_id="<service-principal-uuid>",
    credentials_strategy=aws_sts_wif_strategy
)
# No secrets needed
clusters = w.clusters.list()

Observação

A duração do token de 300 segundos é recomendada. Você pode ajustar até 3600 segundos com base nas necessidades de carga de trabalho.

Para obter um exemplo manual de troca de tokens, consulte Autenticar com um token de provedor de identidade.

Referências de documentação do AWS

  • Last updated on