datavo-cli 0.22.0__py3-none-any.whl

datavo_cli/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ["__version__"]
+
+__version__ = "0.22.0"

datavo_cli/client_factory.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import base64
+import json
+import os
+import time
+from dataclasses import dataclass, replace
+
+from datavo_sdk import (
+    DatavoClient,
+    EntraConfig,
+    EntraTokenProvider,
+    PublicServerConfig,
+    default_msal_cache_path,
+    get_server_config,
+)
+
+from .config_store import DatavoConfig, DatavoProfile, apply_env_overrides, load_config, resolve_profile, update_profile
+
+
+ENTRA_AUTH_PROVIDERS = {"entra", "azure", "azuread", "aad", "easyauth"}
+
+
+@dataclass(frozen=True)
+class RuntimeContext:
+    profile_name: str
+    profile: DatavoProfile
+    server_config: PublicServerConfig | None
+    auth_provider: str
+
+
+def _jwt_expiry(token_text: str) -> int | None:
+    parts = token_text.strip().split(".")
+    if len(parts) < 2:
+        return None
+    payload = parts[1]
+    payload += "=" * (-len(payload) % 4)
+    try:
+        decoded = base64.urlsafe_b64decode(payload.encode("utf-8"))
+        data = json.loads(decoded.decode("utf-8"))
+    except Exception:
+        return None
+    exp = data.get("exp") if isinstance(data, dict) else None
+    return int(exp) if isinstance(exp, (int, float)) else None
+
+
+def _token_needs_refresh(profile: DatavoProfile) -> bool:
+    token = profile.auth_token
+    if not isinstance(token, str) or not token.strip():
+        return True
+    expiry = _jwt_expiry(token)
+    if expiry is None:
+        return False
+    return int(time.time()) >= expiry
+
+
+def try_silent_refresh(profile_name: str, profile: DatavoProfile) -> DatavoProfile | None:
+    if not profile.msal_client_id or not profile.msal_tenant_id or not profile.msal_scopes:
+        return None
+    if not default_msal_cache_path().exists():
+        return None
+    try:
+        provider = EntraTokenProvider(
+            EntraConfig(
+                tenant_id=profile.msal_tenant_id,
+                client_id=profile.msal_client_id,
+                scopes=list(profile.msal_scopes),
+            )
+        )
+        token = provider.acquire_token_silent()
+        if not token:
+            return None
+        update_profile(profile_name, auth_token=token)
+        return replace(profile, auth_token=token)
+    except Exception:
+        return None
+
+
+def resolve_runtime_context(
+    *,
+    profile_name: str | None = None,
+    config: DatavoConfig | None = None,
+    require_url: bool = True,
+) -> RuntimeContext:
+    config = config or load_config()
+    selected_name, selected_profile = resolve_profile(config, profile_name=profile_name)
+    effective_profile = apply_env_overrides(selected_profile)
+    if require_url and not effective_profile.api_base_url:
+        raise RuntimeError("This command requires an API URL. Run `datavo config set <profile> --url <url>`.")
+    server_config = None
+    auth_provider = (os.environ.get("DATAVO_AUTH_PROVIDER") or "").strip().lower()
+    if effective_profile.api_base_url:
+        try:
+            server_config = get_server_config(effective_profile.api_base_url)
+        except Exception:
+            server_config = None
+        if not auth_provider and server_config is not None:
+            auth_provider = (server_config.auth_provider or "").strip().lower()
+    if not auth_provider:
+        auth_provider = "easyauth"
+    if (
+        not os.environ.get("DATAVO_AUTH_TOKEN")
+        and _token_needs_refresh(effective_profile)
+        and selected_name in config.profiles
+    ):
+        refreshed = try_silent_refresh(selected_name, effective_profile)
+        if refreshed is not None:
+            effective_profile = refreshed
+    return RuntimeContext(
+        profile_name=selected_name,
+        profile=effective_profile,
+        server_config=server_config,
+        auth_provider=auth_provider,
+    )
+
+
+def build_client(*, profile_name: str | None = None, config: DatavoConfig | None = None) -> DatavoClient:
+    runtime = resolve_runtime_context(profile_name=profile_name, config=config, require_url=True)
+    return DatavoClient(
+        server_url=runtime.profile.api_base_url or "",
+        auth_token=runtime.profile.auth_token,
+        auto_login=False,
+    )
+
+
+def derive_scopes(
+    *,
+    requested_scopes: list[str] | None,
+    server_config: PublicServerConfig | None,
+) -> list[str]:
+    if requested_scopes:
+        return [scope.strip() for scope in requested_scopes if scope and scope.strip()]
+    env_scopes = (os.environ.get("DATAVO_ENTRA_SCOPES") or "").strip()
+    if env_scopes:
+        return [scope.strip() for scope in env_scopes.split(",") if scope.strip()]
+    audience = (os.environ.get("DATAVO_ENTRA_AUDIENCE") or "").strip()
+    if not audience and server_config is not None:
+        audience = (server_config.entra_audience or "").strip()
+    if not audience:
+        raise RuntimeError("No Entra audience available. Provide --scope or configure entra_audience on the server.")
+    if audience.startswith("api://"):
+        return [f"{audience}/access_as_user"]
+    return [f"api://{audience}/access_as_user"]
+
+
+def derive_msal_profile(
+    *,
+    base_profile: DatavoProfile,
+    server_config: PublicServerConfig | None,
+    tenant_id: str | None,
+    client_id: str | None,
+    scopes: list[str] | None,
+    auth_token: str,
+) -> DatavoProfile:
+    resolved_tenant_id = (tenant_id or os.environ.get("DATAVO_ENTRA_TENANT_ID") or "").strip()
+    if not resolved_tenant_id and server_config is not None:
+        resolved_tenant_id = (server_config.entra_tenant_id or "").strip()
+    if not resolved_tenant_id:
+        raise RuntimeError("No Entra tenant id available. Provide --tenant-id or configure entra_tenant_id on the server.")
+    resolved_client_id = (client_id or os.environ.get("DATAVO_ENTRA_CLIENT_ID") or "").strip()
+    if not resolved_client_id and server_config is not None:
+        resolved_client_id = (server_config.entra_public_client_id or "").strip()
+    if not resolved_client_id:
+        raise RuntimeError("No Entra public client id available. Provide --client-id or configure entra_public_client_id on the server.")
+    resolved_scopes = derive_scopes(requested_scopes=scopes, server_config=server_config)
+    return replace(
+        base_profile,
+        auth_token=auth_token,
+        msal_tenant_id=resolved_tenant_id,
+        msal_client_id=resolved_client_id,
+        msal_scopes=resolved_scopes,
+    )

