Credential refresh can be initialized with the first set of credentials (#921)

`s3_storage` style functions got a new argument:

```
resolve_get_credentials: bool = False
```

If True is passed, pickled repos and sessions will use the result
returned the first time `get_credentials` was called, instead of having
to run it again.

This speeds up short distributed tasks, since they no longer need to
spend time fetching credentials. Tasks will reuse the credentials
gathered by the "coordinator"

Author: paraseba

Cargo.lock

@@ -39,6 +39,7 @@ typetag = "0.2.20"
 serde = { version = "1.0.219", features = ["derive", "rc"] }
 miette = { version = "7.5.0", features = ["fancy"] }
 clap = { version = "4.5", features = ["derive"], optional = true }
+rand = "0.9.0"
 
 [features]
 cli = ["clap", "icechunk/cli"]

icechunk-python/Cargo.toml

@@ -1260,8 +1260,13 @@ class S3Credentials:
         ----------
         pickled_function: bytes
             The pickled function to use to provide credentials.
+        current: S3StaticCredentials
+            The initial credentials. They will be returned the first time credentials
+            are requested and then deleted.
         """
-        def __init__(self, pickled_function: bytes) -> None: ...
+        def __init__(
+            self, pickled_function: bytes, current: S3StaticCredentials | None = None
+        ) -> None: ...
 
 AnyS3Credential = (
     S3Credentials.Static
@@ -1359,7 +1364,9 @@ class GcsCredentials:
 
         This is useful for credentials that have an expiration time, or are otherwise not known ahead of time.
         """
-        def __init__(self, pickled_function: bytes) -> None: ...
+        def __init__(
+            self, pickled_function: bytes, current: GcsBearerCredential | None = None
+        ) -> None: ...
 
 AnyGcsCredential = (
     GcsCredentials.FromEnv | GcsCredentials.Static | GcsCredentials.Refreshable

icechunk-python/python/icechunk/_icechunk_python.pyi

@@ -45,16 +45,23 @@
 
 def s3_refreshable_credentials(
     get_credentials: Callable[[], S3StaticCredentials],
+    scatter_initial_credentials: bool = False,
 ) -> S3Credentials.Refreshable:
     """Create refreshable credentials for S3 and S3 compatible object stores.
 
-
     Parameters
     ----------
     get_credentials: Callable[[], S3StaticCredentials]
         Use this function to get and refresh the credentials. The function must be pickable.
+    scatter_initial_credentials: bool, optional
+        Immediately call and store the value returned by get_credentials. This is useful if the
+        repo or session will be pickled to generate many copies. Passing scatter_initial_credentials=True will
+        ensure all those copies don't need to call get_credentials immediately. After the initial
+        set of credentials has expired, the cached value is no longer used. Notice that credentials
+        obtained are stored, and they can be sent over the network if you pickle the session/repo.
     """
-    return S3Credentials.Refreshable(pickle.dumps(get_credentials))
+    current = get_credentials() if scatter_initial_credentials else None
+    return S3Credentials.Refreshable(pickle.dumps(get_credentials), current)
 
 
 def s3_static_credentials(
@@ -106,6 +113,7 @@ def s3_credentials(
     anonymous: bool | None = None,
     from_env: bool | None = None,
     get_credentials: Callable[[], S3StaticCredentials] | None = None,
+    scatter_initial_credentials: bool = False,
 ) -> AnyS3Credential:
     """Create credentials for S3 and S3 compatible object stores.
 
@@ -127,6 +135,12 @@ def s3_credentials(
         Fetch credentials from the operative system environment
     get_credentials: Callable[[], S3StaticCredentials] | None
         Use this function to get and refresh object store credentials
+    scatter_initial_credentials: bool, optional
+        Immediately call and store the value returned by get_credentials. This is useful if the
+        repo or session will be pickled to generate many copies. Passing scatter_initial_credentials=True will
+        ensure all those copies don't need to call get_credentials immediately. After the initial
+        set of credentials has expired, the cached value is no longer used. Notice that credentials
+        obtained are stored, and they can be sent over the network if you pickle the session/repo.
     """
     if (
         (from_env is None or from_env)
@@ -159,7 +173,9 @@ def s3_credentials(
         and not from_env
         and not anonymous
     ):
-        return s3_refreshable_credentials(get_credentials)
+        return s3_refreshable_credentials(
+            get_credentials, scatter_initial_credentials=scatter_initial_credentials
+        )
 
     if (
         access_key_id
@@ -199,9 +215,24 @@ def gcs_static_credentials(
 
 def gcs_refreshable_credentials(
     get_credentials: Callable[[], GcsBearerCredential],
+    scatter_initial_credentials: bool = False,
 ) -> GcsCredentials.Refreshable:
-    """Create refreshable credentials for Google Cloud Storage object store."""
-    return GcsCredentials.Refreshable(pickle.dumps(get_credentials))
+    """Create refreshable credentials for Google Cloud Storage object store.
+
+    Parameters
+    ----------
+    get_credentials: Callable[[], S3StaticCredentials]
+        Use this function to get and refresh the credentials. The function must be pickable.
+    scatter_initial_credentials: bool, optional
+        Immediately call and store the value returned by get_credentials. This is useful if the
+        repo or session will be pickled to generate many copies. Passing scatter_initial_credentials=True will
+        ensure all those copies don't need to call get_credentials immediately. After the initial
+        set of credentials has expired, the cached value is no longer used. Notice that credentials
+        obtained are stored, and they can be sent over the network if you pickle the session/repo.
+    """
+
+    current = get_credentials() if scatter_initial_credentials else None
+    return GcsCredentials.Refreshable(pickle.dumps(get_credentials), current)
 
 
 def gcs_from_env_credentials() -> GcsCredentials.FromEnv:
@@ -217,6 +248,7 @@ def gcs_credentials(
     bearer_token: str | None = None,
     from_env: bool | None = None,
     get_credentials: Callable[[], GcsBearerCredential] | None = None,
+    scatter_initial_credentials: bool = False,
 ) -> AnyGcsCredential:
     """Create credentials Google Cloud Storage object store.
 
@@ -246,7 +278,9 @@ def gcs_credentials(
         )
 
     if get_credentials is not None:
-        return gcs_refreshable_credentials(get_credentials)
+        return gcs_refreshable_credentials(
+            get_credentials, scatter_initial_credentials=scatter_initial_credentials
+        )
 
     raise ValueError("Conflicting arguments to gcs_credentials function")
 

icechunk-python/python/icechunk/credentials.py

@@ -49,6 +49,7 @@ def s3_store(
     force_path_style: bool = False,
 ) -> ObjectStoreConfig.S3Compatible | ObjectStoreConfig.S3:
     """Build an ObjectStoreConfig instance for S3 or S3 compatible object stores."""
+
     options = S3Options(
         region=region,
         endpoint_url=endpoint_url,
@@ -76,6 +77,7 @@ def s3_storage(
     anonymous: bool | None = None,
     from_env: bool | None = None,
     get_credentials: Callable[[], S3StaticCredentials] | None = None,
+    scatter_initial_credentials: bool = False,
     force_path_style: bool = False,
 ) -> Storage:
     """Create a Storage instance that saves data in S3 or S3 compatible object stores.
@@ -106,6 +108,12 @@ def s3_storage(
         Fetch credentials from the operative system environment
     get_credentials: Callable[[], S3StaticCredentials] | None
         Use this function to get and refresh object store credentials
+    scatter_initial_credentials: bool, optional
+        Immediately call and store the value returned by get_credentials. This is useful if the
+        repo or session will be pickled to generate many copies. Passing scatter_initial_credentials=True will
+        ensure all those copies don't need to call get_credentials immediately. After the initial
+        set of credentials has expired, the cached value is no longer used. Notice that credentials
+        obtained are stored, and they can be sent over the network if you pickle the session/repo.
     force_path_style: bool
         Whether to force using path-style addressing for buckets
     """
@@ -118,6 +126,7 @@ def s3_storage(
         anonymous=anonymous,
         from_env=from_env,
         get_credentials=get_credentials,
+        scatter_initial_credentials=scatter_initial_credentials,
     )
     options = S3Options(
         region=region,
@@ -186,6 +195,7 @@ def tigris_storage(
     anonymous: bool | None = None,
     from_env: bool | None = None,
     get_credentials: Callable[[], S3StaticCredentials] | None = None,
+    scatter_initial_credentials: bool = False,
 ) -> Storage:
     """Create a Storage instance that saves data in Tigris object store.
 
@@ -219,6 +229,12 @@ def tigris_storage(
         Fetch credentials from the operative system environment
     get_credentials: Callable[[], S3StaticCredentials] | None
         Use this function to get and refresh object store credentials
+    scatter_initial_credentials: bool, optional
+        Immediately call and store the value returned by get_credentials. This is useful if the
+        repo or session will be pickled to generate many copies. Passing scatter_initial_credentials=True will
+        ensure all those copies don't need to call get_credentials immediately. After the initial
+        set of credentials has expired, the cached value is no longer used. Notice that credentials
+        obtained are stored, and they can be sent over the network if you pickle the session/repo.
     """
     credentials = s3_credentials(
         access_key_id=access_key_id,
@@ -228,6 +244,7 @@ def tigris_storage(
         anonymous=anonymous,
         from_env=from_env,
         get_credentials=get_credentials,
+        scatter_initial_credentials=scatter_initial_credentials,
     )
     options = S3Options(region=region, endpoint_url=endpoint_url, allow_http=allow_http)
     return Storage.new_tigris(
@@ -254,6 +271,7 @@ def r2_storage(
     anonymous: bool | None = None,
     from_env: bool | None = None,
     get_credentials: Callable[[], S3StaticCredentials] | None = None,
+    scatter_initial_credentials: bool = False,
 ) -> Storage:
     """Create a Storage instance that saves data in Tigris object store.
 
@@ -287,6 +305,12 @@ def r2_storage(
         Fetch credentials from the operative system environment
     get_credentials: Callable[[], S3StaticCredentials] | None
         Use this function to get and refresh object store credentials
+    scatter_initial_credentials: bool, optional
+        Immediately call and store the value returned by get_credentials. This is useful if the
+        repo or session will be pickled to generate many copies. Passing scatter_initial_credentials=True will
+        ensure all those copies don't need to call get_credentials immediately. After the initial
+        set of credentials has expired, the cached value is no longer used. Notice that credentials
+        obtained are stored, and they can be sent over the network if you pickle the session/repo.
     """
     credentials = s3_credentials(
         access_key_id=access_key_id,
@@ -296,6 +320,7 @@ def r2_storage(
         anonymous=anonymous,
         from_env=from_env,
         get_credentials=get_credentials,
+        scatter_initial_credentials=scatter_initial_credentials,
     )
     options = S3Options(region=region, endpoint_url=endpoint_url, allow_http=allow_http)
     return Storage.new_r2(
@@ -318,6 +343,7 @@ def gcs_storage(
     from_env: bool | None = None,
     config: dict[str, str] | None = None,
     get_credentials: Callable[[], GcsBearerCredential] | None = None,
+    scatter_initial_credentials: bool = False,
 ) -> Storage:
     """Create a Storage instance that saves data in Google Cloud Storage object store.
 
@@ -333,6 +359,12 @@ def gcs_storage(
         The bearer token to use for the object store
     get_credentials: Callable[[], GcsBearerCredential] | None
         Use this function to get and refresh object store credentials
+    scatter_initial_credentials: bool, optional
+        Immediately call and store the value returned by get_credentials. This is useful if the
+        repo or session will be pickled to generate many copies. Passing scatter_initial_credentials=True will
+        ensure all those copies don't need to call get_credentials immediately. After the initial
+        set of credentials has expired, the cached value is no longer used. Notice that credentials
+        obtained are stored, and they can be sent over the network if you pickle the session/repo.
     """
     credentials = gcs_credentials(
         service_account_file=service_account_file,
@@ -341,6 +373,7 @@ def gcs_storage(
         bearer_token=bearer_token,
         from_env=from_env,
         get_credentials=get_credentials,
+        scatter_initial_credentials=scatter_initial_credentials,
     )
     return Storage.new_gcs(
         bucket=bucket,