arraylake 1.3.0__tar.gz → 1.5.0__tar.gz

{arraylake-1.3.0 → arraylake-1.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: arraylake
-Version: 1.3.0
+Version: 1.5.0
 Summary: Python client for ArrayLake
 Project-URL: Documentation, https://docs.earthmover.io/
 Project-URL: Changelog, https://docs.earthmover.io/changelog
@@ -14,7 +14,7 @@ Requires-Python: <3.15,>=3.12
 Requires-Dist: click<9.0,>=8.1.3
 Requires-Dist: donfig<1.0,>=0.7
 Requires-Dist: httpx<0.29,>=0.23
-Requires-Dist: icechunk<3.0.0,>=2.1.2
+Requires-Dist: icechunk<3.0.0,>=2.2.2
 Requires-Dist: numpy<3.0,>=1.23
 Requires-Dist: packaging<27.0,>=23.0
 Requires-Dist: pydantic[email]<3.0,>=2.9
@@ -24,15 +24,15 @@ Requires-Dist: ruamel-yaml<1.0,>=0.17
 Requires-Dist: structlog<27.0,>=24.1
 Requires-Dist: typer<1.0,>=0.12.0
 Requires-Dist: zarr<4.0,>=3.1.0
+Provides-Extra: iceberg
+Requires-Dist: pyiceberg<0.13,>=0.11; extra == 'iceberg'
 Provides-Extra: maximal
-Requires-Dist: cf-xarray>=0.10.4; extra == 'maximal'
 Requires-Dist: xarray[accel,etc,io,parallel]>=2024.10.0; extra == 'maximal'
 Provides-Extra: obstore
 Requires-Dist: obstore>=0.9.2; extra == 'obstore'
 Provides-Extra: openeo
-Requires-Dist: openeo<0.40,>=0.35.0; extra == 'openeo'
+Requires-Dist: openeo<0.52,>=0.35.0; extra == 'openeo'
 Provides-Extra: xarray
-Requires-Dist: cf-xarray>=0.10.4; extra == 'xarray'
 Requires-Dist: xarray>=2024.10.0; extra == 'xarray'
 Description-Content-Type: text/markdown
 

{arraylake-1.3.0 → arraylake-1.5.0}/arraylake/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '1.3.0'
-__version_tuple__ = version_tuple = (1, 3, 0)
+__version__ = version = '1.5.0'
+__version_tuple__ = version_tuple = (1, 5, 0)
 
 __commit_id__ = commit_id = None

{arraylake-1.3.0 → arraylake-1.5.0}/arraylake/api_utils.py

@@ -135,10 +135,20 @@ class ClientKey:
 # the global cache of httpx clients
 _GLOBAL_CLIENTS: dict[ClientKey, httpx.AsyncClient] = {}
 
+# one lock per event loop, so concurrent callers on a cold cache build one client, not two
+_CLIENT_CREATION_LOCKS: dict[asyncio.AbstractEventLoop, asyncio.Lock] = {}
+
 # this is a cache to use hold asyncio tasks so they are not garbage collected before finishing
 background_tasks: set[asyncio.Task] = set()
 
 
+def _creation_lock(loop: asyncio.AbstractEventLoop) -> asyncio.Lock:
+    lock = _CLIENT_CREATION_LOCKS.get(loop)
+    if lock is None:
+        lock = _CLIENT_CREATION_LOCKS[loop] = asyncio.Lock()
+    return lock
+
+
 async def get_client(
     loop: asyncio.AbstractEventLoop,
     api_url: str,
@@ -181,7 +191,13 @@ async def get_client(
         proxy,
     )
     await logger.adebug("%i httpx clients present in cache.", len(_GLOBAL_CLIENTS))
-    if key not in _GLOBAL_CLIENTS:
+    if (client := _GLOBAL_CLIENTS.get(key)) is not None:
+        await logger.adebug("Client already present %s. Loop id %s.", key, id(loop))
+        return client
+
+    async with _creation_lock(loop):
+        if (client := _GLOBAL_CLIENTS.get(key)) is not None:
+            return client
         await logger.adebug("Creating new httpx client %s. Loop id %s.", key, id(loop))
         transport = AsyncRetryTransport(
             max_keepalive_connections=max_keepalive_connections,
@@ -199,15 +215,12 @@ async def get_client(
             verify=verify_ssl,
             cert=ssl_cafile,
         ).__aenter__()
-        weakref.finalize(client, close_client, key)
+        weakref.finalize(client, close_client, key, id(client))
         _GLOBAL_CLIENTS[key] = client
-    else:
-        await logger.adebug("Client already present %s. Loop id %s.", key, id(loop))
-    assert key in _GLOBAL_CLIENTS
-    return _GLOBAL_CLIENTS[key]
+        return client
 
 
-def close_client(key: ClientKey):
+def close_client(key: ClientKey, client_id: int) -> None:
     """
     This is a finalizer function that is called when a global client is
     garbage collected. It cleanly closes the client for the specified key.
@@ -224,11 +237,13 @@ def close_client(key: ClientKey):
     Note: logging in this function runs the risk of conflicting with pytest#5502. For
     this reason, we have removed debug log statements.
     """
-    client = _GLOBAL_CLIENTS.pop(key, None)
-
-    if client is None:
+    cached = _GLOBAL_CLIENTS.get(key)
+    # A newer client can own this key. Only the client that registered this finalizer may evict it.
+    if cached is None or id(cached) != client_id:
         return
 
+    client = _GLOBAL_CLIENTS.pop(key)
+
     if client.is_closed:
         return
 

{arraylake-1.3.0 → arraylake-1.5.0}/arraylake/cli/compute.py

@@ -28,7 +28,16 @@ def _services_table(services, org):
     table.add_column("URL", justify="right", style="green", min_width=10)
     table.add_column("Status", justify="right", style="green", min_width=10)
 
-    status_colors = {"available": "green", "progressing": "yellow", "error": "red", "unknown": "grey"}
+    status_colors = {
+        "available": "green",
+        "progressing": "yellow",
+        "error": "red",
+        "quota_exceeded": "red",
+        "sleeping": "cyan",
+        "stopped": "grey",
+        "suspended": "grey",
+        "unknown": "grey",
+    }
 
     for service in services:
         table.add_row(
@@ -38,7 +47,9 @@ def _services_table(services, org):
             str(service.config.is_public),
             service.url,
             service.status,
-            style=status_colors[service.status],
+            # .get, because the API owns this vocabulary and adds to it: an older client must
+            # render a status it has never heard of, not refuse to list anything at all.
+            style=status_colors.get(service.status, "grey"),
         )
 
     return table
@@ -47,7 +58,7 @@ def _services_table(services, org):
 @app.command(name="enable")
 def enable_service(
     org: str = typer.Argument(..., help="The organization to enable the service for"),
-    protocol: str = typer.Argument(..., help="The protocol to enable. Must be one of: dap2, edr, wms, tiles, zarr"),
+    protocol: str = typer.Argument(..., help="The protocol to enable. Must be one of: dap2, edr, wms, tiles, zarr, openeo, sql"),
     sku_code: str = typer.Option(..., "--sku-code", "-s", help="The compute shape to run the service on, e.g. xs"),
     is_public: bool = typer.Option(False, help="Whether the service is public"),
     deployment_name: str | None = typer.Option(

arraylake-1.5.0/arraylake/cli/iceberg.py

@@ -0,0 +1,255 @@
+"""Namespace administration and table discovery for Arraylake's Iceberg catalog."""
+
+import json
+import re
+from enum import StrEnum
+from typing import Any
+
+import typer
+from rich.table import Table
+
+from arraylake import Client
+from arraylake.cli.utils import rich_console
+from arraylake.types import RepoVisibility
+
+app = typer.Typer(help="Manage Iceberg namespaces and discover tables", no_args_is_help=True)
+namespace_app = typer.Typer(help="Manage Iceberg namespaces", no_args_is_help=True)
+table_app = typer.Typer(help="Discover and describe Iceberg tables", no_args_is_help=True)
+app.add_typer(namespace_app, name="namespace")
+app.add_typer(table_app, name="table")
+
+
+class Output(StrEnum):
+    rich = "rich"
+    json = "json"
+
+
+class NamespaceSort(StrEnum):
+    relevance = "relevance"
+    updated = "updated"
+    created = "created"
+    name = "name"
+    tables = "tables"
+
+
+class TableSort(StrEnum):
+    relevance = "relevance"
+    updated = "updated"
+    created = "created"
+    name = "name"
+
+
+class Direction(StrEnum):
+    asc = "asc"
+    desc = "desc"
+
+
+def _table_identifier(namespace: str, table: str) -> str:
+    """Format exact names; quoted components escape double quotes by doubling them."""
+
+    def quote(name: str) -> str:
+        return name if re.fullmatch(r"[A-Za-z0-9_-]+", name) else '"' + name.replace('"', '""') + '"'
+
+    return f"{quote(namespace)}.{quote(table)}"
+
+
+def _parse_table_identifier(identifier: str) -> tuple[str, str]:
+    # Keep this grammar in sync with the app's iceberg-identifier.ts formatter.
+    component = r'(?:[A-Za-z0-9_-]+|"(?:[^"]|"")+")'
+    match = re.fullmatch(rf"({component})\.({component})", identifier)
+    if match is None:
+        raise typer.BadParameter(
+            "Expected NAMESPACE.TABLE. Double-quote names containing dots or special characters, "
+            "and wrap the identifier in single quotes in your shell, e.g. '\"namespace.v2\".table'.",
+            param_hint="IDENTIFIER",
+        )
+    namespace, table = (part[1:-1].replace('""', '"') if part.startswith('"') else part for part in match.groups())
+    return namespace, table
+
+
+def _json_object(value: str | None, option: str) -> dict[str, Any] | None:
+    if value is None:
+        return None
+    try:
+        result = json.loads(value)
+    except ValueError as exc:
+        raise typer.BadParameter("Expected a JSON object", param_hint=option) from exc
+    if not isinstance(result, dict):
+        raise typer.BadParameter("Expected a JSON object", param_hint=option)
+    return result
+
+
+def _show(value: Any, output: Output) -> None:
+    if output == Output.json:
+        typer.echo(json.dumps(value))
+    else:
+        rich_console.print(value)
+
+
+def _show_list(rows: list[dict[str, Any]], output: Output, *, namespaces: bool) -> None:
+    if not namespaces:
+        rows = [{**row, "identifier": _table_identifier(row["namespace"], row["name"])} for row in rows]
+    if output == Output.json:
+        _show(rows, output)
+        return
+    table = Table("Name" if namespaces else "Identifier", "Updated", "Tables" if namespaces else "Storage override")
+    for row in rows:
+        table.add_row(
+            row["name"] if namespaces else row["identifier"],
+            row["updated"],
+            str(row["table_count"] if namespaces else row["has_storage_override"]),
+        )
+    rich_console.print(table)
+
+
+@namespace_app.command("list")
+def list_namespaces(
+    org: str,
+    search: str | None = None,
+    filter_metadata: str | None = typer.Option(None, help="JSON metadata filter"),
+    sort: NamespaceSort | None = None,
+    direction: Direction = Direction.desc,
+    include_ghosts: bool = False,
+    page_size: int = typer.Option(50, min=1, max=100),
+    output: Output = Output.rich,
+):
+    """List namespaces, with optional search, metadata filters, and sorting."""
+    rows = Client().list_iceberg_namespaces_paginated(
+        org,
+        search=search,
+        filter_metadata=_json_object(filter_metadata, "--filter-metadata"),
+        sort=sort.value if sort else None,
+        direction=direction.value,
+        include_ghosts=include_ghosts,
+        page_size=page_size,
+    )
+    _show_list([row.model_dump(mode="json") for row in rows], output, namespaces=True)
+
+
+@namespace_app.command("get")
+def get_namespace(org: str, name: str, output: Output = Output.rich):
+    """Show namespace storage, description, metadata, and visibility."""
+    _show(Client().get_iceberg_namespace(org, name).model_dump(mode="json"), output)
+
+
+@namespace_app.command("create")
+def create_namespace(
+    org: str,
+    name: str,
+    bucket_config_nickname: str | None = None,
+    description: str | None = None,
+    metadata: str | None = typer.Option(None, help="JSON discovery metadata"),
+    visibility: RepoVisibility = RepoVisibility.PRIVATE,
+    properties: str | None = typer.Option(None, help="JSON Iceberg properties (string values)"),
+    output: Output = Output.rich,
+):
+    """Create a namespace on the chosen or default bucket."""
+    props = _json_object(properties, "--properties")
+    if props is not None and any(not isinstance(v, str) for v in props.values()):
+        raise typer.BadParameter("Property values must be strings", param_hint="--properties")
+    result = Client().create_iceberg_namespace(
+        org,
+        name,
+        bucket_config_nickname=bucket_config_nickname,
+        description=description,
+        metadata=_json_object(metadata, "--metadata"),
+        visibility=visibility,
+        properties=props,
+    )
+    _show(result.model_dump(mode="json"), output)
+
+
+@namespace_app.command("update")
+def update_namespace(
+    org: str,
+    name: str,
+    description: str | None = None,
+    add_metadata: str | None = typer.Option(None, help="JSON metadata to add"),
+    remove_metadata: list[str] | None = typer.Option(None, help="Metadata key to remove; repeat for multiple keys"),
+    update_metadata: str | None = typer.Option(None, help="JSON metadata to update"),
+    visibility: RepoVisibility | None = None,
+    output: Output = Output.rich,
+):
+    """Update description, discovery metadata, or visibility. Use --description '' to clear."""
+    result = Client().modify_iceberg_namespace(
+        org,
+        name,
+        description=description,
+        add_metadata=_json_object(add_metadata, "--add-metadata"),
+        remove_metadata=remove_metadata,
+        update_metadata=_json_object(update_metadata, "--update-metadata"),
+        visibility=visibility,
+    )
+    _show(result.model_dump(mode="json"), output)
+
+
+@namespace_app.command("delete")
+def delete_namespace(
+    org: str,
+    name: str,
+    confirm: bool = typer.Option(False, "--confirm", help="Skip the confirmation prompt"),
+    immediate: bool = typer.Option(
+        False, "--immediate", help="Skip the recovery period: the namespace cannot be restored and its name is reusable at once"
+    ),
+    retain_data: bool = typer.Option(False, "--retain-data", help="Leave the bytes under the namespace's storage prefix in the bucket"),
+):
+    """Delete an empty namespace. A soft-delete by default: restore it during its recovery period."""
+    if not confirm:
+        verb = "Permanently delete" if immediate else "Delete"
+        typer.confirm(f"{verb} namespace {org}/{name}?", abort=True)
+    Client().delete_iceberg_namespace(org, name, imsure=True, immediate=immediate, retain_data=retain_data)
+    typer.echo(f"Deleted {org}/{name}")
+
+
+@namespace_app.command("restore")
+def restore_namespace(org: str, name: str):
+    """Restore a soft-deleted namespace."""
+    Client().restore_iceberg_namespace(org, name)
+    typer.echo(f"Restored {org}/{name}")
+
+
+@table_app.command("list")
+def list_tables(
+    org: str,
+    namespace: str,
+    search: str | None = None,
+    sort: TableSort | None = None,
+    direction: Direction = Direction.desc,
+    page_size: int = typer.Option(50, min=1, max=100),
+    output: Output = Output.rich,
+):
+    """List table summaries without fetching table metadata or recording accesses."""
+    rows = Client().list_iceberg_tables_paginated(
+        org,
+        namespace,
+        search=search,
+        sort=sort.value if sort else None,
+        direction=direction.value,
+        page_size=page_size,
+    )
+    _show_list([row.model_dump(mode="json") for row in rows], output, namespaces=False)
+
+
+@table_app.command("describe")
+def describe_table(
+    org: str,
+    identifier: str = typer.Argument(..., help="NAMESPACE.TABLE; double-quote components containing dots or special characters."),
+    output: Output = Output.rich,
+):
+    """Load a table's schema, snapshots, partitioning, and properties. Requires arraylake[iceberg]."""
+    namespace, table = _parse_table_identifier(identifier)
+    catalog = Client().get_iceberg(org)
+    try:
+        loaded = catalog.load_table((namespace, table))
+        _show(
+            {
+                "name": table,
+                "namespace": namespace,
+                "identifier": _table_identifier(namespace, table),
+                "metadata_location": loaded.metadata_location,
+                "metadata": loaded.metadata.model_dump(mode="json", by_alias=True),
+            },
+            output,
+        )
+    finally:
+        catalog.close()

{arraylake-1.3.0 → arraylake-1.5.0}/arraylake/cli/main.py

@@ -6,13 +6,19 @@ from arraylake.api_utils import set_client_name
 from arraylake.cli.auth import auth
 from arraylake.cli.compute import app as compute_app
 from arraylake.cli.config import app as config_app
+from arraylake.cli.iceberg import app as iceberg_app
 from arraylake.cli.org import app as org_app
 from arraylake.cli.repo import app as repo_app
 from arraylake.cli.skills import app as skills_app
 from arraylake.cli.utils import rich_console
 from arraylake.diagnostics import get_diagnostics
 
-app = typer.Typer(no_args_is_help=True, context_settings={"help_option_names": ["-h", "--help"]}, rich_markup_mode="markdown")
+app = typer.Typer(
+    no_args_is_help=True,
+    context_settings={"help_option_names": ["-h", "--help"]},
+    rich_markup_mode="markdown",
+    pretty_exceptions_show_locals=False,
+)
 
 
 def version_callback(value: bool):
@@ -48,6 +54,7 @@ def main(
 app.add_typer(auth, name="auth")
 app.add_typer(org_app, name="org")
 app.add_typer(repo_app, name="repo")
+app.add_typer(iceberg_app, name="iceberg")
 app.add_typer(config_app, name="config")
 app.add_typer(compute_app, name="compute")
 app.add_typer(skills_app, name="skills")

{arraylake-1.3.0 → arraylake-1.5.0}/arraylake/cli/repo.py

@@ -1,31 +1,19 @@
 import json
-import warnings
 from enum import Enum, StrEnum
-from typing import TYPE_CHECKING, Optional
 from uuid import UUID
 
 import typer
 import zarr
 from pydantic import SecretStr
 from rich import print_json
-from rich.prompt import Confirm
 from rich.table import Table
 from rich.text import Text
 
-from arraylake import AsyncClient, Client
+from arraylake import AsyncClient
 from arraylake.cli.utils import coro, rich_console, simple_progress
-from arraylake.compute.doctor import (
-    CFDiagnosisSummary,
-    check_cf_attribute_completeness,
-    fix_cf_noncompliance,
-)
 from arraylake.log_util import get_logger
 from arraylake.types import RepoOperationMode
 
-if TYPE_CHECKING:
-    import xarray as xr  # noqa: F401
-    from xarray import Dataset
-
 app = typer.Typer(help="Manage Arraylake repositories", no_args_is_help=True)
 logger = get_logger(__name__)
 
@@ -141,7 +129,9 @@ async def list_repos(
 async def create(
     repo_name: str = typer.Argument(..., help="Name of repository {ORG}/{REPO_NAME}"),
     bucket_config_nickname: str | None = typer.Option(None, help="Bucket config nickname"),
-    prefix: str = typer.Option(None, help="Optional prefix for Icechunk repo. If not provided, a random ID + the repo name will be used."),
+    prefix: str = typer.Option(
+        None, help="Optional prefix for the repo's storage. If not provided, a random ID + the repo name will be used."
+    ),
     description: str | None = typer.Option(None, help="Description of the repo"),
     metadata: str | None = typer.Option(None, help="Optional metadata for the repo"),
 ):
@@ -315,152 +305,6 @@ async def tree(
         print(_tree)
 
 
-def diagnose_dataset(ds: "Dataset") -> CFDiagnosisSummary:
-    """Checks if a dataset is CF compliant."""
-    rich_console.print("[bold]Diagnosing dataset...[/bold] :stethoscope:")
-    # TODO: pause the spinner while this is running
-    diagnosis: CFDiagnosisSummary = check_cf_attribute_completeness(ds)
-
-    rich_console.print("[bold]Checking compatibilty with compute services...[/bold] :stethoscope:")
-    if diagnosis.compatible_services:
-        for service in diagnosis.compatible_services:
-            rich_console.print(f"  :white_check_mark: {service}")
-    if diagnosis.incompatible_services:
-        for service in diagnosis.incompatible_services:
-            rich_console.print(f"  :x: {service}")
-
-    if diagnosis.is_healthy:
-        # Dataset is already compliant
-        rich_console.print("Dataset is healthy! :apple:")
-    else:
-        rich_console.print("[bold]Dataset is unhealthy[/bold] :face_with_thermometer:")
-        rich_console.print("[bold]Diagnosis:[/bold] :woman_health_worker::clipboard:")
-        if diagnosis.missing_required_keys:
-            rich_console.print("  Missing CF Attributes:")
-            for missing_key in diagnosis.sorted_missing_keys:
-                missing_required = " (required)" if missing_key.required else ""
-                color = "default"
-                if missing_key.required:
-                    color = "red"
-                elif missing_key.proposed_variable:
-                    color = "yellow"
-                rich_console.print(f"    [bold][{color}]- {missing_key.attr_name}: {missing_key.name}{missing_required}[/{color}][/bold]")
-                if missing_key.proposed_variable:
-                    rich_console.print(f"      Possible variable: {missing_key.proposed_variable}")
-        if diagnosis.has_invalid_keys:
-            rich_console.print("  Invalid CF Attributes:")
-            for invalid_key in diagnosis.invalid:
-                rich_console.print(f"    [bold][red]- {invalid_key.attr_name}: {invalid_key.name}: {invalid_key.issue}[/red][/bold]")
-    return diagnosis
-
-
-def treat_dataset(diagnosis: CFDiagnosisSummary, store, group) -> tuple[CFDiagnosisSummary, Optional["Dataset"]]:
-    """Treats a dataset to make it CF compliant."""
-    if diagnosis.is_healthy:
-        rich_console.print("Dataset does not need to be treated :muscle:")
-        return diagnosis, None
-    else:
-        rich_console.print("[bold]Treating dataset... [/bold]:ambulance:")
-        fix_cf_noncompliance(diagnosis, store, group)
-        # Read the data back in
-        import xarray as xr  # noqa: F811
-
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            treated_ds = xr.open_zarr(store, zarr_format=3, group=group, consolidated=False)
-        # Run the compliance check again to see if we actually fixed it
-        new_diagnosis = check_cf_attribute_completeness(treated_ds)
-        if not new_diagnosis.is_healthy:
-            raise ValueError(
-                "[bold]Failed to treat dataset![/bold] :coffin: \n "
-                "Please use the diagnosis to manually update the necessary attributes :hammer:"
-            )
-        return new_diagnosis, treated_ds
-
-
-def get_treatment_plan(og_ds: "Dataset", treated_ds: "Dataset"):
-    """Prints the treatment plan for a dataset by comparing variable attrs."""
-    all_vars = set(og_ds.variables).union(treated_ds.variables)
-
-    for var in all_vars:
-        og_attrs = og_ds[var].attrs if var in og_ds else {}
-        treated_attrs = treated_ds[var].attrs if var in treated_ds else {}
-
-        if og_attrs != treated_attrs:
-            rich_console.print(f"\n  [bold]Treatment plan for variable '{var}' metadata: [/bold]")
-
-            for key in set(og_attrs.keys()).union(treated_attrs.keys()):
-                og_val = og_attrs.get(key, "MISSING")
-                new_val = treated_attrs.get(key, "MISSING")
-
-                if og_val != new_val:
-                    color = "green" if key in treated_attrs else "red"
-                    rich_console.print(f"    - [bold][{color}]{key}:[/{color}][/bold]")
-                    rich_console.print(f"      [bold]Original:[/bold] [red]{og_val}[/red]")
-                    rich_console.print(f"      [bold]Treated:[/bold] [green]{new_val}[/green]")
-
-
-@app.command(name="doctor")
-@coro  # type: ignore
-async def doctor_dataset(
-    repo_name: str = typer.Argument(..., help="Name of repository {ORG}/{REPO_NAME}"),
-    group: str | None = typer.Option(None, help="The path to the zarr group"),
-    treat: bool = typer.Option(default=False, hidden=True, help="Fix and commit the changes automatically"),
-    dry_run: bool = typer.Option(default=False, hidden=True, help="Run the doctor without committing changes"),
-):
-    """**Doctor** a dataset to make it usable by the Arraylake compute engine.
-
-    **Examples**
-
-
-    """
-    if not treat and dry_run:
-        rich_console.print("[bold][yellow]Warning: [/bold][/yellow] Treat must be enabled for dry run mode.")
-
-    group_str = group if group is not None else ""
-    rich_console.print(f"Admitting dataset [bold]{repo_name}/{group_str}[/bold] to the Arraylake hospital :hospital:")
-    client = Client()
-    repo = client.get_repo(repo_name)
-
-    session = repo.writable_session(branch="main") if treat else repo.readonly_session(branch="main")
-    store = session.store
-
-    # with ekg_progress(""):
-    import xarray as xr  # noqa: F811
-
-    with warnings.catch_warnings():
-        warnings.simplefilter("ignore")
-        ds = xr.open_zarr(store, zarr_format=3, group=group, consolidated=False)
-    diagnosis = diagnose_dataset(ds)
-
-    if treat:
-        updated_diagnosis, treated_ds = treat_dataset(diagnosis, store, group=group)
-        # If the dataset was already healthy, treatment is not needed
-        if treated_ds is None:
-            return
-        # Check that the only the dataset was modified using session status
-        diff = session.status()
-        assert len(diff.new_groups) == 0
-        assert len(diff.new_arrays) == 0
-        assert len(diff.deleted_groups) == 0
-        assert len(diff.deleted_arrays) == 0
-        assert len(diff.updated_chunks) == 0
-        assert diff.updated_arrays  # type: ignore[unused-ignore,attr-defined]
-        rich_console.print("[bold]Proposed treatment:[/bold] :pill:")
-        get_treatment_plan(ds, treated_ds)
-
-        if updated_diagnosis.is_healthy and not dry_run:
-            lets_commit = Confirm.ask("Do you want to commit the proposed changes?")
-            if lets_commit:
-                rich_console.print("[bold]Committing changes...[/bold] :sparkles:")
-                message = "Updated metadata attributes for use with Arraylake Flux using repo doctor"
-                snapshot_id = session.commit(message)
-                rich_console.print(f"Dataset successfully treated and committed to the repo with ID {snapshot_id} :dizzy:")
-    else:
-        if not diagnosis.is_healthy:
-            rich_console.print("\nTo manually treat the dataset, use the diagnosis provided above to update the dataset metadata")
-
-
 @app.command(hidden=True)
 @coro  # type: ignore
 async def get_status(