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
There was a problem hiding this comment.
What if we just drop the lock when pickling? When the object is unpickled, it's a brand new instance anyways. It shouldn't have the same lock to begin with.
There was a problem hiding this comment.
Dropping and recreating the lock makes sense. My concern is that the refresh lambda still captures the catalog instance, and the catalog contains additional state like sessions, auth managers, and possibly other locks that may not be pickleable.
So we’d still need to make the refresh call independent of the catalog or reconstruct it after unpickling/
|
|
||
| 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