feat: implement vended credential refresh - #3751
Conversation
350dc82 to
010d6bc
Compare
|
@kevinjqliu I opened this PR as the previous one was closed for being stale. Here I've hopefully addressed your comments!
|
|
+1 (non-binding). GCS support would be pretty to add to this. |
| catalog=self, | ||
| ) | ||
|
|
||
| def _attach_credentials_provider( |
There was a problem hiding this comment.
Can we rename this to _attach_credentials_provider_to_io?
| longest-prefix matching against the freshly vended credentials. | ||
| """ | ||
| if storage_credentials and property_as_bool(self.properties, REFRESH_CREDENTIALS_ENABLED, False): | ||
| io.set_credentials_provider( |
There was a problem hiding this comment.
This breaks pickling support on FileIO. We can't pickle lambdas.
We have some tests for pickling FileIO, but none of them break because this is behind REFRESH_CREDENTIALS_ENABLE
There was a problem hiding this comment.
Good catch, this is a real blocker for this. If we want to keep the concern of refreshing the credentials by calling the catalog.load_credentials then we cannot reliable pickle this object. Specially this we would need to pickle AuthManagers and its locks and it would become a very fragile and prone to break.
I bellieve we should re-visit allowing the credential provider make the API call itself so there is no reference to any lambdas or catalog object dependency
|
|
||
| def _can_refresh(self, location: str) -> bool: | ||
| scheme = urlparse(location).scheme | ||
| refresh_by_scheme = self._needs_refresh_by_scheme.get(scheme) |
There was a problem hiding this comment.
If the TTL on a S3 token is < 300s (the default value), we'll attempt to refresh on every call. That's a lot of possibly unnecessary refreshes.
It seems like we should track when the last refresh was and use that information to determine when we should next refresh.
Java does something similar in scheduleCredentialRefresh
There was a problem hiding this comment.
IIUC your comment, we wont refresh on every call, since the lock will hold until the refresh completes and all other threads will not reach the load_credentials
| The refresh callback returns the full LoadCredentialsResponse so the provider can re-run | ||
| longest-prefix matching against the freshly vended credentials. | ||
| """ | ||
| if storage_credentials and property_as_bool(self.properties, REFRESH_CREDENTIALS_ENABLED, False): |
There was a problem hiding this comment.
We should check that the server has access to the LoadCredentials endpoint. Otherwise, it'll probably throw an error.
(If you're setting this property, your catalog probably has the proper endpoints, but we should verify anyways)
There was a problem hiding this comment.
the self._load_credentials() method performs the check
| return best_match.config if best_match else {} | ||
|
|
||
|
|
||
| class CredentialsProvider: |
There was a problem hiding this comment.
I think this class ultimately needs to less S3-specific and delegate some functionality into scheme-specific handlers. I don't think that has to be part of this PR though, it could be part of a following that adds support for other schemes, such as GCS.
There was a problem hiding this comment.
Overall, I agree. I think we'll have a better idea of the shape of this class once we implement GCS.
|
Thanks for the reviews! I will make the changes this week been AFK |
3 participants
Third time's a charm (prev PR closed due to staleness) :) Closes #3506
Rationale for this change
Implement automatic vended credential refresh in PyIceberg.
Are these changes tested?
Yes,
Are there any user-facing changes?
Yes, users will have the option to set the refresh credentials configurations in the catalog