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.
- {arraylake-1.0.3 → arraylake-1.1.0}/PKG-INFO +1 -1
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/_credential_cache.py +10 -1
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/_version.py +2 -2
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/main.py +2 -0
- arraylake-1.1.0/arraylake/cli/skills.py +208 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/client.py +187 -70
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/compute/types.py +9 -1
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/repos/icechunk/types.py +1 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/repos/icechunk/virtual.py +3 -3
- {arraylake-1.0.3 → arraylake-1.1.0}/pyproject.toml +2 -2
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/conftest.py +10 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/icechunk/test_virtual.py +15 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_credential_cache.py +104 -8
- arraylake-1.1.0/tests/test_skills.py +111 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_virtual_chunks.py +96 -0
- arraylake-1.1.0/uv.lock +2802 -0
- arraylake-1.0.3/uv.lock +0 -2795
- {arraylake-1.0.3 → arraylake-1.1.0}/.gitignore +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/README.md +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/__init__.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/__main__.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/api_utils.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/asyn.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/__init__.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/auth.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/compute.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/config.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/repo.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/cli/utils.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/compute/doctor.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/compute/http_client.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/compute/services.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/config.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/config.yaml +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/credentials.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/diagnostics.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/display/__init__.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/display/repo.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/display/repolist.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/exceptions.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/log_util.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/metastore/__init__.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/metastore/abc.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/metastore/http_metastore.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/py.typed +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/repos/__init__.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/repos/icechunk/__init__.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/repos/icechunk/storage.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/token.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/tuning.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/arraylake/types.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/__init__.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/compute/test_client.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/compute/test_doctor.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/config.yaml +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/icechunk/__init__.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/icechunk/conftest.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/icechunk/test_storage.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_api_utils.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_asyn.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_cli.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_client.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_config.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_credentials.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_diagnostics.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_performance.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_subscriptions.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_token_handler.py +0 -0
- {arraylake-1.0.3 → arraylake-1.1.0}/tests/test_types.py +0 -0
|
@@ -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
|
|
22
|
-
__version_tuple__ = version_tuple = (1,
|
|
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 =
|
|
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 =
|
|
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 =
|
|
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 =
|
|
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 =
|
|
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 =
|
|
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
|
-
|
|
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=
|
|
419
|
-
secret_access_key=
|
|
420
|
-
session_token=
|
|
421
|
-
expires_after=
|
|
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
|
-
|
|
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=
|
|
441
|
-
expires_after=
|
|
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
|
-
|
|
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
|
|
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
|
|
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
|
-
|
|
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
|
|
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
|
|
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=
|
|
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
|
|