datavo_cli/commands/__init__.py

@@ -0,0 +1,36 @@
+from .auth import handle_auth_command
+from .config import handle_config_command
+from .dataset import handle_dataset_command
+from .events import handle_events_command
+from .jobpool import handle_render_dataset_cache_command, handle_warm_command
+from .sample import handle_sample_command
+from .sample_store import handle_sample_store_command
+from .source_import import handle_source_import_command
+from .split_set import handle_split_set_command
+from .stats import handle_stats_command
+from .team import (
+    handle_share_command,
+    handle_teams_command,
+    handle_transfer_command,
+    handle_unshare_command,
+)
+from .user import handle_users_command
+
+__all__ = [
+    "handle_auth_command",
+    "handle_config_command",
+    "handle_dataset_command",
+    "handle_events_command",
+    "handle_render_dataset_cache_command",
+    "handle_sample_command",
+    "handle_sample_store_command",
+    "handle_source_import_command",
+    "handle_split_set_command",
+    "handle_share_command",
+    "handle_stats_command",
+    "handle_teams_command",
+    "handle_transfer_command",
+    "handle_unshare_command",
+    "handle_users_command",
+    "handle_warm_command",
+]

datavo_cli/commands/auth.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from dataclasses import replace
+
+from datavo_sdk import EntraConfig, EntraTokenProvider, default_msal_cache_path
+
+from ..client_factory import ENTRA_AUTH_PROVIDERS, derive_msal_profile, resolve_runtime_context
+from ..config_store import save_profile, update_profile
+from ..render import render_output
+
+
+def _clear_msal_cache() -> None:
+    path = default_msal_cache_path()
+    if path.exists():
+        path.unlink()
+
+
+def _run_msal_login(*, tenant_id: str, client_id: str, scopes: list[str], force_browser: bool) -> str:
+    provider = EntraTokenProvider(
+        EntraConfig(tenant_id=tenant_id, client_id=client_id, scopes=list(scopes))
+    )
+    return provider.login(force_browser=force_browser)
+
+
+def handle_auth_command(args) -> int:
+    output_format = getattr(args, "format", "plain")
+    runtime = resolve_runtime_context(require_url=False)
+    if args.auth_command == "status":
+        render_output(
+            {
+                "profile": runtime.profile_name,
+                "api_base_url": runtime.profile.api_base_url,
+                "auth_provider": runtime.auth_provider,
+                "auth_token_present": bool(runtime.profile.auth_token),
+                "msal_configured": bool(runtime.profile.msal_client_id and runtime.profile.msal_tenant_id and runtime.profile.msal_scopes),
+            },
+            output_format=output_format,
+        )
+        return 0
+    if args.auth_command == "logout":
+        update_profile(runtime.profile_name, auth_token=None)
+        _clear_msal_cache()
+        render_output({"profile": runtime.profile_name, "logged_out": True}, output_format=output_format)
+        return 0
+    if args.auth_command != "login":
+        raise RuntimeError("Missing auth subcommand.")
+
+    if args.token:
+        updated = replace(runtime.profile, auth_token=args.token.strip())
+        save_profile(runtime.profile_name, updated)
+        render_output(
+            {
+                "profile": runtime.profile_name,
+                "auth_provider": runtime.auth_provider,
+                "auth_token_present": True,
+            },
+            output_format=output_format,
+        )
+        return 0
+
+    use_msal = bool(args.msal or args.msal_browser)
+    if not use_msal and runtime.auth_provider in ENTRA_AUTH_PROVIDERS:
+        use_msal = True
+    if not use_msal:
+        raise RuntimeError("Provide --token or use MSAL against an Entra/EasyAuth-enabled Datavo server.")
+
+    msal_profile = derive_msal_profile(
+        base_profile=runtime.profile,
+        server_config=runtime.server_config,
+        tenant_id=args.tenant_id,
+        client_id=args.client_id,
+        scopes=args.scope,
+        auth_token="",
+    )
+    token = _run_msal_login(
+        tenant_id=msal_profile.msal_tenant_id or "",
+        client_id=msal_profile.msal_client_id or "",
+        scopes=msal_profile.msal_scopes or [],
+        force_browser=bool(args.msal_browser),
+    )
+    msal_profile = replace(msal_profile, auth_token=token)
+    save_profile(runtime.profile_name, msal_profile)
+    render_output(
+        {
+            "profile": runtime.profile_name,
+            "auth_provider": runtime.auth_provider,
+            "auth_token_present": True,
+            "msal_tenant_id": msal_profile.msal_tenant_id,
+            "msal_client_id": msal_profile.msal_client_id,
+            "msal_scopes": msal_profile.msal_scopes,
+        },
+        output_format=output_format,
+    )
+    return 0

