SNOW-2436496: Workload Identity Federation using Managed Identities in AKS

[kristianandre]

What is the current behavior?

Right now, it seems like the Snowflake connector and client libraries will try to fetch the token for us, which fails to discover the user-assigned managed identity for workloads in AKS. That is probably the reason you suggest a different flow for AKS that uses service accounts instead instead of Entra ID and managed identites.

What is the desired behavior?

We can discover this user-assigned managed identity and fetch a token for it from our AKS by using the Azure identity libraries. Is it possible to provide the token manually? Does the Snowflake connector have to do it for us?

How would this improve snowflake-connector-python?

It will allow us to use managed identities to authenticate to Snowflake without needing to use service accounts in AKS directly. Since we already use managed identities for everything else, this will allow us to use the same pattern.

References and other background

No response

[sfc-gh-dszmolka]

hi - the team will take a look and consider this request for future planning sessions, no timeline attached. I would guess this is not only a request for this driver, but need the same functionality across all Snowflake drivers.
Disclaimer as always; if this capability is important for your use-case, please reach out to your Snowflake Sales rep., and your account team might be able to help prioritizie this request.

[kristianandre]

Thanks @sfc-gh-dszmolka and @sfc-gh-snow-drivers-warsaw-dl!

I've tried to get around this limitation in several ways.

1.

with connect(
    account="*.west-europe.azure",
    authenticator="WORKLOAD_IDENTITY",
    workload_identity_provider="OIDC",
    token=acquire_token("api://fd3f753b-eed3-462c-b6a7-a4b5bb650aad/.default"),
) as con:
    with con.cursor() as cur:
        cur.execute("SELECT CURRENT_USER();")
        print(cur.fetchone())

I use OIDC because that's the only provider that will actually use the token I give. This fails with the following error message though:

snowflake.connector.errors.DatabaseError: 250001 (08001): None: Failed to connect to DB: *.west-europe.azure.snowflakecomputing.com:443.
An error occurred while attempting to authenticate the JWT with issuer 'https://login.microsoftonline.com/*/v2.0' and subject '*'.
Either the subject or issuer claims were not recognized, or the JWTs signature could not be verified.

I have double-checked and we have set up correct issuer and subject for the user in Snowflake, so most likely this fails due to the provider not being AZURE?

As far as I can tell, the only change for you would be to have a token parameter here, as you do for OIDC. Then in the body of create_azure_attestation(), you first check if token is None and fetch the token as you do today, otherwise one can skip all the way down here.

2.

Second option I tried was creating a custom class:

class CustomAuthByWorkloadIdentity(AuthByWorkloadIdentity):
    def prepare(self, *, conn: SnowflakeConnection | None, **kwargs: Any) -> None:
        """Fetch the token using azure-identity."""
        token = acquire_token("api://fd3f753b-eed3-462c-b6a7-a4b5bb650aad/.default")
        iss, sub = extract_iss_and_sub_without_signature_verification(token)
        self.attestation = WorkloadIdentityAttestation(
            provider=AttestationProvider.AZURE,
            credential=token,
            user_identifier_components={"iss": iss, "sub": sub},
        )

    @property
    def type_(self) -> AuthType:
        return AuthType.WORKLOAD_IDENTITY

with connect(
    account="cb20535.west-europe.azure",
    auth_class=AuthByWorkloadIdentityLynx(),
) as con:
    with con.cursor() as cur:
        cur.execute("SELECT CURRENT_USER();")
        print(cur.fetchone())

But that also raises an error because one is not allowed to provide a custom auth_class of type AuthByWorkloadIdentity, only AuthByKeyPair.

[kristianandre]

I found a workaround:

CREATE USER my_custom_service
  WORKLOAD_IDENTITY = (
    TYPE = OIDC
    ISSUER = 'https://login.microsoftonline.com/<tenant_id>/v2.0'
    SUBJECT = '<managed_identity_object_id>'
    OIDC_AUDIENCE_LIST = ('<snowflake_workload_identity_app_id>')
  )
  TYPE = SERVICE;

[sfc-gh-dszmolka]

thank you for sharing the workaround, this is helpful. this part looks to be documented in https://docs.snowflake.com/en/user-guide/workload-identity-federation#label-wif-oidc-custom-configure-snowflake , although definitely not in a detailed manner like you provided.