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.
Files changed (66) hide show
  1. {arraylake-1.0.2 → arraylake-1.0.3}/.gitignore +1 -0
  2. {arraylake-1.0.2 → arraylake-1.0.3}/PKG-INFO +1 -1
  3. arraylake-1.0.3/arraylake/_credential_cache.py +118 -0
  4. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/_version.py +2 -2
  5. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/client.py +118 -24
  6. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/types.py +5 -0
  7. arraylake-1.0.3/tests/test_credential_cache.py +254 -0
  8. {arraylake-1.0.2 → arraylake-1.0.3}/uv.lock +44 -197
  9. {arraylake-1.0.2 → arraylake-1.0.3}/README.md +0 -0
  10. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/__init__.py +0 -0
  11. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/__main__.py +0 -0
  12. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/api_utils.py +0 -0
  13. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/asyn.py +0 -0
  14. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/__init__.py +0 -0
  15. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/auth.py +0 -0
  16. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/compute.py +0 -0
  17. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/config.py +0 -0
  18. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/main.py +0 -0
  19. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/repo.py +0 -0
  20. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/cli/utils.py +0 -0
  21. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/compute/doctor.py +0 -0
  22. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/compute/http_client.py +0 -0
  23. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/compute/services.py +0 -0
  24. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/compute/types.py +0 -0
  25. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/config.py +0 -0
  26. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/config.yaml +0 -0
  27. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/credentials.py +0 -0
  28. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/diagnostics.py +0 -0
  29. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/display/__init__.py +0 -0
  30. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/display/repo.py +0 -0
  31. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/display/repolist.py +0 -0
  32. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/exceptions.py +0 -0
  33. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/log_util.py +0 -0
  34. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/metastore/__init__.py +0 -0
  35. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/metastore/abc.py +0 -0
  36. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/metastore/http_metastore.py +0 -0
  37. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/py.typed +0 -0
  38. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/repos/__init__.py +0 -0
  39. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/repos/icechunk/__init__.py +0 -0
  40. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/repos/icechunk/storage.py +0 -0
  41. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/repos/icechunk/types.py +0 -0
  42. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/repos/icechunk/virtual.py +0 -0
  43. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/token.py +0 -0
  44. {arraylake-1.0.2 → arraylake-1.0.3}/arraylake/tuning.py +0 -0
  45. {arraylake-1.0.2 → arraylake-1.0.3}/pyproject.toml +0 -0
  46. {arraylake-1.0.2 → arraylake-1.0.3}/tests/__init__.py +0 -0
  47. {arraylake-1.0.2 → arraylake-1.0.3}/tests/compute/test_client.py +0 -0
  48. {arraylake-1.0.2 → arraylake-1.0.3}/tests/compute/test_doctor.py +0 -0
  49. {arraylake-1.0.2 → arraylake-1.0.3}/tests/config.yaml +0 -0
  50. {arraylake-1.0.2 → arraylake-1.0.3}/tests/conftest.py +0 -0
  51. {arraylake-1.0.2 → arraylake-1.0.3}/tests/icechunk/__init__.py +0 -0
  52. {arraylake-1.0.2 → arraylake-1.0.3}/tests/icechunk/conftest.py +0 -0
  53. {arraylake-1.0.2 → arraylake-1.0.3}/tests/icechunk/test_storage.py +0 -0
  54. {arraylake-1.0.2 → arraylake-1.0.3}/tests/icechunk/test_virtual.py +0 -0
  55. {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_api_utils.py +0 -0
  56. {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_asyn.py +0 -0
  57. {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_cli.py +0 -0
  58. {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_client.py +0 -0
  59. {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_config.py +0 -0
  60. {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_credentials.py +0 -0
  61. {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_diagnostics.py +0 -0
  62. {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_performance.py +0 -0
  63. {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_subscriptions.py +0 -0
  64. {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_token_handler.py +0 -0
  65. {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_types.py +0 -0
  66. {arraylake-1.0.2 → arraylake-1.0.3}/tests/test_virtual_chunks.py +0 -0
@@ -149,3 +149,4 @@ client/arraylake/_version.py
149
149
  !/.claude/
150
150
  /.claude/*
151
151
  !/.claude/skills/
152
+ .local/
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: arraylake
3
- Version: 1.0.2
3
+ Version: 1.0.3
4
4
  Summary: Python client for ArrayLake
5
5
  Project-URL: Documentation, https://docs.earthmover.io/
6
6
  Project-URL: Changelog, https://docs.earthmover.io/changelog
@@ -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.2'
22
- __version_tuple__ = version_tuple = (1, 0, 2)
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
- mstore = self._metastore_for_org(org)
218
- s3_creds = await mstore.get_s3_bucket_credentials_from_repo(repo_name)
219
- return s3_creds
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
- mstore = self._metastore_for_org(org)
232
- gcs_creds = await mstore.get_gs_bucket_credentials_from_repo(repo_name)
233
- return gcs_creds
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
- mstore = self._metastore_for_org(org)
246
- azure_creds = await mstore.get_azure_container_credentials_from_repo(repo_name)
247
- return azure_creds
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
- mstore = self._metastore_for_org(org)
266
- bucket_id = await self._bucket_id_for_nickname(mstore, nickname)
267
- s3_creds = await mstore.get_s3_bucket_credentials_from_bucket(bucket_id, access=access)
268
- return s3_creds
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
- mstore = self._metastore_for_org(org)
287
- bucket_id = await self._bucket_id_for_nickname(mstore, nickname)
288
- gcs_creds = await mstore.get_gs_bucket_credentials_from_bucket(bucket_id, access=access)
289
- return gcs_creds
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
- mstore = self._metastore_for_org(org)
308
- bucket_id = await self._bucket_id_for_nickname(mstore, nickname)
309
- azure_creds = await mstore.get_azure_container_credentials_from_bucket(bucket_id, access=access)
310
- return azure_creds
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"