datavo_cli/commands/config.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import sys
+from urllib.parse import urlparse
+
+from datavo_sdk import fetch_server_config, save_server_cache
+
+from ..config_store import DatavoProfile, get_active_profile_name, load_config, remove_profile, save_profile, set_active_profile
+from ..render import render_output
+
+
+def _derive_profile_name(url: str) -> str:
+    parsed = urlparse(url if "://" in url else f"https://{url}")
+    hostname = parsed.hostname
+    if not hostname:
+        raise RuntimeError(f"Cannot derive a profile name from URL {url!r}; pass an explicit profile name.")
+    return hostname
+
+
+def handle_config_command(args) -> int:
+    output_format = getattr(args, "format", "plain")
+    if args.config_command == "show":
+        config = load_config()
+        profile_name = args.profile or get_active_profile_name(config)
+        if not profile_name:
+            raise RuntimeError("No active Datavo profile.")
+        if profile_name not in config.profiles:
+            raise RuntimeError(f"Unknown Datavo profile: {profile_name}")
+        profile = config.profiles[profile_name]
+        render_output(
+            {
+                "active_profile": config.active_profile,
+                "profile": profile_name,
+                "api_base_url": profile.api_base_url,
+                "auth_token_present": bool(profile.auth_token),
+                "msal_tenant_id": profile.msal_tenant_id,
+                "msal_client_id": profile.msal_client_id,
+                "msal_scopes": profile.msal_scopes,
+            },
+            output_format=output_format,
+        )
+        return 0
+    if args.config_command == "list":
+        config = load_config()
+        render_output(
+            {
+                "profiles": [
+                    {
+                        "name": name,
+                        "api_base_url": profile.api_base_url,
+                        "active": name == config.active_profile,
+                    }
+                    for name, profile in sorted(config.profiles.items())
+                ]
+            },
+            output_format=output_format,
+        )
+        return 0
+    if args.config_command == "use":
+        updated = set_active_profile(args.profile)
+        render_output({"active_profile": updated.active_profile}, output_format=output_format)
+        return 0
+    if args.config_command == "set":
+        profile_name = args.profile or _derive_profile_name(args.url)
+        config = load_config()
+        existing = config.profiles.get(profile_name, DatavoProfile())
+        profile = DatavoProfile(
+            api_base_url=args.url.rstrip("/"),
+            auth_token=existing.auth_token,
+            msal_tenant_id=existing.msal_tenant_id,
+            msal_client_id=existing.msal_client_id,
+            msal_scopes=existing.msal_scopes,
+        )
+        save_profile(profile_name, profile, set_active=True)
+        try:
+            public_config = fetch_server_config(profile.api_base_url or "")
+            save_server_cache(profile.api_base_url or "", public_config)
+        except Exception as exc:
+            print(f"warning: failed to fetch public config from {profile.api_base_url}: {exc}", file=sys.stderr)
+        render_output(
+            {
+                "active_profile": profile_name,
+                "profile": profile_name,
+                "api_base_url": profile.api_base_url,
+            },
+            output_format=output_format,
+        )
+        return 0
+    if args.config_command == "remove":
+        updated = remove_profile(args.profile)
+        render_output({"active_profile": updated.active_profile}, output_format=output_format)
+        return 0
+    raise RuntimeError("Missing config subcommand.")

