arraylake 1.0.3__tar.gz → 1.1.0__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 (69) hide show
  1. {arraylake-1.0.3 → arraylake-1.1.0}/PKG-INFO +1 -1
  2. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/_credential_cache.py +10 -1
  3. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/_version.py +2 -2
  4. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/main.py +2 -0
  5. arraylake-1.1.0/arraylake/cli/skills.py +208 -0
  6. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/client.py +187 -70
  7. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/compute/types.py +9 -1
  8. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/repos/icechunk/types.py +1 -0
  9. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/repos/icechunk/virtual.py +3 -3
  10. {arraylake-1.0.3 → arraylake-1.1.0}/pyproject.toml +2 -2
  11. {arraylake-1.0.3 → arraylake-1.1.0}/tests/conftest.py +10 -0
  12. {arraylake-1.0.3 → arraylake-1.1.0}/tests/icechunk/test_virtual.py +15 -0
  13. {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_credential_cache.py +104 -8
  14. arraylake-1.1.0/tests/test_skills.py +111 -0
  15. {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_virtual_chunks.py +96 -0
  16. arraylake-1.1.0/uv.lock +2802 -0
  17. arraylake-1.0.3/uv.lock +0 -2795
  18. {arraylake-1.0.3 → arraylake-1.1.0}/.gitignore +0 -0
  19. {arraylake-1.0.3 → arraylake-1.1.0}/README.md +0 -0
  20. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/__init__.py +0 -0
  21. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/__main__.py +0 -0
  22. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/api_utils.py +0 -0
  23. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/asyn.py +0 -0
  24. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/__init__.py +0 -0
  25. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/auth.py +0 -0
  26. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/compute.py +0 -0
  27. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/config.py +0 -0
  28. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/repo.py +0 -0
  29. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/utils.py +0 -0
  30. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/compute/doctor.py +0 -0
  31. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/compute/http_client.py +0 -0
  32. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/compute/services.py +0 -0
  33. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/config.py +0 -0
  34. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/config.yaml +0 -0
  35. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/credentials.py +0 -0
  36. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/diagnostics.py +0 -0
  37. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/display/__init__.py +0 -0
  38. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/display/repo.py +0 -0
  39. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/display/repolist.py +0 -0
  40. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/exceptions.py +0 -0
  41. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/log_util.py +0 -0
  42. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/metastore/__init__.py +0 -0
  43. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/metastore/abc.py +0 -0
  44. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/metastore/http_metastore.py +0 -0
  45. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/py.typed +0 -0
  46. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/repos/__init__.py +0 -0
  47. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/repos/icechunk/__init__.py +0 -0
  48. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/repos/icechunk/storage.py +0 -0
  49. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/token.py +0 -0
  50. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/tuning.py +0 -0
  51. {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/types.py +0 -0
  52. {arraylake-1.0.3 → arraylake-1.1.0}/tests/__init__.py +0 -0
  53. {arraylake-1.0.3 → arraylake-1.1.0}/tests/compute/test_client.py +0 -0
  54. {arraylake-1.0.3 → arraylake-1.1.0}/tests/compute/test_doctor.py +0 -0
  55. {arraylake-1.0.3 → arraylake-1.1.0}/tests/config.yaml +0 -0
  56. {arraylake-1.0.3 → arraylake-1.1.0}/tests/icechunk/__init__.py +0 -0
  57. {arraylake-1.0.3 → arraylake-1.1.0}/tests/icechunk/conftest.py +0 -0
  58. {arraylake-1.0.3 → arraylake-1.1.0}/tests/icechunk/test_storage.py +0 -0
  59. {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_api_utils.py +0 -0
  60. {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_asyn.py +0 -0
  61. {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_cli.py +0 -0
  62. {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_client.py +0 -0
  63. {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_config.py +0 -0
  64. {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_credentials.py +0 -0
  65. {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_diagnostics.py +0 -0
  66. {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_performance.py +0 -0
  67. {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_subscriptions.py +0 -0
  68. {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_token_handler.py +0 -0
  69. {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_types.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: arraylake
3
- Version: 1.0.3
3
+ Version: 1.1.0
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
@@ -41,7 +41,7 @@ CLOCK_SKEW_SAFETY = timedelta(seconds=30)
41
41
  class CredentialCacheKey:
42
42
  api_url: str
43
43
  auth_key: int
44
- scope: Literal["repo", "bucket"]
44
+ scope: Literal["repo", "bucket", "vcc"]
45
45
  org: OrgName
46
46
  identifier: str
47
47
  platform: Platform
@@ -113,6 +113,15 @@ async def get_or_refresh(
113
113
  _INFLIGHT.pop(key, None)
114
114
 
115
115
 
116
+ def populate(key: CredentialCacheKey, creds: TempCredentials) -> None:
117
+ exp = creds.expiration
118
+ if exp is not None and exp.tzinfo is None:
119
+ exp = exp.replace(tzinfo=UTC)
120
+ if _fresh(_CACHE.get(key), datetime.now(UTC)) is not None:
121
+ return
122
+ _CACHE[key] = _Entry(creds=creds, expires_at=exp)
123
+
124
+
116
125
  def clear() -> None:
117
126
  _CACHE.clear()
118
127
  _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.3'
22
- __version_tuple__ = version_tuple = (1, 0, 3)
21
+ __version__ = version = '1.1.0'
22
+ __version_tuple__ = version_tuple = (1, 1, 0)
23
23
 
24
24
  __commit_id__ = commit_id = None
@@ -6,6 +6,7 @@ from arraylake.cli.auth import auth
6
6
  from arraylake.cli.compute import app as compute_app
7
7
  from arraylake.cli.config import app as config_app
8
8
  from arraylake.cli.repo import app as repo_app
9
+ from arraylake.cli.skills import app as skills_app
9
10
  from arraylake.cli.utils import rich_console
10
11
  from arraylake.diagnostics import get_diagnostics
11
12
 
@@ -44,3 +45,4 @@ app.add_typer(auth, name="auth")
44
45
  app.add_typer(repo_app, name="repo")
45
46
  app.add_typer(config_app, name="config")
46
47
  app.add_typer(compute_app, name="compute")
48
+ app.add_typer(skills_app, name="skills")
@@ -0,0 +1,208 @@
1
+ """Install Arraylake Agent Skills into agent-specific skill directories.
2
+
3
+ Skills follow the open Agent Skills spec (https://agentskills.io) and are
4
+ hosted as static content on the Earthmover docs site. This command fetches
5
+ the manifest, downloads each skill's files, and writes them into the
6
+ directories agent skill loaders read from:
7
+
8
+ - `~/.claude/skills/` (Claude Code, opencode via compat shim)
9
+ - `~/.agents/skills/` (Codex, Cursor, Gemini CLI, and others)
10
+
11
+ With `--project`, installs into `.claude/skills/` and `.agents/skills/`
12
+ under the current working directory instead.
13
+
14
+ Skill content ships independently of the `arraylake` client release. Re-run
15
+ `arraylake skills install` to pick up new content; no `pip install -U`
16
+ required.
17
+ """
18
+
19
+ from __future__ import annotations
20
+
21
+ import shutil
22
+ import tempfile
23
+ from dataclasses import dataclass
24
+ from pathlib import Path
25
+
26
+ import httpx
27
+ import typer
28
+ from packaging.version import InvalidVersion, Version
29
+
30
+ import arraylake
31
+ from arraylake.cli.utils import rich_console
32
+ from arraylake.config import config
33
+
34
+ app = typer.Typer(help="Install Arraylake Agent Skills for AI coding agents")
35
+
36
+ DEFAULT_SKILLS_URI = "https://docs.earthmover.io/skills"
37
+
38
+
39
+ @dataclass(frozen=True)
40
+ class SkillEntry:
41
+ name: str
42
+ version: str
43
+ min_client_version: str | None
44
+ description: str
45
+ files: tuple[str, ...]
46
+
47
+
48
+ def _target_dirs(project: bool) -> list[Path]:
49
+ base = Path.cwd() if project else Path.home()
50
+ return [base / ".claude" / "skills", base / ".agents" / "skills"]
51
+
52
+
53
+ def _skills_uri() -> str:
54
+ uri: str = config.get("skills.uri", DEFAULT_SKILLS_URI)
55
+ return uri.rstrip("/")
56
+
57
+
58
+ def _client_version_ok(min_required: str | None) -> bool:
59
+ if not min_required:
60
+ return True
61
+ try:
62
+ return Version(arraylake.__version__) >= Version(min_required)
63
+ except InvalidVersion:
64
+ # Treat unparseable client versions (e.g. local dev builds without
65
+ # a tag) as "good enough" — don't block users on tooling weirdness.
66
+ return True
67
+
68
+
69
+ def _parse_manifest(payload: dict) -> list[SkillEntry]:
70
+ entries: list[SkillEntry] = []
71
+ for raw in payload.get("skills", []):
72
+ entries.append(
73
+ SkillEntry(
74
+ name=raw["name"],
75
+ version=str(raw.get("version", "0.0")),
76
+ min_client_version=raw.get("min_client_version"),
77
+ description=str(raw.get("description", "")),
78
+ files=tuple(raw.get("files", [])),
79
+ )
80
+ )
81
+ return entries
82
+
83
+
84
+ def _fetch_manifest(uri: str) -> list[SkillEntry]:
85
+ url = f"{uri}/manifest.json"
86
+ with httpx.Client(timeout=20.0, follow_redirects=True) as client:
87
+ resp = client.get(url)
88
+ resp.raise_for_status()
89
+ return _parse_manifest(resp.json())
90
+
91
+
92
+ def _empty_temporary_sibling(dest: Path, prefix: str) -> Path:
93
+ path = Path(tempfile.mkdtemp(prefix=prefix, dir=dest))
94
+ path.rmdir()
95
+ return path
96
+
97
+
98
+ def _remove_path(path: Path) -> None:
99
+ if path.is_symlink() or path.is_file():
100
+ path.unlink()
101
+ elif path.exists():
102
+ shutil.rmtree(path)
103
+
104
+
105
+ def _install_skill(uri: str, entry: SkillEntry, dest: Path) -> str:
106
+ """Download all of `entry.files` into `dest/<name>/...`.
107
+
108
+ Returns one-word status: written | conflict.
109
+ """
110
+ skill_root = dest / entry.name
111
+ if skill_root.exists() and not skill_root.is_dir():
112
+ return "conflict"
113
+
114
+ tmp_root = Path(tempfile.mkdtemp(prefix=".arraylake-skill-", dir=dest))
115
+ backup_root: Path | None = None
116
+
117
+ try:
118
+ with httpx.Client(timeout=30.0, follow_redirects=True) as client:
119
+ for rel in entry.files:
120
+ if ".." in rel.split("/") or rel.startswith("/"):
121
+ raise RuntimeError(f"Refusing to write file outside skill root: {rel}")
122
+ file_url = f"{uri}/{entry.name}/{rel}"
123
+ resp = client.get(file_url)
124
+ resp.raise_for_status()
125
+ target = tmp_root / rel
126
+ target.parent.mkdir(parents=True, exist_ok=True)
127
+ target.write_bytes(resp.content)
128
+
129
+ if skill_root.exists() or skill_root.is_symlink():
130
+ # Pre-existing install, including legacy wheel-bundled symlinks.
131
+ # Move it aside only after every replacement file has downloaded.
132
+ backup_root = _empty_temporary_sibling(dest, ".arraylake-skill-backup-")
133
+ skill_root.replace(backup_root)
134
+
135
+ tmp_root.replace(skill_root)
136
+ if backup_root is not None:
137
+ _remove_path(backup_root)
138
+ backup_root = None
139
+ return "written"
140
+ except Exception:
141
+ if backup_root is not None and (backup_root.exists() or backup_root.is_symlink()):
142
+ if skill_root.exists() or skill_root.is_symlink():
143
+ _remove_path(skill_root)
144
+ backup_root.replace(skill_root)
145
+ backup_root = None
146
+ raise
147
+ finally:
148
+ if tmp_root.exists() or tmp_root.is_symlink():
149
+ _remove_path(tmp_root)
150
+
151
+
152
+ @app.command()
153
+ def install(
154
+ project: bool = typer.Option(False, "--project", help="Install into the current directory instead of $HOME."),
155
+ url: str | None = typer.Option(
156
+ None,
157
+ "--url",
158
+ help=f"Override the skills server URL. Default: skills.uri config or {DEFAULT_SKILLS_URI}.",
159
+ ),
160
+ ) -> None:
161
+ """**Install** Arraylake skills into agent skill directories.
162
+
163
+ Default: user scope (`~/.claude/skills/` and `~/.agents/skills/`).
164
+ With `--project`: project scope under the current working directory.
165
+
166
+ Fetches the latest published skills from the Earthmover docs site
167
+ (configurable via `--url` or the `skills.uri` config key). Re-run this
168
+ to pick up content updates — they ship independently of the `arraylake`
169
+ package release cycle.
170
+ """
171
+ uri = (url or _skills_uri()).rstrip("/")
172
+ rich_console.print(f"Fetching skill manifest from [cyan]{uri}/manifest.json[/cyan]")
173
+ try:
174
+ entries = _fetch_manifest(uri)
175
+ except httpx.HTTPError as err:
176
+ rich_console.print(f"[red]Failed to fetch skills from {uri}: {err}[/red]")
177
+ raise typer.Exit(code=1) from err
178
+
179
+ if not entries:
180
+ rich_console.print("[yellow]Manifest contained no skills; nothing to install.[/yellow]")
181
+ return
182
+
183
+ ok = True
184
+ for target in _target_dirs(project):
185
+ target.mkdir(parents=True, exist_ok=True)
186
+ for entry in entries:
187
+ if not _client_version_ok(entry.min_client_version):
188
+ rich_console.print(
189
+ f"[yellow]skip[/yellow] {target / entry.name} "
190
+ f"(skill requires arraylake>={entry.min_client_version}, "
191
+ f"installed {arraylake.__version__})"
192
+ )
193
+ continue
194
+ try:
195
+ status = _install_skill(uri, entry, target)
196
+ except httpx.HTTPError as err:
197
+ rich_console.print(f"[red]error[/red] {target / entry.name}: {err}")
198
+ ok = False
199
+ continue
200
+ if status == "conflict":
201
+ ok = False
202
+ rich_console.print(f"[yellow]skip[/yellow] {target / entry.name} (path exists and is not a directory)")
203
+ else:
204
+ rich_console.print(f"[green]{status:8}[/green]{target / entry.name}")
205
+
206
+ if not ok:
207
+ rich_console.print("\n[yellow]One or more skills were skipped. Resolve the conflicts above and re-run.[/yellow]")
208
+ raise typer.Exit(code=1)
@@ -26,7 +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
+ from arraylake._credential_cache import CredentialCacheKey, get_or_refresh, populate
30
30
  from arraylake.asyn import async_gather_tasks, asyncio_run, sync
31
31
  from arraylake.compute.services import AsyncComputeClient, ComputeClient
32
32
  from arraylake.config import config as arraylake_config
@@ -66,10 +66,12 @@ from arraylake.types import (
66
66
  ExplicitVirtualChunkAccessPolicyResponse,
67
67
  GSCredentials,
68
68
  NewBucket,
69
+ OpenRepoResponse,
69
70
  OptimizationConfig,
70
71
  OrgActions,
71
72
  OrgAndRepoName,
72
73
  OrgName,
74
+ Platform,
73
75
  RepoActions,
74
76
  RepoKind,
75
77
  RepoMetadataT,
@@ -79,6 +81,7 @@ from arraylake.types import (
79
81
  S3Credentials,
80
82
  StorageOptions,
81
83
  TempCredentials,
84
+ VirtualChunkCredentials,
82
85
  validate_name,
83
86
  validate_org_and_repo_name,
84
87
  )
@@ -213,6 +216,25 @@ class AsyncClient:
213
216
  break
214
217
  page += 1
215
218
 
219
+ def _credential_cache_key(
220
+ self,
221
+ *,
222
+ scope: Literal["repo", "bucket", "vcc"],
223
+ org: OrgName,
224
+ identifier: str,
225
+ platform: Platform,
226
+ access: Literal["read", "write"] = "read",
227
+ ) -> CredentialCacheKey:
228
+ return CredentialCacheKey(
229
+ api_url=self.service_uri,
230
+ auth_key=hash(self.token),
231
+ scope=scope,
232
+ org=org,
233
+ identifier=identifier,
234
+ platform=platform,
235
+ access=access,
236
+ )
237
+
216
238
  async def _get_s3_delegated_credentials_from_repo(self, org: OrgName, repo_name: RepoName) -> S3Credentials:
217
239
  """Get delegated credentials for a repo's S3 bucket.
218
240
 
@@ -223,14 +245,7 @@ class AsyncClient:
223
245
  Returns:
224
246
  S3Credentials: Temporary credentials for the S3 bucket.
225
247
  """
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
- )
248
+ key = self._credential_cache_key(scope="repo", org=org, identifier=repo_name, platform="s3")
234
249
 
235
250
  async def fetch() -> S3Credentials:
236
251
  mstore = self._metastore_for_org(org)
@@ -250,14 +265,7 @@ class AsyncClient:
250
265
  Returns:
251
266
  GSCredentials: Temporary credentials for the GCS bucket.
252
267
  """
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
- )
268
+ key = self._credential_cache_key(scope="repo", org=org, identifier=repo_name, platform="gs")
261
269
 
262
270
  async def fetch() -> GSCredentials:
263
271
  mstore = self._metastore_for_org(org)
@@ -277,14 +285,7 @@ class AsyncClient:
277
285
  Returns:
278
286
  AzureCredentials: Temporary credentials for the Azure Blob Storage container.
279
287
  """
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
+ key = self._credential_cache_key(scope="repo", org=org, identifier=repo_name, platform="azure")
288
289
 
289
290
  async def fetch() -> AzureCredentials:
290
291
  mstore = self._metastore_for_org(org)
@@ -310,15 +311,7 @@ class AsyncClient:
310
311
  Returns:
311
312
  S3Credentials: Temporary credentials for the S3 bucket.
312
313
  """
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
- )
314
+ key = self._credential_cache_key(scope="bucket", org=org, identifier=nickname, platform="s3", access=access)
322
315
 
323
316
  async def fetch() -> S3Credentials:
324
317
  mstore = self._metastore_for_org(org)
@@ -345,15 +338,7 @@ class AsyncClient:
345
338
  Returns:
346
339
  GSCredentials: Temporary credentials for the GCS bucket.
347
340
  """
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
- )
341
+ key = self._credential_cache_key(scope="bucket", org=org, identifier=nickname, platform="gs", access=access)
357
342
 
358
343
  async def fetch() -> GSCredentials:
359
344
  mstore = self._metastore_for_org(org)
@@ -380,15 +365,7 @@ class AsyncClient:
380
365
  Returns:
381
366
  AzureCredentials: Temporary credentials for the Azure Blob Storage container.
382
367
  """
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
- )
368
+ key = self._credential_cache_key(scope="bucket", org=org, identifier=nickname, platform="azure", access=access)
392
369
 
393
370
  async def fetch() -> AzureCredentials:
394
371
  mstore = self._metastore_for_org(org)
@@ -413,12 +390,13 @@ class AsyncClient:
413
390
  Returns:
414
391
  Callable: Function that returns a S3StaticCredentials object.
415
392
  """
416
- s3_credentials = asyncio_run(self._get_s3_delegated_credentials_from_repo(org, repo_name), timeout=None)
393
+ credentials = asyncio_run(self._refresh_repo_bucket_credentials(org, repo_name, "s3"), timeout=None)
394
+ assert isinstance(credentials, S3Credentials)
417
395
  return icechunk.S3StaticCredentials(
418
- access_key_id=s3_credentials.aws_access_key_id,
419
- secret_access_key=s3_credentials.aws_secret_access_key,
420
- session_token=s3_credentials.aws_session_token,
421
- expires_after=s3_credentials.expiration,
396
+ access_key_id=credentials.aws_access_key_id,
397
+ secret_access_key=credentials.aws_secret_access_key,
398
+ session_token=credentials.aws_session_token,
399
+ expires_after=credentials.expiration,
422
400
  )
423
401
 
424
402
  def _get_icechunk_gcs_credentials_refresh_function_for_repo(self, org: OrgName, repo_name: RepoName) -> icechunk.GcsBearerCredential:
@@ -435,10 +413,11 @@ class AsyncClient:
435
413
  Returns:
436
414
  Callable: Function that returns a GcsBearerCredential object.
437
415
  """
438
- gcs_credentials = asyncio_run(self._get_gcs_delegated_credentials_from_repo(org, repo_name), timeout=None)
416
+ credentials = asyncio_run(self._refresh_repo_bucket_credentials(org, repo_name, "gs"), timeout=None)
417
+ assert isinstance(credentials, GSCredentials)
439
418
  return icechunk.GcsBearerCredential(
440
- bearer=gcs_credentials.access_token,
441
- expires_after=gcs_credentials.expiration.replace(tzinfo=UTC) if gcs_credentials.expiration else None,
419
+ bearer=credentials.access_token,
420
+ expires_after=credentials.expiration.replace(tzinfo=UTC) if credentials.expiration else None,
442
421
  )
443
422
 
444
423
  def _get_icechunk_azure_credentials_refresh_function_for_repo(
@@ -497,6 +476,138 @@ class AsyncClient:
497
476
  ) -> icechunk.AzureStaticCredentials:
498
477
  raise NotImplementedError("Azure credential refreshes not implemented in Icechunk.")
499
478
 
479
+ def _populate_credential_cache(self, org: OrgName, repo_name: RepoName, response: OpenRepoResponse) -> None:
480
+ """Seed the credential cache with every delegated credential from an open-repo response.
481
+
482
+ A single ``/open`` returns fresh credentials for the repo's own bucket and for all of
483
+ its virtual chunk containers, so we cache them all (each keyed independently and honoring
484
+ its own expiry). The next refresh of any one of them is then served from the cache until
485
+ it individually expires, rather than each triggering its own ``/open``.
486
+ """
487
+ if not self._cache_credentials:
488
+ return
489
+ if response.repo_credentials is not None:
490
+ platform = self._delegated_credential_platform(response.repo_bucket)
491
+ if platform is not None:
492
+ populate(
493
+ self._credential_cache_key(scope="repo", org=org, identifier=repo_name, platform=platform),
494
+ response.repo_credentials,
495
+ )
496
+ for vcc_prefix, vcc_creds in response.virtual_chunk_credentials.items():
497
+ if vcc_creds.credentials is not None and vcc_creds.credentials.expiration is not None:
498
+ populate(
499
+ self._credential_cache_key(scope="vcc", org=org, identifier=vcc_prefix, platform=vcc_creds.platform),
500
+ vcc_creds.credentials,
501
+ )
502
+
503
+ async def _open_repo_and_cache(self, org: OrgName, repo_name: RepoName) -> OpenRepoResponse:
504
+ """Open the repo once and seed the credential cache with all of its delegated credentials."""
505
+ mstore = self._metastore_for_org(org)
506
+ response = await mstore.open_repo(repo_name)
507
+ self._populate_credential_cache(org, repo_name, response)
508
+ return response
509
+
510
+ async def _refresh_repo_bucket_credentials(self, org: OrgName, repo_name: RepoName, platform: Platform) -> TempCredentials:
511
+ """Refresh the repo's own bucket credentials via ``/open``, going through the cache.
512
+
513
+ Because ``/open`` also seeds the virtual chunk container credentials, a single repo-bucket
514
+ refresh lands fresh credentials for the repo bucket and every virtual bucket at once.
515
+ """
516
+ key = self._credential_cache_key(scope="repo", org=org, identifier=repo_name, platform=platform)
517
+
518
+ async def fetch() -> TempCredentials:
519
+ response = await self._open_repo_and_cache(org, repo_name)
520
+ if response.repo_credentials is None:
521
+ raise ValueError(f"Repo '{org}/{repo_name}' returned no refreshable credentials.")
522
+ return response.repo_credentials
523
+
524
+ return await get_or_refresh(key, fetch, use_cache=self._cache_credentials)
525
+
526
+ async def _get_vcc_credentials_from_repo(
527
+ self, org: OrgName, repo_name: RepoName, vcc_prefix: BucketPrefix, platform: Platform
528
+ ) -> TempCredentials:
529
+ """Re-fetch credentials for a single virtual chunk container, going through the cache.
530
+
531
+ On a cache miss this re-opens the repo, which both refreshes the requested container's
532
+ credentials and seeds the cache for the repo bucket and all other containers. Reusing the
533
+ repo-open path means refreshes inherit the same authorization and resolution as the initial
534
+ open, including cross-org (filtered subscription / marketplace) virtual chunks — the server
535
+ never exposes the provider's org or bucket to us, only the credentials.
536
+ """
537
+ key = self._credential_cache_key(scope="vcc", org=org, identifier=vcc_prefix, platform=platform)
538
+
539
+ async def fetch() -> TempCredentials:
540
+ response = await self._open_repo_and_cache(org, repo_name)
541
+ vcc_creds = response.virtual_chunk_credentials.get(vcc_prefix)
542
+ if vcc_creds is None or vcc_creds.credentials is None:
543
+ raise ValueError(
544
+ f"Could not refresh credentials for virtual chunk container prefix {vcc_prefix!r} in repo "
545
+ f"'{org}/{repo_name}'; the repo may no longer authorize access to it."
546
+ )
547
+ return vcc_creds.credentials
548
+
549
+ return await get_or_refresh(key, fetch, use_cache=self._cache_credentials)
550
+
551
+ def _get_icechunk_s3_vcc_credentials_refresh_function(
552
+ self, org: OrgName, repo_name: RepoName, vcc_prefix: BucketPrefix, platform: Platform
553
+ ) -> icechunk.S3StaticCredentials:
554
+ """Synchronous, no-arg S3 refresh function for a virtual chunk container.
555
+
556
+ Icechunk calls this (synchronously, with no event loop) when the container's
557
+ credentials expire, so it bridges to the async fetch via ``asyncio_run``.
558
+ """
559
+ credentials = asyncio_run(self._get_vcc_credentials_from_repo(org, repo_name, vcc_prefix, platform), timeout=None)
560
+ assert isinstance(credentials, S3Credentials)
561
+ return icechunk.S3StaticCredentials(
562
+ access_key_id=credentials.aws_access_key_id,
563
+ secret_access_key=credentials.aws_secret_access_key,
564
+ session_token=credentials.aws_session_token,
565
+ expires_after=credentials.expiration,
566
+ )
567
+
568
+ def _get_icechunk_gcs_vcc_credentials_refresh_function(
569
+ self, org: OrgName, repo_name: RepoName, vcc_prefix: BucketPrefix, platform: Platform
570
+ ) -> icechunk.GcsBearerCredential:
571
+ """Synchronous, no-arg GCS refresh function for a virtual chunk container.
572
+
573
+ Icechunk calls this (synchronously, with no event loop) when the container's
574
+ credentials expire, so it bridges to the async fetch via ``asyncio_run``.
575
+ """
576
+ credentials = asyncio_run(self._get_vcc_credentials_from_repo(org, repo_name, vcc_prefix, platform), timeout=None)
577
+ assert isinstance(credentials, GSCredentials)
578
+ return icechunk.GcsBearerCredential(
579
+ bearer=credentials.access_token,
580
+ expires_after=credentials.expiration.replace(tzinfo=UTC) if credentials.expiration else None,
581
+ )
582
+
583
+ def _maybe_get_vcc_credential_refresh_func(
584
+ self, vcc_creds: VirtualChunkCredentials, org: OrgName, repo_name: RepoName, vcc_prefix: BucketPrefix
585
+ ) -> Callable | None:
586
+ """Build a refresh closure for a server-provided virtual chunk container, if its
587
+ credentials are refreshable.
588
+
589
+ Delegated/temporary credentials carry an expiration; static (long-lived), anonymous,
590
+ and bucket-policy credentials do not and are left as-is. Azure is not yet supported by
591
+ Icechunk so it never gets a refresh function.
592
+ """
593
+ credentials = vcc_creds.credentials
594
+ if credentials is None or credentials.expiration is None:
595
+ return None
596
+ if vcc_creds.platform in ("s3", "s3-compatible", "minio"):
597
+ return partial(self._get_icechunk_s3_vcc_credentials_refresh_function, org, repo_name, vcc_prefix, vcc_creds.platform)
598
+ elif vcc_creds.platform == "gs":
599
+ return partial(self._get_icechunk_gcs_vcc_credentials_refresh_function, org, repo_name, vcc_prefix, vcc_creds.platform)
600
+ return None
601
+
602
+ def _delegated_credential_platform(self, bucket: BucketResponse) -> Platform | None:
603
+ if bucket.platform == "s3" or _is_r2_bucket(bucket):
604
+ return "s3"
605
+ if bucket.platform == "gs":
606
+ return "gs"
607
+ if bucket.platform == "azure":
608
+ return "azure"
609
+ return None
610
+
500
611
  async def _maybe_get_credentials_for_icechunk(
501
612
  self,
502
613
  bucket: BucketResponse,
@@ -513,17 +624,18 @@ class AsyncClient:
513
624
  repo-scoped credential fetches always inherit the caller's permissions on the repo.
514
625
  """
515
626
  if _use_delegated_credentials(bucket):
516
- if bucket.platform == "s3" or _is_r2_bucket(bucket):
627
+ platform = self._delegated_credential_platform(bucket)
628
+ if platform == "s3":
517
629
  if repo_name:
518
630
  return await self._get_s3_delegated_credentials_from_repo(org, repo_name)
519
631
  else:
520
632
  return await self._get_s3_delegated_credentials_from_bucket(org, bucket.nickname, access=bucket_access)
521
- elif bucket.platform == "gs":
633
+ elif platform == "gs":
522
634
  if repo_name:
523
635
  return await self._get_gcs_delegated_credentials_from_repo(org, repo_name)
524
636
  else:
525
637
  return await self._get_gcs_delegated_credentials_from_bucket(org, bucket.nickname, access=bucket_access)
526
- elif bucket.platform == "azure":
638
+ elif platform == "azure":
527
639
  if repo_name:
528
640
  return await self._get_azure_delegated_credentials_from_repo(org, repo_name)
529
641
  else:
@@ -543,17 +655,18 @@ class AsyncClient:
543
655
  Returns None if delegated credentials are not configured for the bucket.
544
656
  """
545
657
  if _use_delegated_credentials(bucket):
546
- if bucket.platform == "s3" or _is_r2_bucket(bucket):
658
+ platform = self._delegated_credential_platform(bucket)
659
+ if platform == "s3":
547
660
  if repo_name:
548
661
  return partial(self._get_icechunk_s3_credentials_refresh_function_for_repo, org, repo_name)
549
662
  else:
550
663
  return partial(self._get_icechunk_s3_credentials_refresh_function_for_bucket, org, bucket.nickname)
551
- elif bucket.platform == "gs":
664
+ elif platform == "gs":
552
665
  if repo_name:
553
666
  return partial(self._get_icechunk_gcs_credentials_refresh_function_for_repo, org, repo_name)
554
667
  else:
555
668
  return partial(self._get_icechunk_gcs_credentials_refresh_function_for_bucket, org, bucket.nickname)
556
- elif bucket.platform == "azure":
669
+ elif platform == "azure":
557
670
  # Credential refreshes currently not implemented in Icechunk, return None
558
671
  return None
559
672
  else:
@@ -630,6 +743,9 @@ class AsyncClient:
630
743
  bucket=response.repo_bucket, org=org, repo_name=repo_name
631
744
  )
632
745
  credentials = response.repo_credentials if credential_refresh_func is None else None
746
+ # Seed the cache with the repo bucket credentials and every virtual chunk container's
747
+ # credentials, so the first refresh of any of them is served from this single open.
748
+ self._populate_credential_cache(org, repo_name, response)
633
749
  icechunk_storage = _get_icechunk_storage_obj(
634
750
  bucket_config=response.repo_bucket,
635
751
  prefix=response.repo_prefix,
@@ -682,17 +798,18 @@ class AsyncClient:
682
798
  container.name = vcc_creds.runtime_vcc_name
683
799
  config.set_virtual_chunk_container(container)
684
800
 
685
- # TODO: build refresh closures here using vcc_creds.org, vcc_creds.bucket_nickname,
686
- # and vcc_creds.platform so that credentials can be refreshed when they expire. (See https://github.com/earth-mover/arraylake/issues/5122)
687
801
  # TODO: support Azure virtual chunk credentials
688
802
  if isinstance(vcc_creds.credentials, AzureCredentials):
689
803
  raise NotImplementedError(f"Azure virtual chunk credentials are not yet supported. (prefix: {vcc_prefix}).")
804
+ # Refresh delegated/temporary credentials by re-opening the repo; pass static
805
+ # credentials only when they aren't refreshable (icechunk forbids supplying both).
806
+ refresh_func = self._maybe_get_vcc_credential_refresh_func(vcc_creds, org, repo_name, vcc_prefix)
690
807
  # TODO: fix return type of get_icechunk_container_credentials (returns GcsBearerCredential
691
808
  # which is not in icechunk's AnyGcsCredential union — icechunk typing gap)
692
809
  credentials_map[vcc_prefix] = get_icechunk_container_credentials( # type: ignore[assignment]
693
810
  bucket_platform=vcc_creds.platform,
694
- credentials=vcc_creds.credentials,
695
- credential_refresh_func=None,
811
+ credentials=None if refresh_func is not None else vcc_creds.credentials,
812
+ credential_refresh_func=refresh_func,
696
813
  )
697
814
  vcc_credentials = icechunk.containers_credentials(credentials_map)
698
815