kaxil commented on code in PR #75810:
URL: https://github.com/apache/airflow/pull/75810#discussion_r4473329254
##########
providers/snowflake/docs/connections/snowflake.rst:
##########
@@ -189,3 +189,30 @@ and is granted a role with access to the target objects.
"workload_identity_provider": "GCP"
}
}'
+
+``SnowflakeSqlApiHook``, ``SnowflakeCortexAgentHook`` and the
``pydanticai_snowflake`` connection type call
+Snowflake's REST APIs, which accept the workload's identity token directly
instead of a connector session.
+The token is requested from the cloud and reused until shortly before it
expires. A connection that also
+sets a private key uses key-pair authentication for these REST hooks. Each
provider needs its cloud's client
+library installed:
+
+* ``AWS``: ``boto3`` with ``botocore>=1.41.0`` for ``sts:GetWebIdentityToken``
support.
+ The workload's IAM role must be allowed to
+ call it, and AWS outbound identity federation must be enabled for the
account. This requirement applies
+ only to the REST hooks; ``SnowflakeHook`` uses a signed
``GetCallerIdentity`` request by default and
+ does not require outbound identity federation. The call is not available on
Review Comment:
With the new sentence in between, "The call is not available on the global
STS endpoint" now reads as if it's about `GetCallerIdentity`. Naming
`GetWebIdentityToken` there, or moving the `SnowflakeHook` sentence to the end
of the bullet, would fix it.
##########
providers/snowflake/tests/unit/snowflake/utils/test__rest_auth.py:
##########
@@ -148,20 +166,170 @@ def
test_get_token_key_pair_raises_without_private_key(self, mock_conn_params):
@mock.patch.object(SnowflakeHook, "_get_conn_params", autospec=True)
def test_get_token_workload_identity_with_private_key_uses_key_pair(self,
mock_conn_params):
- mock_conn_params.return_value = {**CONN_PARAMS_KEYPAIR,
"workload_identity_provider": "AWS"}
+ mock_conn_params.return_value = {
+ **CONN_PARAMS_KEYPAIR,
+ "authenticator": "WORKLOAD_IDENTITY",
+ "workload_identity_provider": "AWS",
+ }
key = rsa.generate_private_key(backend=default_backend(),
public_exponent=65537, key_size=2048)
hook = SnowflakeHook(snowflake_conn_id="mock_conn_id")
provider = SnowflakeRestTokenProvider(hook, private_key_loader=lambda:
key)
assert provider.get_token().token_type == "KEYPAIR_JWT"
@mock.patch.object(SnowflakeHook, "_get_conn_params", autospec=True)
- def test_get_token_workload_identity_without_private_key_raises(self,
mock_conn_params):
- mock_conn_params.return_value = {**CONN_PARAMS_KEYPAIR,
"workload_identity_provider": "AWS"}
+ @mock.patch(FETCH_TOKEN.format("gcp"), autospec=True)
+ def test_private_key_absence_cached_across_wif_refresh(self, fetch,
conn_params, time_machine):
+ time_machine.move_to("2024-01-01T00:00:00Z", tick=False)
+ conn_params.return_value = CONN_PARAMS_WIF
+ fetch.return_value = _make_jwt(exp=int(time.time()) + 3600)
+ loader = mock.Mock(spec=SnowflakeHook().get_private_key,
return_value=None)
+ provider = SnowflakeRestTokenProvider(SnowflakeHook(),
private_key_loader=loader)
+ provider.get_token()
+ provider.get_token()
+ time_machine.shift(timedelta(hours=1))
+ provider.get_token()
+ loader.assert_called_once_with()
+ assert fetch.call_count == 2
+
+ @mock.patch.object(SnowflakeHook, "_get_conn_params", autospec=True)
+ def test_private_key_loader_retries_after_error(self, conn_params):
+ conn_params.return_value = {**CONN_PARAMS_WIF,
"workload_identity_provider": "OIDC", "token": "token"}
+ loader = mock.Mock(spec=SnowflakeHook().get_private_key,
side_effect=[OSError("unavailable"), None])
+ provider = SnowflakeRestTokenProvider(SnowflakeHook(),
private_key_loader=loader)
+ with pytest.raises(OSError, match="unavailable"):
+ provider.get_token()
+ assert provider.get_token().token == "WIF.OIDC.token"
+ assert loader.call_count == 2
+
+ @pytest.mark.parametrize("provider_name", ["GCP", "AZURE", "AWS", "gcp"])
Review Comment:
`provider_name` gets a lowercase `gcp` case here, but no test uses a
lowercase `authenticator`, so dropping the `.upper()` in `get_token` would go
unnoticed. The connector upper-cases it, so `workload_identity` works for
`SnowflakeHook`. Could this get a `workload_identity` case too?
##########
providers/snowflake/src/airflow/providers/snowflake/utils/_rest_auth.py:
##########
@@ -154,6 +181,61 @@ def get_token(self) -> SnowflakeRestToken:
token = self._jwt_generator.get_token()
return SnowflakeRestToken(token=token, token_type="KEYPAIR_JWT")
+ def _get_workload_identity_token(self, conn_config: dict[str, Any]) ->
SnowflakeRestToken:
+ """
+ Return a Workload Identity Federation token, fetching a new one when
the cached one expires.
+
+ The caller must hold ``self._lock``.
+ """
+ _validate_wif_host(conn_config)
+ if self._wif_token is not None and time.time() <
self._wif_token_expires_at:
+ return self._wif_token
+
+ provider_name = str(conn_config["workload_identity_provider"]).upper()
+ if provider_name == "AWS":
+ credential = _fetch_aws_identity_token()
Review Comment:
For deferred SQL API tasks this runs on the triggerer's event loop, since
`get_sql_api_query_status_async` builds headers synchronously. With the 300s
default on `GetWebIdentityToken` that's a blocking STS call every ~270s per
task. Can we pass `DurationSeconds=3600` here for now? Moving
`get_request_url_header_params` to `asyncio.to_thread` (like `cortex_model.py`)
is the real fix, but OAuth refresh has the same problem on main, so that one
can be a follow-up.
##########
providers/snowflake/docs/operators/cortex_agent.rst:
##########
@@ -51,9 +51,14 @@ Authentication
credentials grant, or ``azure_conn_id``, as on the Snowflake connection
generally.
- **PAT (Programmatic Access Token)**: set ``authenticator`` to
``programmatic_access_token`` and
put the PAT value in the connection's Password field.
-- **Key-pair JWT**: the default when neither of the above is set. Configure
+- **Key-pair JWT**: used when neither OAuth nor PAT is selected and a private
key is configured,
+ taking precedence over WIF. Configure
``private_key_file`` or ``private_key_content`` (optionally with a
passphrase in Password), and
set Login to the Snowflake user name the key is registered to.
+- **Workload Identity Federation (WIF)**: when no private key is configured,
set
Review Comment:
The `SnowflakeCortexAgentHook` docstring still says it authenticates "the
same three ways", with key-pair as "the default when neither of the above is
set", and `SnowflakeSqlApiHook.get_headers` still lists only OAuth, PAT and
JWT. Worth updating both to match this list?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