datavo_cli/commands/dataset.py

@@ -0,0 +1,198 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from datavo_sdk import (
+    AddKeyGroupRequest,
+    DatasetCreateRequest,
+    DatasetPreviewRequest,
+    ShardPlan,
+    SourceSpec,
+    WorksetCreateRequest,
+    WorksetIssueStreamTicketsRequest,
+    WorksetOpenRequest,
+    WorksetShardSelector,
+)
+
+from ..client_factory import build_client
+from ..render import render_output
+from ..utils import load_model
+
+
+def _is_workset_expressible(request: DatasetCreateRequest) -> bool:
+    """Whether this composition can be selected as a workset first.
+
+    Only store-backed sources can. A workset takes its scope from its first source's
+    ``sample_store``, so a ``dataset_scope`` source makes a workset that is a dataset
+    *read* — explicitly not convertible. It has no workset representation, so it keeps
+    going straight to ``/datasets``.
+    """
+    return bool(request.sources) and all(
+        source.kind == "sample_store" for source in request.sources
+    )
+
+
+def _create_dataset_workset_first(client, request: DatasetCreateRequest):
+    """Create a dataset the canonical way: build the workset, then convert it.
+
+    The spec file keeps its shape, but a composition in it is no longer POSTed to
+    ``/datasets`` directly — the CLI selects it as a workset first and keeps that
+    selection with ``convert_workset_to_dataset``, which is the path a user is told
+    to take by hand.
+
+    Three details make the round-trip lossless. A spec that already names
+    ``from_workset`` is a convert request as written, so it passes straight through.
+    The workset carries the composition only, so ``shard_plan.sort_by`` is held back
+    and applied at convert time — a workset does not order its shards, and its
+    validator rejects the field. And a composition a workset cannot express falls
+    back to the direct create (see :func:`_is_workset_expressible`).
+    """
+    if request.from_workset is not None:
+        return client.convert_workset_to_dataset(
+            request.from_workset,
+            request.name,
+            shard_plan=request.shard_plan,
+            key_groups=request.key_groups or None,
+            scalars=request.scalars or None,
+            ttl_days=request.ttl_days,
+        )
+    if not _is_workset_expressible(request):
+        return client.create_dataset(request)
+
+    selection_plan = request.shard_plan.model_copy(update={"sort_by": []})
+    workset = client.create_workset(
+        WorksetCreateRequest(
+            sources=request.sources,
+            key_groups=request.key_groups,
+            shard_plan=selection_plan,
+        )
+    )
+    client.wait_for_workset(workset.workset_id)
+    return client.convert_workset_to_dataset(
+        workset.workset_id,
+        request.name,
+        shard_plan=request.shard_plan,
+        key_groups=request.key_groups or None,
+        scalars=request.scalars or None,
+        ttl_days=request.ttl_days,
+    )
+
+
+def handle_dataset_command(args) -> int:
+    output_format = getattr(args, "format", "plain")
+    with build_client() as client:
+        if args.dataset_command == "preview":
+            response = client.preview_dataset(load_model(args.spec, DatasetPreviewRequest))
+            render_output(response, output_format=output_format)
+            return 0
+        if args.dataset_command == "create":
+            # Creation is worker-driven for every source kind: the POST returns a
+            # `planning` dataset and the worker materializes it asynchronously.
+            # --wait polls to `ready` for callers that want the materialized dataset.
+            response = _create_dataset_workset_first(
+                client, load_model(args.spec, DatasetCreateRequest)
+            )
+            if getattr(args, "wait", False):
+                response = client.wait_for_dataset(response.name)
+            render_output(response, output_format=output_format)
+            return 0
+        if args.dataset_command == "list":
+            response = client.list_datasets(
+                query=args.query,
+                keys=args.keys,
+                include_archived=args.include_archived,
+                include_split_children=args.include_split_children,
+                offset=args.offset, limit=args.limit,
+            )
+            render_output(response, output_format=output_format)
+            return 0
+        if args.dataset_command == "get":
+            response = client.get_dataset(args.name)
+            render_output(response, output_format=output_format)
+            return 0
+        if args.dataset_command == "archive":
+            response = client.archive_dataset(args.name)
+            render_output(response, output_format=output_format)
+            return 0
+        if args.dataset_command == "restore":
+            response = client.restore_dataset(args.name)
+            render_output(response, output_format=output_format)
+            return 0
+        if args.dataset_command == "delete":
+            response = client.delete_dataset(args.name)
+            render_output(response, output_format=output_format)
+            return 0
+        if args.dataset_command == "add-key-group":
+            response = client.add_key_group(args.name, load_model(args.spec, AddKeyGroupRequest))
+            render_output(response, output_format=output_format)
+            return 0
+        if args.dataset_command == "pull":
+            output_dir = Path(args.output_dir)
+            dataset = client.get_dataset(args.name)
+            key_groups = args.key_groups or [group.name for group in dataset.key_groups]
+            key_group_by_name = {group.name: group for group in dataset.key_groups}
+            all_shards = []
+            for key_group in key_groups:
+                group = key_group_by_name.get(key_group)
+                if group is None:
+                    raise RuntimeError(f"dataset {args.name} does not define key group {key_group}")
+                workset = client.create_workset(
+                    # retire-workset-scope: a dataset read is a `dataset_scope` source.
+                    WorksetCreateRequest(
+                        sources=[
+                            SourceSpec(
+                                kind="dataset_scope",
+                                dataset_name=args.name,
+                                keys=list(group.keys),
+                            )
+                        ],
+                        shard_plan=ShardPlan(shard_size=100),
+                    )
+                )
+                client.wait_for_workset(workset.workset_id)
+                batch_ordinal = 0
+                while True:
+                    opened = client.open_workset_stream(
+                        workset.workset_id,
+                        WorksetOpenRequest(batch_ordinal=batch_ordinal, limit=200),
+                    )
+                    if not opened.shards:
+                        break
+                    for shard_ref in opened.shards:
+                        issued = client.issue_stream_tickets(
+                            workset.workset_id,
+                            WorksetIssueStreamTicketsRequest(
+                                shards=[
+                                    WorksetShardSelector(
+                                        shard_index=shard_ref.shard_index,
+                                        key_group=shard_ref.key_group,
+                                    )
+                                ],
+                            ),
+                        )
+                        if len(issued.tickets) != 1:
+                            raise RuntimeError(f"Datavo did not issue a stream ticket for shard {shard_ref.shard_index}")
+                        shard = issued.tickets[0]
+                        if shard.content_hash != shard_ref.content_hash:
+                            raise RuntimeError(
+                                f"Datavo stream ticket for shard {shard_ref.shard_index} has unexpected content hash"
+                            )
+                        filename = f"{shard_ref.shard_index:05d}-{shard_ref.content_hash}.tar"
+                        destination = output_dir / key_group / filename
+                        client.download_file(shard.download_url, destination, overwrite=args.overwrite)
+                        all_shards.append(
+                            {
+                                "key_group": key_group,
+                                **shard.model_dump(mode="json", by_alias=True),
+                            }
+                        )
+                    batch_ordinal += 1
+            render_output(
+                {
+                    "dataset_name": args.name,
+                    "shards": all_shards,
+                },
+                output_format=output_format,
+            )
+            return 0
+    raise RuntimeError("Missing dataset subcommand.")

