arraylake 1.0.2__tar.gz → 1.0.3__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {arraylake-1.0.2 → arraylake-1.0.3}/.gitignore +1 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/PKG-INFO +1 -1
- arraylake-1.0.3/arraylake/_credential_cache.py +118 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/_version.py +2 -2
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/client.py +118 -24
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/types.py +5 -0
- arraylake-1.0.3/tests/test_credential_cache.py +254 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/uv.lock +44 -197
- {arraylake-1.0.2 → arraylake-1.0.3}/README.md +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/__init__.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/__main__.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/api_utils.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/asyn.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/__init__.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/auth.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/compute.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/config.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/main.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/repo.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/utils.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/compute/doctor.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/compute/http_client.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/compute/services.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/compute/types.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/config.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/config.yaml +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/credentials.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/diagnostics.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/display/__init__.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/display/repo.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/display/repolist.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/exceptions.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/log_util.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/metastore/__init__.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/metastore/abc.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/metastore/http_metastore.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/py.typed +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/repos/__init__.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/repos/icechunk/__init__.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/repos/icechunk/storage.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/repos/icechunk/types.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/repos/icechunk/virtual.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/token.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/tuning.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/pyproject.toml +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/__init__.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/compute/test_client.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/compute/test_doctor.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/config.yaml +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/conftest.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/icechunk/__init__.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/icechunk/conftest.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/icechunk/test_storage.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/icechunk/test_virtual.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_api_utils.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_asyn.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_cli.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_client.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_config.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_credentials.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_diagnostics.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_performance.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_subscriptions.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_token_handler.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_types.py +0 -0
- {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_virtual_chunks.py +0 -0
|
@@ -0,0 +1,118 @@
|
|
|
1
|
+
"""Process-wide cache for delegated bucket credentials.
|
|
2
|
+
|
|
3
|
+
When users pickle an icechunk Session and ship it to many workers, every
|
|
4
|
+
unpickled task can independently call the arraylake server for fresh
|
|
5
|
+
credentials when the embedded ones expire. This module gives the client a
|
|
6
|
+
single per-process cache so all consumers in a worker process share one
|
|
7
|
+
in-flight refresh.
|
|
8
|
+
|
|
9
|
+
Validity follows ``Cache-Control: max-age`` semantics: cached creds are
|
|
10
|
+
served while ``expires_at - CLOCK_SKEW_SAFETY > now``. The 30s safety only
|
|
11
|
+
protects against worker/server NTP-level clock drift; it is not a
|
|
12
|
+
pre-refresh window.
|
|
13
|
+
|
|
14
|
+
The cache is loop-agnostic: callers may live on any asyncio event loop in
|
|
15
|
+
the process. In-flight refreshes are coalesced via the singleflight
|
|
16
|
+
pattern using a single ``concurrent.futures.Future`` per key as the shared
|
|
17
|
+
result channel; each waiter bridges into its own loop via
|
|
18
|
+
``asyncio.wrap_future``. The race for who runs the fetcher is resolved by
|
|
19
|
+
``dict.setdefault``, which is atomic under the CPython GIL.
|
|
20
|
+
"""
|
|
21
|
+
|
|
22
|
+
from __future__ import annotations
|
|
23
|
+
|
|
24
|
+
import asyncio
|
|
25
|
+
import concurrent.futures
|
|
26
|
+
from collections.abc import Awaitable, Callable
|
|
27
|
+
from dataclasses import dataclass
|
|
28
|
+
from datetime import UTC, datetime, timedelta
|
|
29
|
+
from typing import Literal
|
|
30
|
+
|
|
31
|
+
from arraylake.log_util import get_logger
|
|
32
|
+
from arraylake.types import OrgName, Platform, TempCredentials
|
|
33
|
+
|
|
34
|
+
logger = get_logger(__name__)
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
CLOCK_SKEW_SAFETY = timedelta(seconds=30)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
@dataclass(eq=True, frozen=True)
|
|
41
|
+
class CredentialCacheKey:
|
|
42
|
+
api_url: str
|
|
43
|
+
auth_key: int
|
|
44
|
+
scope: Literal["repo", "bucket"]
|
|
45
|
+
org: OrgName
|
|
46
|
+
identifier: str
|
|
47
|
+
platform: Platform
|
|
48
|
+
access: Literal["read", "write"] = "read"
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
@dataclass
|
|
52
|
+
class _Entry:
|
|
53
|
+
creds: TempCredentials
|
|
54
|
+
expires_at: datetime | None
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
_CACHE: dict[CredentialCacheKey, _Entry] = {}
|
|
58
|
+
_INFLIGHT: dict[CredentialCacheKey, concurrent.futures.Future[TempCredentials]] = {}
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def _fresh(entry: _Entry | None, now: datetime) -> TempCredentials | None:
|
|
62
|
+
if entry is None or entry.expires_at is None:
|
|
63
|
+
return None
|
|
64
|
+
if entry.expires_at - CLOCK_SKEW_SAFETY > now:
|
|
65
|
+
return entry.creds
|
|
66
|
+
return None
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
async def get_or_refresh(
|
|
70
|
+
key: CredentialCacheKey,
|
|
71
|
+
fetcher: Callable[[], Awaitable[TempCredentials]],
|
|
72
|
+
*,
|
|
73
|
+
use_cache: bool = True,
|
|
74
|
+
) -> TempCredentials:
|
|
75
|
+
if not use_cache:
|
|
76
|
+
return await fetcher()
|
|
77
|
+
|
|
78
|
+
cached = _fresh(_CACHE.get(key), datetime.now(UTC))
|
|
79
|
+
if cached is not None:
|
|
80
|
+
return cached
|
|
81
|
+
|
|
82
|
+
my_future: concurrent.futures.Future[TempCredentials] = concurrent.futures.Future()
|
|
83
|
+
existing = _INFLIGHT.setdefault(key, my_future)
|
|
84
|
+
if existing is not my_future:
|
|
85
|
+
return await asyncio.wrap_future(existing)
|
|
86
|
+
|
|
87
|
+
try:
|
|
88
|
+
cached = _fresh(_CACHE.get(key), datetime.now(UTC))
|
|
89
|
+
if cached is not None:
|
|
90
|
+
my_future.set_result(cached)
|
|
91
|
+
return cached
|
|
92
|
+
|
|
93
|
+
creds = await fetcher()
|
|
94
|
+
exp = creds.expiration
|
|
95
|
+
if exp is not None and exp.tzinfo is None:
|
|
96
|
+
exp = exp.replace(tzinfo=UTC)
|
|
97
|
+
now = datetime.now(UTC)
|
|
98
|
+
if exp is not None and exp - CLOCK_SKEW_SAFETY <= now:
|
|
99
|
+
logger.warning(
|
|
100
|
+
"delegated credentials returned with lifetime <= clock-skew safety; server contract violation",
|
|
101
|
+
remaining_seconds=(exp - now).total_seconds(),
|
|
102
|
+
api_url=key.api_url,
|
|
103
|
+
org=key.org,
|
|
104
|
+
identifier=key.identifier,
|
|
105
|
+
)
|
|
106
|
+
_CACHE[key] = _Entry(creds=creds, expires_at=exp)
|
|
107
|
+
my_future.set_result(creds)
|
|
108
|
+
return creds
|
|
109
|
+
except BaseException as e:
|
|
110
|
+
my_future.set_exception(e)
|
|
111
|
+
raise
|
|
112
|
+
finally:
|
|
113
|
+
_INFLIGHT.pop(key, None)
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
def clear() -> None:
|
|
117
|
+
_CACHE.clear()
|
|
118
|
+
_INFLIGHT.clear()
|
|
@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
|
|
|
18
18
|
commit_id: str | None
|
|
19
19
|
__commit_id__: str | None
|
|
20
20
|
|
|
21
|
-
__version__ = version = '1.0.
|
|
22
|
-
__version_tuple__ = version_tuple = (1, 0,
|
|
21
|
+
__version__ = version = '1.0.3'
|
|
22
|
+
__version_tuple__ = version_tuple = (1, 0, 3)
|
|
23
23
|
|
|
24
24
|
__commit_id__ = commit_id = None
|
|
@@ -26,6 +26,7 @@ import icechunk
|
|
|
26
26
|
from icechunk import IcechunkError, RepositoryConfig
|
|
27
27
|
from icechunk import Repository as IcechunkRepository
|
|
28
28
|
|
|
29
|
+
from arraylake._credential_cache import CredentialCacheKey, get_or_refresh
|
|
29
30
|
from arraylake.asyn import async_gather_tasks, asyncio_run, sync
|
|
30
31
|
from arraylake.compute.services import AsyncComputeClient, ComputeClient
|
|
31
32
|
from arraylake.config import config as arraylake_config
|
|
@@ -114,12 +115,19 @@ class AsyncClient:
|
|
|
114
115
|
[Optional] The service URI to target.
|
|
115
116
|
token:
|
|
116
117
|
[Optional] API token for service account authentication.
|
|
118
|
+
cache_credentials:
|
|
119
|
+
[Optional] When True (default), delegated bucket credentials are
|
|
120
|
+
served from a process-wide cache so all consumers in a worker
|
|
121
|
+
share one in-flight refresh per ``(auth, target, access)``. Set
|
|
122
|
+
to False for server-side contexts that should not retain
|
|
123
|
+
credentials in memory — every call hits the server.
|
|
117
124
|
"""
|
|
118
125
|
|
|
119
126
|
_service_uri: str | None
|
|
120
127
|
token: str | None
|
|
128
|
+
_cache_credentials: bool
|
|
121
129
|
|
|
122
|
-
def __init__(self, service_uri: str | None = None, token: str | None = None) -> None:
|
|
130
|
+
def __init__(self, service_uri: str | None = None, token: str | None = None, cache_credentials: bool = True) -> None:
|
|
123
131
|
if service_uri is not None:
|
|
124
132
|
_validate_service_uri(service_uri)
|
|
125
133
|
|
|
@@ -133,6 +141,7 @@ class AsyncClient:
|
|
|
133
141
|
raise ValueError("Invalid token provided. Tokens must start with ema_ or be a JWT token.")
|
|
134
142
|
|
|
135
143
|
self.token = token
|
|
144
|
+
self._cache_credentials = cache_credentials
|
|
136
145
|
|
|
137
146
|
@property
|
|
138
147
|
def service_uri(self) -> str:
|
|
@@ -214,9 +223,22 @@ class AsyncClient:
|
|
|
214
223
|
Returns:
|
|
215
224
|
S3Credentials: Temporary credentials for the S3 bucket.
|
|
216
225
|
"""
|
|
217
|
-
|
|
218
|
-
|
|
219
|
-
|
|
226
|
+
key = CredentialCacheKey(
|
|
227
|
+
api_url=self.service_uri,
|
|
228
|
+
auth_key=hash(self.token),
|
|
229
|
+
scope="repo",
|
|
230
|
+
org=org,
|
|
231
|
+
identifier=repo_name,
|
|
232
|
+
platform="s3",
|
|
233
|
+
)
|
|
234
|
+
|
|
235
|
+
async def fetch() -> S3Credentials:
|
|
236
|
+
mstore = self._metastore_for_org(org)
|
|
237
|
+
return await mstore.get_s3_bucket_credentials_from_repo(repo_name)
|
|
238
|
+
|
|
239
|
+
creds = await get_or_refresh(key, fetch, use_cache=self._cache_credentials)
|
|
240
|
+
assert isinstance(creds, S3Credentials)
|
|
241
|
+
return creds
|
|
220
242
|
|
|
221
243
|
async def _get_gcs_delegated_credentials_from_repo(self, org: OrgName, repo_name: RepoName) -> GSCredentials:
|
|
222
244
|
"""Get delegated credentials for a repo's GCS bucket.
|
|
@@ -228,9 +250,22 @@ class AsyncClient:
|
|
|
228
250
|
Returns:
|
|
229
251
|
GSCredentials: Temporary credentials for the GCS bucket.
|
|
230
252
|
"""
|
|
231
|
-
|
|
232
|
-
|
|
233
|
-
|
|
253
|
+
key = CredentialCacheKey(
|
|
254
|
+
api_url=self.service_uri,
|
|
255
|
+
auth_key=hash(self.token),
|
|
256
|
+
scope="repo",
|
|
257
|
+
org=org,
|
|
258
|
+
identifier=repo_name,
|
|
259
|
+
platform="gs",
|
|
260
|
+
)
|
|
261
|
+
|
|
262
|
+
async def fetch() -> GSCredentials:
|
|
263
|
+
mstore = self._metastore_for_org(org)
|
|
264
|
+
return await mstore.get_gs_bucket_credentials_from_repo(repo_name)
|
|
265
|
+
|
|
266
|
+
creds = await get_or_refresh(key, fetch, use_cache=self._cache_credentials)
|
|
267
|
+
assert isinstance(creds, GSCredentials)
|
|
268
|
+
return creds
|
|
234
269
|
|
|
235
270
|
async def _get_azure_delegated_credentials_from_repo(self, org: OrgName, repo_name: RepoName) -> AzureCredentials:
|
|
236
271
|
"""Get delegated credentials for a repo's Azure Blob Storage container.
|
|
@@ -242,9 +277,22 @@ class AsyncClient:
|
|
|
242
277
|
Returns:
|
|
243
278
|
AzureCredentials: Temporary credentials for the Azure Blob Storage container.
|
|
244
279
|
"""
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
|
|
280
|
+
key = CredentialCacheKey(
|
|
281
|
+
api_url=self.service_uri,
|
|
282
|
+
auth_key=hash(self.token),
|
|
283
|
+
scope="repo",
|
|
284
|
+
org=org,
|
|
285
|
+
identifier=repo_name,
|
|
286
|
+
platform="azure",
|
|
287
|
+
)
|
|
288
|
+
|
|
289
|
+
async def fetch() -> AzureCredentials:
|
|
290
|
+
mstore = self._metastore_for_org(org)
|
|
291
|
+
return await mstore.get_azure_container_credentials_from_repo(repo_name)
|
|
292
|
+
|
|
293
|
+
creds = await get_or_refresh(key, fetch, use_cache=self._cache_credentials)
|
|
294
|
+
assert isinstance(creds, AzureCredentials)
|
|
295
|
+
return creds
|
|
248
296
|
|
|
249
297
|
async def _get_s3_delegated_credentials_from_bucket(
|
|
250
298
|
self, org: OrgName, nickname: BucketNickname, access: Literal["read", "write"] = "read"
|
|
@@ -262,10 +310,24 @@ class AsyncClient:
|
|
|
262
310
|
Returns:
|
|
263
311
|
S3Credentials: Temporary credentials for the S3 bucket.
|
|
264
312
|
"""
|
|
265
|
-
|
|
266
|
-
|
|
267
|
-
|
|
268
|
-
|
|
313
|
+
key = CredentialCacheKey(
|
|
314
|
+
api_url=self.service_uri,
|
|
315
|
+
auth_key=hash(self.token),
|
|
316
|
+
scope="bucket",
|
|
317
|
+
org=org,
|
|
318
|
+
identifier=nickname,
|
|
319
|
+
platform="s3",
|
|
320
|
+
access=access,
|
|
321
|
+
)
|
|
322
|
+
|
|
323
|
+
async def fetch() -> S3Credentials:
|
|
324
|
+
mstore = self._metastore_for_org(org)
|
|
325
|
+
bucket_id = await self._bucket_id_for_nickname(mstore, nickname)
|
|
326
|
+
return await mstore.get_s3_bucket_credentials_from_bucket(bucket_id, access=access)
|
|
327
|
+
|
|
328
|
+
creds = await get_or_refresh(key, fetch, use_cache=self._cache_credentials)
|
|
329
|
+
assert isinstance(creds, S3Credentials)
|
|
330
|
+
return creds
|
|
269
331
|
|
|
270
332
|
async def _get_gcs_delegated_credentials_from_bucket(
|
|
271
333
|
self, org: RepoName, nickname: BucketNickname, access: Literal["read", "write"] = "read"
|
|
@@ -283,10 +345,24 @@ class AsyncClient:
|
|
|
283
345
|
Returns:
|
|
284
346
|
GSCredentials: Temporary credentials for the GCS bucket.
|
|
285
347
|
"""
|
|
286
|
-
|
|
287
|
-
|
|
288
|
-
|
|
289
|
-
|
|
348
|
+
key = CredentialCacheKey(
|
|
349
|
+
api_url=self.service_uri,
|
|
350
|
+
auth_key=hash(self.token),
|
|
351
|
+
scope="bucket",
|
|
352
|
+
org=org,
|
|
353
|
+
identifier=nickname,
|
|
354
|
+
platform="gs",
|
|
355
|
+
access=access,
|
|
356
|
+
)
|
|
357
|
+
|
|
358
|
+
async def fetch() -> GSCredentials:
|
|
359
|
+
mstore = self._metastore_for_org(org)
|
|
360
|
+
bucket_id = await self._bucket_id_for_nickname(mstore, nickname)
|
|
361
|
+
return await mstore.get_gs_bucket_credentials_from_bucket(bucket_id, access=access)
|
|
362
|
+
|
|
363
|
+
creds = await get_or_refresh(key, fetch, use_cache=self._cache_credentials)
|
|
364
|
+
assert isinstance(creds, GSCredentials)
|
|
365
|
+
return creds
|
|
290
366
|
|
|
291
367
|
async def _get_azure_delegated_credentials_from_bucket(
|
|
292
368
|
self, org: RepoName, nickname: BucketNickname, access: Literal["read", "write"] = "read"
|
|
@@ -304,10 +380,24 @@ class AsyncClient:
|
|
|
304
380
|
Returns:
|
|
305
381
|
AzureCredentials: Temporary credentials for the Azure Blob Storage container.
|
|
306
382
|
"""
|
|
307
|
-
|
|
308
|
-
|
|
309
|
-
|
|
310
|
-
|
|
383
|
+
key = CredentialCacheKey(
|
|
384
|
+
api_url=self.service_uri,
|
|
385
|
+
auth_key=hash(self.token),
|
|
386
|
+
scope="bucket",
|
|
387
|
+
org=org,
|
|
388
|
+
identifier=nickname,
|
|
389
|
+
platform="azure",
|
|
390
|
+
access=access,
|
|
391
|
+
)
|
|
392
|
+
|
|
393
|
+
async def fetch() -> AzureCredentials:
|
|
394
|
+
mstore = self._metastore_for_org(org)
|
|
395
|
+
bucket_id = await self._bucket_id_for_nickname(mstore, nickname)
|
|
396
|
+
return await mstore.get_azure_container_credentials_from_bucket(bucket_id, access=access)
|
|
397
|
+
|
|
398
|
+
creds = await get_or_refresh(key, fetch, use_cache=self._cache_credentials)
|
|
399
|
+
assert isinstance(creds, AzureCredentials)
|
|
400
|
+
return creds
|
|
311
401
|
|
|
312
402
|
def _get_icechunk_s3_credentials_refresh_function_for_repo(self, org: OrgName, repo_name: RepoName) -> icechunk.S3StaticCredentials:
|
|
313
403
|
"""Get a function that returns S3 credentials for the given org and repo
|
|
@@ -1626,12 +1716,16 @@ class Client:
|
|
|
1626
1716
|
[Optional] The service URI to target.
|
|
1627
1717
|
token (str):
|
|
1628
1718
|
[Optional] API token for service account authentication.
|
|
1719
|
+
cache_credentials (bool):
|
|
1720
|
+
[Optional] When True (default), delegated bucket credentials are
|
|
1721
|
+
served from a process-wide cache. Set to False for server-side
|
|
1722
|
+
contexts that should not retain credentials in memory.
|
|
1629
1723
|
"""
|
|
1630
1724
|
|
|
1631
1725
|
aclient: AsyncClient
|
|
1632
1726
|
|
|
1633
|
-
def __init__(self, service_uri: str | None = None, token: str | None = None) -> None:
|
|
1634
|
-
self.aclient = AsyncClient(service_uri=service_uri, token=token)
|
|
1727
|
+
def __init__(self, service_uri: str | None = None, token: str | None = None, cache_credentials: bool = True) -> None:
|
|
1728
|
+
self.aclient = AsyncClient(service_uri=service_uri, token=token, cache_credentials=cache_credentials)
|
|
1635
1729
|
|
|
1636
1730
|
@property
|
|
1637
1731
|
def service_uri(self) -> str:
|
|
@@ -554,6 +554,7 @@ class BucketModel(BaseModel):
|
|
|
554
554
|
self,
|
|
555
555
|
credentials: TempCredentials | None,
|
|
556
556
|
prefix: str | None = None,
|
|
557
|
+
retry_config: Any | None = None,
|
|
557
558
|
) -> obs.store.ObjectStore:
|
|
558
559
|
"""Build an obstore ObjectStore for this bucket using the given credentials.
|
|
559
560
|
|
|
@@ -609,6 +610,7 @@ class BucketModel(BaseModel):
|
|
|
609
610
|
prefix=effective_prefix,
|
|
610
611
|
config=s3_config,
|
|
611
612
|
client_options=s3_client_options or None,
|
|
613
|
+
retry_config=retry_config,
|
|
612
614
|
)
|
|
613
615
|
elif self.platform == "gs":
|
|
614
616
|
if credentials is None:
|
|
@@ -616,6 +618,7 @@ class BucketModel(BaseModel):
|
|
|
616
618
|
bucket=self.name,
|
|
617
619
|
prefix=effective_prefix,
|
|
618
620
|
config={"skip_signature": True},
|
|
621
|
+
retry_config=retry_config,
|
|
619
622
|
)
|
|
620
623
|
elif isinstance(credentials, GSCredentials):
|
|
621
624
|
# google-auth returns naive datetimes (utcnow with tzinfo stripped) but
|
|
@@ -632,6 +635,7 @@ class BucketModel(BaseModel):
|
|
|
632
635
|
bucket=self.name,
|
|
633
636
|
prefix=effective_prefix,
|
|
634
637
|
credential_provider=gcs_credential_provider,
|
|
638
|
+
retry_config=retry_config,
|
|
635
639
|
)
|
|
636
640
|
else:
|
|
637
641
|
raise ValueError(f"Expected GSCredentials or None for platform 'gs', got {type(credentials).__name__}")
|
|
@@ -649,6 +653,7 @@ class BucketModel(BaseModel):
|
|
|
649
653
|
container_name=self.name,
|
|
650
654
|
prefix=effective_prefix,
|
|
651
655
|
config=azure_config,
|
|
656
|
+
retry_config=retry_config,
|
|
652
657
|
)
|
|
653
658
|
else:
|
|
654
659
|
raise ValueError(f"Unsupported platform: {self.platform}")
|
|
@@ -0,0 +1,254 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
import threading
|
|
5
|
+
from datetime import UTC, datetime, timedelta
|
|
6
|
+
from unittest.mock import AsyncMock, MagicMock, patch
|
|
7
|
+
|
|
8
|
+
import pytest
|
|
9
|
+
|
|
10
|
+
from arraylake import AsyncClient, _credential_cache
|
|
11
|
+
from arraylake._credential_cache import (
|
|
12
|
+
CLOCK_SKEW_SAFETY,
|
|
13
|
+
CredentialCacheKey,
|
|
14
|
+
clear,
|
|
15
|
+
get_or_refresh,
|
|
16
|
+
)
|
|
17
|
+
from arraylake.asyn import asyncio_run
|
|
18
|
+
from arraylake.types import GSCredentials, S3Credentials
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def _key(identifier: str = "repo-a", auth_key: int = 1) -> CredentialCacheKey:
|
|
22
|
+
return CredentialCacheKey(
|
|
23
|
+
api_url="https://api.test",
|
|
24
|
+
auth_key=auth_key,
|
|
25
|
+
scope="repo",
|
|
26
|
+
org="org-a",
|
|
27
|
+
identifier=identifier,
|
|
28
|
+
platform="s3",
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def _creds(remaining: timedelta | None = timedelta(hours=1)) -> S3Credentials:
|
|
33
|
+
return S3Credentials(
|
|
34
|
+
aws_access_key_id="AKIA",
|
|
35
|
+
aws_secret_access_key="secret",
|
|
36
|
+
aws_session_token="token",
|
|
37
|
+
expiration=datetime.now(UTC) + remaining if remaining is not None else None,
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
@pytest.fixture(autouse=True)
|
|
42
|
+
def _reset_cache():
|
|
43
|
+
clear()
|
|
44
|
+
yield
|
|
45
|
+
clear()
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def test_validity_boundary():
|
|
49
|
+
"""Cache freshness follows expires_at - CLOCK_SKEW_SAFETY. Each case runs
|
|
50
|
+
two back-to-back lookups; count is 1 (cache hit) or 2 (refetched)."""
|
|
51
|
+
cases = [
|
|
52
|
+
("hour-lifetime", timedelta(hours=1), 1),
|
|
53
|
+
("hard-expired", timedelta(seconds=-5), 2),
|
|
54
|
+
("inside-clock-skew", CLOCK_SKEW_SAFETY // 2, 2),
|
|
55
|
+
("just-past-clock-skew", CLOCK_SKEW_SAFETY + timedelta(seconds=10), 1),
|
|
56
|
+
("none-never-cached", None, 2),
|
|
57
|
+
]
|
|
58
|
+
for name, remaining, expected in cases:
|
|
59
|
+
clear()
|
|
60
|
+
calls = 0
|
|
61
|
+
|
|
62
|
+
async def fetch(_remaining=remaining) -> S3Credentials:
|
|
63
|
+
nonlocal calls
|
|
64
|
+
calls += 1
|
|
65
|
+
return _creds(remaining=_remaining)
|
|
66
|
+
|
|
67
|
+
async def body() -> None:
|
|
68
|
+
await get_or_refresh(_key(), fetch)
|
|
69
|
+
await get_or_refresh(_key(), fetch)
|
|
70
|
+
|
|
71
|
+
asyncio_run(body())
|
|
72
|
+
assert calls == expected, f"case {name}: expected {expected} fetches, got {calls}"
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
def test_keys_isolate_entries():
|
|
76
|
+
"""Two distinct cache keys hold independent entries."""
|
|
77
|
+
cases = [
|
|
78
|
+
("different-identifier", _key("repo-a"), _key("repo-b")),
|
|
79
|
+
("different-auth-key", _key(auth_key=1), _key(auth_key=2)),
|
|
80
|
+
]
|
|
81
|
+
for name, key_a, key_b in cases:
|
|
82
|
+
clear()
|
|
83
|
+
calls = 0
|
|
84
|
+
|
|
85
|
+
async def fetch() -> S3Credentials:
|
|
86
|
+
nonlocal calls
|
|
87
|
+
calls += 1
|
|
88
|
+
return _creds()
|
|
89
|
+
|
|
90
|
+
async def body() -> None:
|
|
91
|
+
await get_or_refresh(key_a, fetch)
|
|
92
|
+
await get_or_refresh(key_b, fetch)
|
|
93
|
+
|
|
94
|
+
asyncio_run(body())
|
|
95
|
+
assert calls == 2, f"case {name}: expected 2 fetches, got {calls}"
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def test_async_single_flight():
|
|
99
|
+
started = asyncio.Event()
|
|
100
|
+
release = asyncio.Event()
|
|
101
|
+
calls = 0
|
|
102
|
+
|
|
103
|
+
async def fetch() -> S3Credentials:
|
|
104
|
+
nonlocal calls
|
|
105
|
+
calls += 1
|
|
106
|
+
started.set()
|
|
107
|
+
await release.wait()
|
|
108
|
+
return _creds()
|
|
109
|
+
|
|
110
|
+
async def body() -> None:
|
|
111
|
+
tasks = [asyncio.create_task(get_or_refresh(_key(), fetch)) for _ in range(50)]
|
|
112
|
+
await started.wait()
|
|
113
|
+
release.set()
|
|
114
|
+
results = await asyncio.gather(*tasks)
|
|
115
|
+
for r in results:
|
|
116
|
+
assert results[0] is r
|
|
117
|
+
|
|
118
|
+
asyncio_run(body())
|
|
119
|
+
assert calls == 1
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
def test_cross_thread_single_flight():
|
|
123
|
+
release = threading.Event()
|
|
124
|
+
fetch_thread_arrived = threading.Event()
|
|
125
|
+
calls = 0
|
|
126
|
+
|
|
127
|
+
async def fetch() -> S3Credentials:
|
|
128
|
+
nonlocal calls
|
|
129
|
+
calls += 1
|
|
130
|
+
fetch_thread_arrived.set()
|
|
131
|
+
await asyncio.get_running_loop().run_in_executor(None, release.wait)
|
|
132
|
+
return _creds()
|
|
133
|
+
|
|
134
|
+
def worker() -> None:
|
|
135
|
+
asyncio_run(get_or_refresh(_key(), fetch))
|
|
136
|
+
|
|
137
|
+
threads = [threading.Thread(target=worker) for _ in range(30)]
|
|
138
|
+
for t in threads:
|
|
139
|
+
t.start()
|
|
140
|
+
assert fetch_thread_arrived.wait(timeout=5), "fetcher never started"
|
|
141
|
+
release.set()
|
|
142
|
+
for t in threads:
|
|
143
|
+
t.join(timeout=10)
|
|
144
|
+
assert calls == 1
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def test_works_from_foreign_loop():
|
|
148
|
+
"""Regression: previously get_or_refresh asserted that the caller was
|
|
149
|
+
running on the arraylake background loop, breaking any user who awaited
|
|
150
|
+
AsyncClient methods from their own asyncio loop. Calling from a fresh
|
|
151
|
+
foreign loop must complete normally."""
|
|
152
|
+
|
|
153
|
+
async def fetch() -> S3Credentials:
|
|
154
|
+
return _creds()
|
|
155
|
+
|
|
156
|
+
foreign_loop = asyncio.new_event_loop()
|
|
157
|
+
try:
|
|
158
|
+
creds = foreign_loop.run_until_complete(get_or_refresh(_key(), fetch))
|
|
159
|
+
finally:
|
|
160
|
+
foreign_loop.close()
|
|
161
|
+
assert creds.aws_access_key_id == "AKIA"
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
def test_fetcher_exception_does_not_poison_cache():
|
|
165
|
+
attempts = 0
|
|
166
|
+
|
|
167
|
+
async def flaky() -> S3Credentials:
|
|
168
|
+
nonlocal attempts
|
|
169
|
+
attempts += 1
|
|
170
|
+
if attempts == 1:
|
|
171
|
+
raise RuntimeError("transient")
|
|
172
|
+
return _creds()
|
|
173
|
+
|
|
174
|
+
async def body() -> None:
|
|
175
|
+
with pytest.raises(RuntimeError):
|
|
176
|
+
await get_or_refresh(_key(), flaky)
|
|
177
|
+
creds = await get_or_refresh(_key(), flaky)
|
|
178
|
+
assert creds.aws_access_key_id == "AKIA"
|
|
179
|
+
|
|
180
|
+
asyncio_run(body())
|
|
181
|
+
assert attempts == 2
|
|
182
|
+
|
|
183
|
+
|
|
184
|
+
def test_clear_drops_entries():
|
|
185
|
+
calls = 0
|
|
186
|
+
|
|
187
|
+
async def fetch() -> S3Credentials:
|
|
188
|
+
nonlocal calls
|
|
189
|
+
calls += 1
|
|
190
|
+
return _creds()
|
|
191
|
+
|
|
192
|
+
async def body() -> None:
|
|
193
|
+
await get_or_refresh(_key(), fetch)
|
|
194
|
+
assert _credential_cache._CACHE != {}
|
|
195
|
+
clear()
|
|
196
|
+
assert _credential_cache._CACHE == {}
|
|
197
|
+
await get_or_refresh(_key(), fetch)
|
|
198
|
+
|
|
199
|
+
asyncio_run(body())
|
|
200
|
+
assert calls == 2
|
|
201
|
+
|
|
202
|
+
|
|
203
|
+
def test_delegated_credentials_are_process_cached():
|
|
204
|
+
"""Two back-to-back calls to the same delegated-credential fetcher must
|
|
205
|
+
only hit the metastore once. This is the core guarantee for the
|
|
206
|
+
pickled-icechunk-session-across-dask-workers scenario."""
|
|
207
|
+
expiry = datetime.now(UTC) + timedelta(hours=1)
|
|
208
|
+
s3_creds = S3Credentials(
|
|
209
|
+
aws_access_key_id="AKIA",
|
|
210
|
+
aws_secret_access_key="secret",
|
|
211
|
+
aws_session_token="token",
|
|
212
|
+
expiration=expiry,
|
|
213
|
+
)
|
|
214
|
+
gs_creds = GSCredentials(access_token="bearer", principal="sa@x", expiration=expiry)
|
|
215
|
+
aclient = AsyncClient("https://test.example", token="ema_fake_token")
|
|
216
|
+
aclient_no_cache = AsyncClient("https://test.example", token="ema_fake_token", cache_credentials=False)
|
|
217
|
+
|
|
218
|
+
cases = [
|
|
219
|
+
("s3-repo", "s3", "repo", aclient, "_get_s3_delegated_credentials_from_repo", ("org-a", "repo-a"), 1),
|
|
220
|
+
("gs-repo", "gs", "repo", aclient, "_get_gcs_delegated_credentials_from_repo", ("org-a", "repo-a"), 1),
|
|
221
|
+
("s3-bucket", "s3", "bucket", aclient, "_get_s3_delegated_credentials_from_bucket", ("org-a", "bucket-a"), 1),
|
|
222
|
+
("gs-bucket", "gs", "bucket", aclient, "_get_gcs_delegated_credentials_from_bucket", ("org-a", "bucket-a"), 1),
|
|
223
|
+
("s3-repo-nocache", "s3", "repo", aclient_no_cache, "_get_s3_delegated_credentials_from_repo", ("org-a", "repo-a"), 2),
|
|
224
|
+
("gs-bucket-nocache", "gs", "bucket", aclient_no_cache, "_get_gcs_delegated_credentials_from_bucket", ("org-a", "bucket-a"), 2),
|
|
225
|
+
]
|
|
226
|
+
for name, platform, scope, client, method_name, args, expected_calls in cases:
|
|
227
|
+
clear()
|
|
228
|
+
creds = s3_creds if platform == "s3" else gs_creds
|
|
229
|
+
mstore = MagicMock()
|
|
230
|
+
mstore.get_s3_bucket_credentials_from_repo = AsyncMock(return_value=creds)
|
|
231
|
+
mstore.get_gs_bucket_credentials_from_repo = AsyncMock(return_value=creds)
|
|
232
|
+
mstore.get_s3_bucket_credentials_from_bucket = AsyncMock(return_value=creds)
|
|
233
|
+
mstore.get_gs_bucket_credentials_from_bucket = AsyncMock(return_value=creds)
|
|
234
|
+
|
|
235
|
+
method = getattr(client, method_name)
|
|
236
|
+
|
|
237
|
+
async def invoke(_method=method, _args=args) -> None:
|
|
238
|
+
await _method(*_args)
|
|
239
|
+
await _method(*_args)
|
|
240
|
+
|
|
241
|
+
with (
|
|
242
|
+
patch.object(AsyncClient, "_metastore_for_org", return_value=mstore),
|
|
243
|
+
patch.object(AsyncClient, "_bucket_id_for_nickname", new=AsyncMock(return_value="bucket-uuid")),
|
|
244
|
+
):
|
|
245
|
+
asyncio_run(invoke())
|
|
246
|
+
|
|
247
|
+
mstore_method = f"get_{'gs' if platform == 'gs' else 's3'}_bucket_credentials_from_{scope}"
|
|
248
|
+
assert getattr(mstore, mstore_method).await_count == expected_calls, (
|
|
249
|
+
f"case {name}: expected {expected_calls} metastore call(s), got {getattr(mstore, mstore_method).await_count}"
|
|
250
|
+
)
|
|
251
|
+
if expected_calls == 1:
|
|
252
|
+
assert _credential_cache._CACHE != {}, f"case {name}: expected cache to be populated"
|
|
253
|
+
else:
|
|
254
|
+
assert _credential_cache._CACHE == {}, f"case {name}: expected cache to remain empty"
|