datavo_cli/commands/events.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from datavo_sdk import DatavoEventCreateRequest
+
+from ..client_factory import build_client
+from ..render import render_output
+
+
+def _parse_params(values: list[str]) -> dict[str, str]:
+    params: dict[str, str] = {}
+    for value in values:
+        if "=" not in value:
+            raise ValueError("--param values must use key=value")
+        key, param_value = value.split("=", 1)
+        key = key.strip()
+        if not key:
+            raise ValueError("--param keys must not be empty")
+        params[key] = param_value
+    return params
+
+
+def handle_events_command(args) -> int:
+    output_format = getattr(args, "format", "plain")
+    with build_client() as client:
+        if args.events_command == "create":
+            response = client.create_event(
+                DatavoEventCreateRequest(
+                    event_type=args.event_type,
+                    source=args.source,
+                    jobpool_run_id=args.jobpool_run_id,
+                    idempotency_key=args.idempotency_key,
+                    params=_parse_params(args.param),
+                )
+            )
+            render_output(response, output_format=output_format)
+            return 0
+        if args.events_command == "list":
+            response = client.list_events(
+                event_type=args.event_type,
+                param_filters=_parse_params(args.param),
+            )
+            render_output(response, output_format=output_format)
+            return 0
+        if args.events_command == "get":
+            response = client.get_event(args.event_id)
+            render_output(response, output_format=output_format)
+            return 0
+    raise RuntimeError("Missing events subcommand.")