diff --git a/docs/adr/11-code-execution.md b/docs/adr/11-code-execution.md index dfb168e..8b361a7 100644 --- a/docs/adr/11-code-execution.md +++ b/docs/adr/11-code-execution.md @@ -2,7 +2,7 @@ ## Status -Accepted; the containment mechanism revised once. (2026-08) The sandbox was a profile pointing at the user's own database through a read-only role. It is now a disposable **copy**, after the original cost a maintainer his database ([#73](https://github.com/aiidateam/aiida-agents/issues/73)). The original reasoning is kept below, marked, because the mistake in it is the instructive part. +Accepted; revised twice, both times by using the thing. (2026-08) The sandbox was a profile pointing at the user's own database through a read-only role; it became a disposable **copy** after the original cost a maintainer his database ([#73](https://github.com/aiidateam/aiida-agents/issues/73)); and the copy is now a **scratch profile the agent may write to**, which is what a disposable copy was always for. Each superseded version is kept below, marked, because the mistakes in them are the instructive part. ## Context @@ -16,7 +16,7 @@ That raises the obvious objection. Executing model-written Python against a rese **Generated code runs against a disposable copy of the user's storage.** -`aiida-agents sandbox init` copies the profile's storage and registers a profile pointing at the copy; `sandbox check` verifies the copy shares no database and no repository with any real profile; `refresh` rebuilds it when it has drifted, and `teardown` removes it. A read-only PostgreSQL role over the copy is available as a second layer, and is no longer the mechanism. +`aiida-agents sandbox init` copies the profile's storage and registers a profile pointing at the copy; `sandbox check` verifies the copy shares no database and no repository with any real profile; `teardown` removes it. A read-only PostgreSQL role over the copy is available as a second layer, and is no longer the mechanism. > **Revised (2026-08).** This originally read: *"Generated code runs against a profile whose PostgreSQL role holds no write privilege. The sandbox profile points at **the same database** as the user's own. A scratch profile would be safer and useless: an empty database cannot answer any question worth asking about someone's data."* > @@ -28,12 +28,28 @@ That raises the obvious objection. Executing model-written Python against a rese > > The rule that replaces it is one sentence: **a sandbox profile must never share deletable storage with a real one.** It lives in `sandbox/copy.py` as `shares_storage`, it fails closed, and `init`, `check`, `teardown` and `doctor` all ask it rather than each re-deciding. +**That rule is only worth its wording if the comparison is exact**, and the first implementation was not. Locations were tagged strings, so a directory and an archive at one path compared as *separate* whenever the two profiles disagreed about the kind of thing there; a `file://` repository URI had its scheme stripped by hand, and since `Path.as_uri()` percent-encodes, a repository under a path containing a space never matched the directory it named; and one storage nested inside another read as separate right up until `teardown` removed the parent recursively and took both. Locations are now `PathLocation` and `DatabaseLocation`, `shares_storage` takes two `ProfileStorage` values so a caller cannot pair one profile's backend with another's config, containment counts as overlap, and anything unreadable fails closed. Separation is proved twice: `init` refuses a layout that would copy a source into itself *before* writing anything, and `run_aiida_code` asks the same question again at run time, because the setting it trusts is a profile name and nothing stops it naming the user's own profile. + +**The copy is also a scratch profile, and writes belong in it.** + +> **Revised again (2026-08).** The consequence below originally read: *"Writes remain impossible from this path. A user who wants to submit is told so and routed to the Execution agent, which asks first."* +> +> That was inherited from the read-only role and was never true of the copy. The static guard blocks the write calls it knows; on SQLite nothing sits beneath it, so a call it does not recognise (`Group.collection.get_or_create` was the one found by trying it) succeeds against the copy. The model was being told writes were impossible when they were merely invisible, which is the worse of the two errors: it invites the agent to report having created something the user will never find. +> +> The list was briefly widened to cover the calls it missed, and that was reversed the same day. Blocking `store()` to protect a copy is protecting the wrong thing. A researcher iterating on inputs they are unsure of *wants* somewhere to be wrong five times, and five excepted workflows, a batch submitted with the wrong parameters, and a deletion that took too much all belong in a profile that is thrown away rather than in the one they do their work in. +> +> So the copy is where iteration happens. What the guard is still for is everything that leaves the machine. + +**The sandbox contains data, not actions.** The copy carries the `Computer` rows and the `AuthInfo` beside them, so a calculation submitted from it runs on the user's real cluster, under their credentials, spending their allocation, and leaves its remote work directory behind when the node is deleted. Deliberate, because inputs cannot be validated against a fake machine, and the sharpest limit of the whole design: the provenance is sandboxed and the compute is not. + +**Work done in the sandbox exists nowhere else**, which is what makes `refresh` a trap and why there is none. Getting it back out is a promote step that has not been built. The mechanism it wants is [`verdi collab`](https://github.com/aiidateam/aiida-core/pull/7516): cursors and UUID-manifest negotiation already answer "what is new here", sync is additive so a deletion in the sandbox cannot propagate, and the receiver decides what enters it. + Two further layers sit above it, and neither is containment: - A **static guard** rejects imports outside an allowlist, calls that write, and the builtins (`getattr`, `exec`, `open`) that would step around either rule. It turns the common mistakes into a readable message instead of a permission error. It is a pre-check and must never be relied on: a one-line bypass reaching `os` through an allowed module survived review and was found by dogfooding. - A **subprocess** with a timeout, a scrubbed environment (it inherited the user's API keys until #73), resource limits, and its own process group so a timeout takes with it whatever the snippet spawned. -**The execution tool is not approval-gated**, which is the decision most worth arguing with. It holds no capability to write, so there is nothing for an approval to protect. And the alternative is worse than it looks: an approval prompt showing twenty lines of unexecuted Python asks a researcher to audit code under time pressure, which is a far weaker check than the one it appears to be. Letting the code run where it can do no harm and showing what it *returned* converts the same decision into one made from evidence. +**The execution tool is not approval-gated**, which is the decision most worth arguing with, and which the second revision weakens: it once held no capability to write, so there was nothing for an approval to protect. What it now holds is the ability to write to a copy, which is still nothing to protect, and the ability to spend the user's allocation on their real cluster, which is not. Whether that stays ungated is an open decision rather than a settled one. And the alternative is worse than it looks: an approval prompt showing twenty lines of unexecuted Python asks a researcher to audit code under time pressure, which is a far weaker check than the one it appears to be. Letting the code run where it can do no harm and showing what it *returned* converts the same decision into one made from evidence. **Everything that genuinely writes stays where it was**: on the Execution agent, behind `requires_approval=True`, where the preview is a resolved input a researcher can judge in seconds. @@ -46,10 +62,14 @@ Two further layers sit above it, and neither is containment: - A question needing an unanticipated combination of filters is answerable without a new tool. - Answers come with output that was actually produced, and the agent fixes its own mistakes from real tracebacks before the user sees them. - **The feature is inert until someone runs `sandbox init`.** `run_aiida_code` refuses to run when no sandbox profile is configured, and says the snippet is unverified rather than falling back to the user's writable profile. Silently falling back is the one failure that would make all of the above worthless. -- **The copy drifts.** It is a snapshot, so anything the user has run since is missing until `sandbox refresh`. That is the price of not sharing their storage, and it is the right way round: a stale answer is visible and correctable, a destroyed database is not. -- **Copying is not free.** A large provenance database takes time and disk to copy, which is why the copy has a lifecycle (`init`/`refresh`/`teardown`) rather than being made per query. +- **The copy drifts.** It is a snapshot, so anything the user has run since is missing until it is rebuilt, which is `teardown` then `init`. That is the price of not sharing their storage, and it is the right way round: a stale answer is visible and correctable, a destroyed database is not. +- **Copying is not free.** A large provenance database takes time and disk to copy, which is why the copy has a lifecycle (`init`/`teardown`) rather than being made per query. `init` says what the copy will cost and asks before making it. +- **There is deliberately no `refresh`.** A word promising to bring the copy up to date would hide both the cost, which is the whole repository again, and the loss, which is anything the sandbox holds that the source does not. The version worth having syncs incrementally against the source, and that wants [`verdi collab`](https://github.com/aiidateam/aiida-core/pull/7516) underneath it. - **The containment covers the database and nothing else.** Code that gets past the guard can still read the filesystem and reach the network. The environment scrub and the rlimits narrow that; OS-level isolation (`bwrap`, `nsjail`, a container) would close it, and is not implemented. -- Writes remain impossible from this path. A user who wants to submit is told so and routed to the Execution agent, which asks first. +- ~~Writes remain impossible from this path.~~ **Revised (2026-08), see the Decision above.** They are impossible against the *user's* profile, which is the property that matters. Against the copy they are the point. +- **PKs collide once both sides write.** The copy preserves them, so for everything that existed at copy time sandbox `pk 12` is the user's `pk 12`; anything created afterwards diverges while reusing the same numbers on both sides. The agent must report UUIDs for what it creates, or send the user to a node that is not the one it meant. +- **The static guard is doing more than a pre-check should.** On SQLite nothing sits beneath it for database writes. Its own docstring says not to rely on it, and for that one category we currently do. +- **`copytree` takes no consistent snapshot.** A profile being written while it is copied can produce a torn one. aiida-core's `StorageBackend.backup` is the supported answer and is not a drop-in: it wants the source profile loaded and produces its own versioned folder layout. ## Alternatives considered @@ -62,3 +82,9 @@ Two further layers sit above it, and neither is containment: **A restricted interpreter instead of a database role.** Rejected: Python is not sandboxable in-process, and a guard that claimed to be one would be believed. **No execution: generate code and let the user run it.** Rejected as the whole product. It leaves the correctness problem entirely with the user, and correctness is the thing the feature exists to provide. + +**Keeping the sandbox read-only by naming every write.** Tried and reversed within a day (see the second revision). The list is unclosable in principle, since `get_or_create`, `add_nodes` and `base.extras.set` were all missing and nothing says they were the last; and closing it would forbid the iteration the copy exists to make cheap. + +**Building the sandbox from an archive instead of copying bytes.** Not adopted, and the strongest remaining alternative. `verdi archive create` reads through the ORM, so unlike `copytree` it snapshots consistently, and `include_authinfos` defaults to false, which would make reaching a cluster an explicit choice rather than something inherited. A `core.sqlite_zip` profile from that archive would refuse writes in the backend itself, which is the enforced read-only the original design claimed and never had; an import into a fresh `core.sqlite_dos` stays writable but costs a full export plus a full import and renumbers every PK. Worth having as a `--read-only` mode rather than as a replacement, since the zip cannot be the scratch profile. + +**Promoting work out by timestamp.** Rejected: `ctime` survives an import, so anything the agent imported into the sandbox keeps its original time and would be missed. A PK watermark recorded at `init` is exact and needs only the sandbox loaded; a UUID set difference is exact without stored state but needs both profiles open, and AiiDA gives one profile per process. diff --git a/docs/adr/README.md b/docs/adr/README.md index 00d0246..50374d2 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -27,11 +27,11 @@ Numbering follows build/dependency order, not chronology: **01–06** are the pa | [08](/docs/adr/08-human-in-the-loop-before-writes.md) | Enforced human-in-the-loop confirmation before any write/submit | | [09](/docs/adr/09-agent-orchestration.md) | Agent orchestration: a planner over two specialists; specialists are never wrapped | | [10](/docs/adr/10-plugin-extensibility.md) | Plugin extensibility through one `aiida_agents.plugins` entry point | -| [11](/docs/adr/11-code-execution.md) | Executing generated code against a write-refusing database role | +| [11](/docs/adr/11-code-execution.md) | Executing generated code against a disposable copy of the user's storage | | 11 | Agent-run provenance: persist agent decisions/traces in AiiDA's provenance graph (exploratory) | ADR-01 is in effect; ADR-02 and ADR-03 are seeds with direction confirmed (2026-05-22). -ADR-04 through ADR-10 are written; ADR-11 is still exploratory and has no record yet. +ADR-04 through ADR-11 are written; the exploratory agent-run-provenance entry still has no record. ADR-09 supersedes ADR-04's future-architecture table and settles the agent-to-agent question ADR-04 left open; ADR-06 and ADR-07 carry Revision sections where reality diverged from the original decision. For how the pieces fit together rather than why each was chosen, see [Architecture](/docs/architecture.md); to add your own, see [Extending](/docs/extending.md). diff --git a/src/aiida_agents/_settings.py b/src/aiida_agents/_settings.py index 6271514..04da626 100644 --- a/src/aiida_agents/_settings.py +++ b/src/aiida_agents/_settings.py @@ -194,10 +194,10 @@ class SandboxSettings(_Base): """Where and how generated code is executed (``AIIDA_AGENTS_*``).""" sandbox_profile: str = "agents-sandbox" - """AiiDA profile generated code runs against. Must be one whose database - role cannot write --- ``aiida-agents sandbox init`` creates one, and - ``sandbox check`` proves it. Nothing downstream can tell a read-only - profile from a writable one, so this setting is the whole safety boundary.""" + """AiiDA profile generated code runs against. Must be a disposable copy of + the user's storage --- ``aiida-agents sandbox init`` makes one, and + ``sandbox check`` proves it shares nothing. Nothing downstream can tell a + copy from the real thing, so this setting is the whole safety boundary.""" sandbox_timeout: float = Field(default=30.0, gt=0) """Seconds a snippet may run before it is killed. A query over a large diff --git a/src/aiida_agents/agents/codegen/__init__.py b/src/aiida_agents/agents/codegen/__init__.py index a8371db..e36f6d3 100644 --- a/src/aiida_agents/agents/codegen/__init__.py +++ b/src/aiida_agents/agents/codegen/__init__.py @@ -17,8 +17,8 @@ What it does *not* have is a write tool, and that is deliberate. Executing generated code sounds like the most dangerous thing in this project, and it would be if it ran against the user's live profile. It does not: it runs -against one whose database role has no INSERT privilege, so the danger is -handled at the database rather than by asking a user to vet Python at a prompt. +against a disposable copy of their storage, so the danger is handled by what +the code can reach rather than by asking a user to vet Python at a prompt. Approving twenty lines of unexecuted code is a weak check. Reading what the code actually returned is a real one --- which is why this agent is not gated and every tool that genuinely writes still is, on the Execution agent, where diff --git a/src/aiida_agents/agents/codegen/prompt.md b/src/aiida_agents/agents/codegen/prompt.md index 999a8ab..ee27018 100644 --- a/src/aiida_agents/agents/codegen/prompt.md +++ b/src/aiida_agents/agents/codegen/prompt.md @@ -43,11 +43,19 @@ output you did not see. user's actual database, so the results are their real data. Do not hedge them as "example" or "illustrative" values. -**You cannot write, and must not pretend otherwise.** The profile you run -against has no write privilege at the database level. If the user wants to -submit a workflow, import a structure or delete something, say that is the -Execution agent's job and that it will ask for their approval first. Never -report having submitted, stored or changed anything. +**Nothing you run reaches the user's own profile.** You run against a copy of +their storage, and writes are refused before the code runs. If one ever gets +past that, it changes the copy and nothing else: the user cannot see it, and it +is thrown away the next time the sandbox is rebuilt. So never report having +submitted, stored or changed anything, even if the output says you did. If the +user wants to submit a workflow, import a structure or delete something, say +that is the Execution agent's job and that it will ask for their approval +first. + +The copy protects their AiiDA storage, not their machine. The filesystem and +the network are real, so do not treat reading files or reaching the network as +harmless just because the profile is a copy: neither belongs in a snippet that +answers a question about provenance. ## Style diff --git a/src/aiida_agents/cli/doctor.py b/src/aiida_agents/cli/doctor.py index 9b0ca43..25d80ac 100644 --- a/src/aiida_agents/cli/doctor.py +++ b/src/aiida_agents/cli/doctor.py @@ -143,11 +143,13 @@ def _check_sandbox() -> _DiagnosticRow: sharing = profiles_sharing_storage(config, name) if sharing: + # No remedy named here: `check` is where the two cases are told + # apart, and only one of them is the user's to fix. return _DiagnosticRow( label, False, - f"shares storage with {', '.join(sharing)}; deleting it would " - "destroy that data. Rebuild with `aiida-agents sandbox refresh`", + "; ".join(entry.describe() for entry in sharing) + + "; `aiida-agents sandbox check` spells out what to do", ) return _DiagnosticRow(label, True, "own copy, shared with nothing") except Exception as exc: diff --git a/src/aiida_agents/cli/sandbox.py b/src/aiida_agents/cli/sandbox.py index 4186ae4..2eef7df 100644 --- a/src/aiida_agents/cli/sandbox.py +++ b/src/aiida_agents/cli/sandbox.py @@ -1,12 +1,20 @@ """``aiida-agents sandbox`` --- build and verify the profile generated code reads. -Four commands over one idea: the sandbox is a **disposable copy** of the user's -storage, never the storage itself. +Three commands over one idea: the sandbox is a **disposable copy** of the +user's storage, never the storage itself. ``init`` makes the copy and registers a profile for it. ``check`` proves the copy shares nothing with a real profile, which is the property that matters. -``refresh`` rebuilds it when it has drifted. ``teardown`` removes it, and can -do so safely precisely because nothing else is pointing at what it deletes. +``teardown`` removes it, and can do so safely precisely because nothing else is +pointing at what it deletes. + +There is deliberately no ``refresh``. Rebuilding is ``teardown`` then ``init``, +which is two commands that say what they do, where a single word promising to +bring the copy up to date would hide both the cost (the whole repository again) +and the loss (anything the sandbox holds that the source does not). The version +worth having is an incremental sync against the source rather than a fresh +copy, and that wants `verdi collab +`_ underneath it. The design this replaces pointed the sandbox at the same database through a read-only role, and cost a maintainer his data when he deleted the sandbox @@ -16,13 +24,29 @@ from __future__ import annotations +import contextlib import secrets import shutil +import typing as t +from collections.abc import Callable, Iterator +from pathlib import Path import rich_click as click +from rich.filesize import decimal +from rich.progress import ( + BarColumn, + DownloadColumn, + Progress, + SpinnerColumn, + TextColumn, + TimeElapsedColumn, +) from aiida_agents.cli.output import console +if t.TYPE_CHECKING: + from aiida_agents.sandbox.copy import ShellCommand + __all__ = ["sandbox"] #: Default name for the read-only role and the profile that uses it. Named for @@ -31,6 +55,141 @@ DEFAULT_PROFILE = "agents-sandbox" +def _agreed_to_copy(source: Path, target: Path, size: int, *, yes: bool) -> bool: + """Say what the copy costs, and get agreement unless ``yes``. + + The copy is the whole repository, which on a real profile runs to + gigabytes. Announcing it after the fact is how somebody finds out from + ``df``, so the size, both paths and the room left go out before anything is + written. + """ + anchor = target + while not anchor.exists(): + anchor = anchor.parent + + console.print(f"This copies [bold]{decimal(size)}[/bold]") + # `soft_wrap` so a path longer than the terminal breaks where the terminal + # breaks it, rather than being reflowed with a space in the middle of a + # directory name, which is a path nobody can paste back. + console.print(f" from [cyan]{source}[/cyan]", soft_wrap=True) + console.print(f" to [cyan]{target}[/cyan]", soft_wrap=True) + console.print(f" [dim]{decimal(shutil.disk_usage(anchor).free)} free there[/dim]") + return yes or click.confirm("Proceed?") + + +@contextlib.contextmanager +def _copy_progress(size: int) -> Iterator[Callable[[int], None]]: + """A progress bar counting bytes, yielding the callback that advances it. + + Bytes rather than files: a packed disk-objectstore is a handful of very + large files, so a file count would sit at nothing and then finish. + """ + with Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(), + DownloadColumn(), + TimeElapsedColumn(), + console=console, + ) as progress: + task = progress.add_task("Copying storage", total=size) + yield lambda advance: progress.advance(task, advance) + + +def _refuse_if_sharing(source: object, sandbox: object, profile: object) -> None: + """Stop unless the sandbox's storage is provably the source's alone. + + Called before the copy as well as before registration: a source directory + that contains the sandbox root would otherwise be copied into itself, and + the refusal would arrive after the gigabytes had landed. + """ + from aiida_agents.sandbox.copy import shares_storage + + if shares_storage(source, sandbox): # type: ignore[arg-type] + raise click.ClickException( + "Refusing to build a sandbox that shares storage with " + f"{profile.name!r}. Deleting it would take the real data with it." # type: ignore[attr-defined] + ) + + +def _copy_postgres( + config: dict[str, t.Any], + database: str, + commands: list[ShellCommand], + *, + yes: bool, + role: str, +) -> bool: + """Make the PostgreSQL copy, or print the commands that would make it. + + Returns: + True when the copy is in place and the profile may be registered. + """ + from aiida_agents.sandbox.postgres import ( + PostgresUnavailableError, + copy_database, + create_database, + ) + + # The first command creates the database, which a superuser connection does + # here instead; the second is the one that moves the data, as the profile's + # own user, who owns the database by then. + _, copy = commands + + host = config.get("database_hostname") or "localhost" + console.print( + f"This copies the database [bold]{config.get('database_name')}[/bold]" + ) + console.print( + f" to [cyan]{database}[/cyan] on [cyan]{host}[/cyan]", soft_wrap=True + ) + console.print( + " [dim]Creating a database needs a PostgreSQL superuser, so this may " + "ask for a password.[/dim]" + ) + if not yes and not click.confirm("Proceed?"): + return False + + try: + create_database(config, database) + with console.status("Copying the database ..."): + copy_database(copy, config) + except PostgresUnavailableError as exc: + console.print(f"\n[yellow]![/yellow] Could not do it here: {exc}") + console.print("Run these instead, as a role that may create databases:\n") + for command in commands: + console.print(f"[dim]# {command.explanation}[/dim]") + # `soft_wrap`, or rich reflows the command at the terminal width + # and pasting it runs the tail as a command of its own. This text + # exists to be pasted; a line break in it is a broken command. + console.print(f"[dim]{command.as_shell()}[/dim]\n", soft_wrap=True) + _print_readonly_role(database, role) + return False + + console.print(f"[green]✓[/green] Copied into [cyan]{database}[/cyan]") + _print_readonly_role(database, role) + return True + + +def _print_readonly_role(database: str, role: str) -> None: + """The optional second layer, still printed rather than run. + + Unlike the copy, this one is genuinely optional: the copy is the mechanism, + and a role that cannot write is belt and braces over it. Somebody who wants + it can paste four statements; nobody needs them to get a working sandbox. + """ + from aiida_agents.sandbox.setup import readonly_role_sql + + console.print( + "\n[bold]Optional, as a second layer:[/bold] read the copy through a " + "role that cannot write it.\n" + ) + password = secrets.token_urlsafe(24) + console.print( + f"[dim]{readonly_role_sql(database, role, password)}[/dim]\n", soft_wrap=True + ) + + def _source_profile(name: str | None) -> object: from aiida.manage.configuration import get_config @@ -63,7 +222,8 @@ def sandbox() -> None: show_default=True, help="PostgreSQL role to read the copy through. Optional hardening.", ) -def init(profile: str | None, sandbox_name: str, role: str) -> None: +@click.option("--yes", is_flag=True, help="Do not ask before copying.") +def init(profile: str | None, sandbox_name: str, role: str, yes: bool) -> None: """Copy the profile's storage and register a sandbox profile for it. On SQLite this is done for you: the storage directory is copied and the @@ -77,13 +237,14 @@ def init(profile: str | None, sandbox_name: str, role: str) -> None: from aiida_agents.sandbox.copy import ( SUPPORTED_BACKENDS, - SandboxStorage, + ProfileStorage, copy_sqlite_storage, postgres_copy_commands, + postgres_sandbox_storage, register_profile, sandbox_profile_dictionary, sandbox_storage_root, - shares_storage, + storage_size, ) config = get_config() @@ -98,67 +259,45 @@ def init(profile: str | None, sandbox_name: str, role: str) -> None: ) if sandbox_name in {existing.name for existing in config.profiles}: raise click.ClickException( - f"Profile {sandbox_name!r} already exists. Use `sandbox refresh` to " - "rebuild it, or `sandbox teardown` to remove it first." + f"Profile {sandbox_name!r} already exists. Remove it with `sandbox " + "teardown` first, then run this again to rebuild it." ) root = sandbox_storage_root(sandbox_name) if backend == "core.sqlite_dos": - from pathlib import Path - + storage = Path(storage_config["filepath"]) target = root / "storage" - console.print(f"Copying storage to [cyan]{target}[/cyan] ...") + new_storage = ProfileStorage(backend, {"filepath": str(target)}) + # Before the copy, not after. A source directory that contains the + # sandbox root would be copied into itself, and the check below would + # then refuse to register what had just been written to disk. + _refuse_if_sharing(ProfileStorage(backend, storage_config), new_storage, source) try: - copy_sqlite_storage(Path(storage_config["filepath"]), target) + size = storage_size(storage) + if not _agreed_to_copy(storage, target, size, yes=yes): + return + with _copy_progress(size) as advance: + copy_sqlite_storage(storage, target, progress=advance) except (FileNotFoundError, FileExistsError, OSError) as exc: raise click.ClickException(str(exc)) from exc - new_storage = SandboxStorage(backend, {"filepath": str(target)}) else: sandbox_database = ( f"{storage_config.get('database_name', 'aiida')}_agents_sandbox" ) repository = root / "repository" - console.print( - "[bold]The copy lives in PostgreSQL, so these two run as a role that " - "may create databases:[/bold]\n" + new_storage = postgres_sandbox_storage( + storage_config, database=sandbox_database, repository=repository ) - for explanation, command in postgres_copy_commands( - storage_config, sandbox_database + commands = postgres_copy_commands(storage_config, sandbox_database) + if not _database_exists(new_storage.config) and not _copy_postgres( + storage_config, sandbox_database, commands, yes=yes, role=role ): - console.print(f"[dim]# {explanation}[/dim]") - console.print(f"[dim]{command}[/dim]\n") - console.print( - "[bold]Optional, as a second layer:[/bold] read the copy through a " - "role that cannot write it.\n" - ) - from aiida_agents.sandbox.setup import readonly_role_sql - - password = secrets.token_urlsafe(24) - console.print( - f"[dim]{readonly_role_sql(sandbox_database, role, password)}[/dim]\n" - ) - console.print( - f"Then rerun [cyan]aiida-agents sandbox init --sandbox-name " - f"{sandbox_name}[/cyan] to register the profile." - ) - new_storage = SandboxStorage( - backend, - { - **storage_config, - "database_name": sandbox_database, - "repository_uri": str(repository), - }, - ) - if not _database_exists(new_storage.config): return - # The whole point of the exercise, checked rather than assumed. - if shares_storage(backend, storage_config, new_storage.backend, new_storage.config): - raise click.ClickException( - "Refusing to register a sandbox that shares storage with " - f"{source.name!r}. Deleting it would take the real data with it." # type: ignore[attr-defined] - ) + # The whole point of the exercise, checked rather than assumed. Checked + # again here because the PostgreSQL path builds its storage further down. + _refuse_if_sharing(ProfileStorage(backend, storage_config), new_storage, source) register_profile( config, @@ -209,7 +348,7 @@ def check(profile: str) -> None: """ from aiida.manage.configuration import get_config - from aiida_agents.sandbox.copy import profiles_sharing_storage + from aiida_agents.sandbox.copy import Overlap, profiles_sharing_storage from aiida_agents.sandbox.setup import verify_read_only config = get_config() @@ -220,11 +359,28 @@ def check(profile: str) -> None: failures = profiles_sharing_storage(config, profile) if failures: - console.print( - f"[red]✗[/red] {profile!r} shares storage with " - f"{', '.join(repr(name) for name in failures)}. Deleting it would " - "destroy their data too. This must not be used as a sandbox." - ) + # A proven overlap and a config that could not be read both stop the + # command, and saying which one it was is the difference between "your + # sandbox is pointed at your own storage" and "some profile here does + # not say where its storage is". Only the first is about their data. + proven = any(entry.overlap is Overlap.SHARED for entry in failures) + verdict = "must not be used" if proven else "cannot be cleared for use" + console.print(f"[red]✗[/red] {profile!r} {verdict} as a sandbox:") + for entry in failures: + console.print(f" {entry.describe()}") + # Both remedies when there is one of each: `init` compares the copy + # against its source alone, so a rebuild clears the overlap and leaves + # the unreadable profile failing the next check for the same reason. + if proven: + console.print( + f" Delete it with [cyan]verdi profile delete --keep-data " + f"{profile}[/cyan], then rebuild with [cyan]aiida-agents " + "sandbox init[/cyan]." + ) + if any(entry.overlap is Overlap.UNKNOWN for entry in failures): + console.print( + " What cannot be read is treated as sharing until it can be ruled out." + ) raise SystemExit(1) console.print(f"[green]✓[/green] {profile!r} shares no storage with any profile") @@ -235,24 +391,6 @@ def check(profile: str) -> None: console.print(f"{mark} {result.detail}") -@sandbox.command("refresh") -@click.option( - "--profile", default=DEFAULT_PROFILE, show_default=True, help="Sandbox to rebuild." -) -@click.option( - "--source", default=None, help="Profile to copy from. Default: the default profile." -) -@click.pass_context -def refresh(ctx: click.Context, profile: str, source: str | None) -> None: - """Rebuild the copy from the source profile. - - The copy drifts the moment the user runs anything, which is the accepted - cost of not sharing their storage. This is how it catches up. - """ - ctx.invoke(teardown, profile=profile, yes=True) - ctx.invoke(init, profile=source, sandbox_name=profile, role=DEFAULT_ROLE) - - @sandbox.command("teardown") @click.option( "--profile", default=DEFAULT_PROFILE, show_default=True, help="Sandbox to remove." @@ -268,8 +406,11 @@ def teardown(profile: str, yes: bool) -> None: from aiida.manage.configuration import get_config from aiida_agents.sandbox.copy import ( + postgres_drop_command, + profile_storage, profiles_sharing_storage, sandbox_storage_root, + storage_size, ) config = get_config() @@ -279,15 +420,41 @@ def teardown(profile: str, yes: bool) -> None: sharing = profiles_sharing_storage(config, profile) if sharing: + reasons = "\n".join(f" {entry.describe()}" for entry in sharing) raise click.ClickException( - f"{profile!r} shares storage with {', '.join(sharing)}. Refusing to " - "delete its data. Remove the profile by hand if you are certain." + f"Refusing to delete {profile!r} and its storage:\n{reasons}\n" + "Remove the profile by hand if you are certain." ) - if not yes and not click.confirm(f"Delete profile {profile!r} and its copy?"): + # Named rather than described. "Its copy" reads as a copy of the profile, + # which is the one thing this does not delete: the profile entry is the + # cheap half, and the storage it points at is what takes the disk with it. + storage = profile_storage(config.get_profile(profile)) + root = sandbox_storage_root(profile) + console.print(f"This removes the profile [cyan]{profile}[/cyan], and") + if root.is_dir(): + console.print( + f" [bold]{decimal(storage_size(root))}[/bold] of copied storage at " + f"[cyan]{root}[/cyan]", + soft_wrap=True, + ) + else: + console.print(f" nothing on disk: [cyan]{root}[/cyan] is not there") + console.print( + " [dim]Anything the sandbox holds and the source does not goes with it.[/dim]" + ) + if not yes and not click.confirm("Delete?"): return config.delete_profile(profile, delete_storage=False) config.store() # type: ignore[no-untyped-call] - shutil.rmtree(sandbox_storage_root(profile), ignore_errors=True) - console.print(f"[green]✓[/green] Removed {profile!r} and its copied storage") + shutil.rmtree(root, ignore_errors=True) + console.print(f"[green]✓[/green] Removed {profile!r} and the storage it copied") + + # The PostgreSQL copy lives in the server, which nothing here reaches. + # `init` printed the commands that made it, so this prints the one that + # undoes them rather than leaving a database nobody is accounting for. + if storage.backend == "core.psql_dos": + drop = postgres_drop_command(storage.config) + console.print(f"\n[dim]# {drop.explanation}[/dim]") + console.print(f"[dim]{drop.as_shell()}[/dim]", soft_wrap=True) diff --git a/src/aiida_agents/sandbox/copy.py b/src/aiida_agents/sandbox/copy.py index f1c685f..5f60397 100644 --- a/src/aiida_agents/sandbox/copy.py +++ b/src/aiida_agents/sandbox/copy.py @@ -36,36 +36,107 @@ from __future__ import annotations +import enum +import shlex import shutil import typing as t import uuid +from collections.abc import Callable from dataclasses import dataclass +from functools import partial from pathlib import Path +from urllib.parse import urlparse +from urllib.request import url2pathname + +from rich.filesize import decimal +from typing_extensions import assert_never __all__ = [ - "SandboxStorage", + "FILEPATH_BACKENDS", + "SUPPORTED_BACKENDS", + "DatabaseLocation", + "Location", + "Overlap", + "PathLocation", + "ProfileStorage", + "SharingProfile", + "ShellCommand", + "StorageConfig", + "ConfigLike", + "ProfileLike", + "profile_storage", "profiles_sharing_storage", "register_profile", "copy_sqlite_storage", "postgres_copy_commands", + "postgres_drop_command", + "postgres_sandbox_storage", "sandbox_profile_dictionary", "sandbox_storage_root", "shares_storage", "storage_locations", + "storage_overlap", + "storage_size", ] -#: Backends this module knows how to copy. Anything else is refused by name -#: rather than attempted, because a copy that half worked would be worse than -#: no sandbox: it would look like containment. -SUPPORTED_BACKENDS = frozenset({"core.psql_dos", "core.sqlite_dos"}) +# Backends this module knows how to copy. Anything else is refused by name +# rather than attempted, because a copy that half worked would be worse than +# no sandbox: it would look like containment. +SUPPORTED_BACKENDS: t.Final = frozenset({"core.psql_dos", "core.sqlite_dos"}) + +# Backends whose whole storage is the one path under `filepath`, and so can be +# compared without loading them. Deliberately wider than SUPPORTED_BACKENDS: an +# imported archive cannot be copied into a sandbox, but it can be told apart +# from one, and until it could, every config holding an archive read as sharing +# storage with the sandbox. +FILEPATH_BACKENDS: t.Final = frozenset({"core.sqlite_dos", "core.sqlite_zip"}) + +# Read and written a chunk at a time so the copy can be reported as it goes. +# Large enough that the reporting costs nothing against gigabytes of packs. +_COPY_CHUNK_BYTES: t.Final = 4 * 1024 * 1024 + +#: A profile's ``storage.config`` as it comes out of ``config.json``: whatever +#: the backend put there, which is why every read of it is defensive. +StorageConfig: t.TypeAlias = dict[str, t.Any] + + +class ProfileLike(t.Protocol): + """The three things this module reads off an AiiDA ``Profile``. + + A protocol rather than the class itself, because importing ``aiida`` at + module scope is what the lazy imports below exist to avoid, and because a + test that hands this a stub then has the stub checked rather than trusted. + """ + + @property + def name(self) -> str: ... + + @property + def storage_backend(self) -> str: ... + + @property + def storage_config(self) -> StorageConfig | None: ... + + +class ConfigLike(t.Protocol): + """The one thing this module reads off an AiiDA ``Config``.""" + + @property + def profiles(self) -> t.Sequence[ProfileLike]: ... @dataclass(frozen=True) -class SandboxStorage: - """Where a sandbox's copied storage lives.""" +class ProfileStorage: + """A profile's storage: the backend that reads it, and where it keeps data. + + The pair every comparison here works on, rather than the two of them passed + side by side. ``shares_storage(a_backend, b_config, ...)`` is a mistake the + type system cannot see, and it is the one comparison standing between a + ``--delete-data`` and somebody's work. + """ backend: str - config: dict[str, t.Any] + config: StorageConfig def sandbox_storage_root(sandbox_name: str) -> Path: @@ -80,43 +151,197 @@ def sandbox_storage_root(sandbox_name: str) -> Path: return Path(AiiDAConfigDir.get()) / "agents-sandbox" / sandbox_name -def storage_locations(backend: str, config: dict[str, t.Any]) -> frozenset[str]: - """Everything ``verdi profile delete --delete-data`` would destroy. +class Overlap(enum.Enum): + """How one profile's storage compares with another's. - Canonical strings rather than paths so that two profiles can be compared - without either being loadable. An empty set means "I could not tell", which - :func:`shares_storage` treats as sharing. + ``UNKNOWN`` is not a shade of ``SEPARATE``. It is acted on exactly like + ``SHARED``, and the two are kept apart only so the user can be told which + one they have: "these are the same directory" and "there is not enough in + these configurations to tell" call for different next steps. + + The second is the rarer of the two and not hypothetical: a storage backend + from a plugin, or a profile whose configuration was written by something + other than ``verdi`` and does not say where its storage is. aiida-core + guards the same ground, raising on a ``repository_uri`` that is missing, + relative, or not a ``file://`` URL. + """ + + SEPARATE = "separate" + SHARED = "shared" + UNKNOWN = "unknown" + + +@dataclass(frozen=True, kw_only=True) +class SharingProfile: + """A profile that could not be proved separate from the one being checked. + + ``overlap`` cannot be ``SEPARATE``: a profile proved separate is not one of + these at all. Keyword-only because ``name`` and ``backend`` are both + strings, and swapping them would build a sentence that reads fine and names + the wrong thing. """ - if backend == "core.sqlite_dos": - filepath = config.get("filepath") + + name: str + backend: str + overlap: t.Literal[Overlap.SHARED, Overlap.UNKNOWN] + + def describe(self) -> str: + """The reason, as a clause that follows the checked profile's name. + + Here rather than at the call sites because ``sandbox check``, + ``sandbox teardown`` and ``doctor`` all report it, and three wordings + of one finding would drift. + """ + if self.overlap is Overlap.SHARED: + return ( + f"shares storage with {self.name!r}, so deleting either " + "destroys the other's data" + ) + if self.overlap is Overlap.UNKNOWN: + # Not "its backend cannot be read": the unreadable half is as often + # an incomplete configuration, and it can be either profile's. + return ( + f"cannot be told apart from {self.name!r} (storage backend " + f"{self.backend!r}): the two configurations do not say enough " + "to prove they are separate" + ) + assert_never(self.overlap) # pragma: no cover + + +def _storage_path(value: object) -> Path | None: + """A configuration value read as a place on disk, or ``None`` if it is not one. + + The one place a configuration value becomes a comparable path. These come + from ``config.json``, where a path is a string; a missing key or anything + else is not one. Neither is a relative path: :meth:`Path.resolve` would + answer it against the working directory, so the same two profiles would + compare differently depending on where the command was run. + """ + if not isinstance(value, str): + return None + # Postgres writes its repository as a `file://` URL, the SQLite backends a + # plain path. Both name a place on disk, and the URL is decoded the way + # `aiida.storage.psql_dos.backend.get_filepath_container` decodes it --- + # `Path.as_uri` percent-encodes, so stripping the scheme by hand leaves + # `/home/u/My%20Drive/...`, which never matches the directory it names. + text = url2pathname(urlparse(value).path) if value.startswith("file://") else value + path = Path(text) + return path.resolve() if path.is_absolute() else None + + +def _database_location(config: StorageConfig) -> DatabaseLocation | None: + """The database a Postgres profile uses, or ``None`` if it cannot be read. + + The same rule as the paths, for the same reason: ``config.json`` can hold + anything, and a value of the wrong type has to fail closed rather than + raise out of the comparison. These end up in a set, so an unhashable + hostname took ``frozenset`` down with a ``TypeError``. + """ + name = config.get("database_name") + if not isinstance(name, str) or not name: + return None + host = config.get("database_hostname") or "localhost" + port = config.get("database_port") or 5432 + if not isinstance(host, str) or not isinstance(port, int): + return None + return DatabaseLocation(host=host, port=port, name=name) + + +@dataclass(frozen=True) +class PathLocation: + """Storage kept somewhere on disk: a directory, or a single archive file.""" + + path: Path + + def overlaps(self, other: Location) -> bool: + """Whether destroying this place would destroy ``other`` as well. + + Containment counts, not only equality. ``sandbox teardown`` removes its + root with :func:`shutil.rmtree`, so a profile whose storage sits inside + another's goes with it, and comparing the two paths for equality would + call them separate right up until the moment one deleted the other. + """ + if not isinstance(other, PathLocation): + return False return ( - frozenset({f"dir:{Path(filepath).resolve()}"}) if filepath else frozenset() + self.path == other.path + or self.path.is_relative_to(other.path) + or other.path.is_relative_to(self.path) ) - if backend == "core.psql_dos": - locations = set() - database = config.get("database_name") - if database: - host = config.get("database_hostname") or "localhost" - port = config.get("database_port") or 5432 - locations.add(f"pg://{host}:{port}/{database}") - repository = config.get("repository_uri") - if repository: - locations.add( - f"dir:{Path(str(repository).removeprefix('file://')).resolve()}" - ) + +@dataclass(frozen=True) +class DatabaseLocation: + """Storage kept in a database on a server, named the way a client reaches it.""" + + host: str + port: int + name: str + + def overlaps(self, other: Location) -> bool: + """Whether this is the same database. Two of them nest in no sense.""" + return self == other + + +#: Somewhere a profile keeps data. A Postgres profile has two, one of each. +Location: t.TypeAlias = PathLocation | DatabaseLocation + + +def storage_locations(storage: ProfileStorage) -> frozenset[Location]: + """Everywhere ``verdi profile delete --delete-data`` would destroy. + + Read from the configuration so that two profiles can be compared without + either being loadable. An empty set means "I could not tell", which + :func:`shares_storage` treats as sharing. + """ + # core.sqlite_dos keeps its database and container in one directory; + # core.sqlite_zip is a single archive file, read-only in every respect but + # this one --- `SqliteZipBackend.delete` unlinks it, so `--delete-data` + # destroys it just the same. One kind of location for both, because two + # profiles naming one path share it whatever each of them calls it. + if storage.backend in FILEPATH_BACKENDS: + path = _storage_path(storage.config.get("filepath")) + return frozenset({PathLocation(path)}) if path is not None else frozenset() + + if storage.backend == "core.psql_dos": + locations: set[Location] = set() + database = _database_location(storage.config) + if database is not None: + locations.add(database) + repository = _storage_path(storage.config.get("repository_uri")) + if repository is not None: + locations.add(PathLocation(repository)) # A Postgres profile with neither is not something we can reason about. return frozenset(locations) if len(locations) == 2 else frozenset() return frozenset() -def shares_storage( - backend_a: str, - config_a: dict[str, t.Any], - backend_b: str, - config_b: dict[str, t.Any], -) -> bool: +def profile_storage(profile: ProfileLike) -> ProfileStorage: + """The storage of an AiiDA profile, as this module compares it. + + The one place AiiDA's profile object is read, so a missing + ``storage_config`` becomes an empty one here rather than at four call sites. + """ + return ProfileStorage(profile.storage_backend, profile.storage_config or {}) + + +def storage_overlap(a: ProfileStorage, b: ProfileStorage) -> Overlap: + """How two profiles' storage compares, including *why* when it is not separate. + + The comparison :func:`shares_storage` answers yes-or-no. Callers that + report the finding to a user need the distinction :class:`Overlap` draws, + and taking it from here keeps one implementation of the rule. + """ + locations_a = storage_locations(a) + locations_b = storage_locations(b) + if not locations_a or not locations_b: + return Overlap.UNKNOWN + shared = any(one.overlaps(other) for one in locations_a for other in locations_b) + return Overlap.SHARED if shared else Overlap.SEPARATE + + +def shares_storage(a: ProfileStorage, b: ProfileStorage) -> bool: """Whether deleting one profile's data would destroy the other's. **Fails closed.** A backend this module does not understand, or a config @@ -128,14 +353,10 @@ def shares_storage( Sharing *either* the database or the repository counts. A sandbox with its own database but the real repository still loses the user their files. """ - locations_a = storage_locations(backend_a, config_a) - locations_b = storage_locations(backend_b, config_b) - if not locations_a or not locations_b: - return True - return bool(locations_a & locations_b) + return storage_overlap(a, b) is not Overlap.SEPARATE -def profiles_sharing_storage(config: t.Any, name: str) -> list[str]: +def profiles_sharing_storage(config: ConfigLike, name: str) -> list[SharingProfile]: """Every other profile whose data would go with ``name``'s. The one implementation of this question. ``sandbox check`` asks it before @@ -145,24 +366,27 @@ def profiles_sharing_storage(config: t.Any, name: str) -> list[str]: somebody deletes. Returns: - Profile names, empty when the sandbox is genuinely self-contained. + One entry per profile that could not be proved separate, carrying which + of the two reasons it was. Empty when the sandbox is genuinely + self-contained. """ profiles = {profile.name: profile for profile in config.profiles} - target = profiles[name] - return [ - other.name - for other in config.profiles - if other.name != name - and shares_storage( - target.storage_backend, - target.storage_config or {}, - other.storage_backend, - other.storage_config or {}, - ) - ] + target = profile_storage(profiles[name]) + found = [] + for other in config.profiles: + if other.name == name: + continue + overlap = storage_overlap(target, profile_storage(other)) + if overlap is not Overlap.SEPARATE: + found.append( + SharingProfile( + name=other.name, backend=other.storage_backend, overlap=overlap + ) + ) + return found -def register_profile(config: t.Any, name: str, dictionary: dict[str, t.Any]) -> None: +def register_profile(config: t.Any, name: str, dictionary: StorageConfig) -> None: """Add a profile to the AiiDA configuration and persist it. Here rather than inline at the call sites so AiiDA's untyped configuration @@ -174,7 +398,41 @@ def register_profile(config: t.Any, name: str, dictionary: dict[str, t.Any]) -> config.store() -def copy_sqlite_storage(source: Path, target: Path) -> None: +def storage_size(path: Path) -> int: + """Total bytes of the files under ``path``. + + What a copy of it will write, and so what the disk must have room for and + what a progress bar counts up to. + + Raises: + FileNotFoundError: If ``path`` is not a directory. + """ + if not path.is_dir(): + msg = f"storage directory {path} does not exist" + raise FileNotFoundError(msg) + return sum(item.stat().st_size for item in path.rglob("*") if item.is_file()) + + +def _copy_file( + source: str, target: str, *, progress: Callable[[int], None] | None +) -> None: + """``shutil.copy2`` in chunks, reporting each one as it lands. + + ``copytree`` calls its copy function once per file, and a packed + disk-objectstore is two files of several gigabytes, so counting files would + leave the bar at nothing until it is nearly finished. + """ + with open(source, "rb") as reader, open(target, "wb") as writer: + while chunk := reader.read(_COPY_CHUNK_BYTES): + writer.write(chunk) + if progress is not None: + progress(len(chunk)) + shutil.copystat(source, target) + + +def copy_sqlite_storage( + source: Path, target: Path, *, progress: Callable[[int], None] | None = None +) -> None: """Copy a ``core.sqlite_dos`` storage directory wholesale. The whole directory, not just ``database.sqlite``: the disk-objectstore @@ -185,20 +443,29 @@ def copy_sqlite_storage(source: Path, target: Path) -> None: Args: source: The source profile's ``filepath``. target: Where the copy goes. Must not already exist. + progress: Called with the number of bytes as each chunk lands, for a + caller showing how far along the copy is. Raises: FileNotFoundError: If ``source`` is not a directory. FileExistsError: If ``target`` exists --- refreshing is an explicit teardown, never an overwrite of something already in use. + OSError: If the filesystem has less room than the source occupies. """ - if not source.is_dir(): - msg = f"storage directory {source} does not exist" - raise FileNotFoundError(msg) + required = storage_size(source) if target.exists(): msg = f"{target} already exists; tear the sandbox down before rebuilding it" raise FileExistsError(msg) target.parent.mkdir(parents=True, exist_ok=True) + free = shutil.disk_usage(target.parent).free + if free < required: + msg = ( + f"copying {source} needs {decimal(required)} and {target.parent} has " + f"{decimal(free)} free" + ) + raise OSError(msg) + # Everything, verbatim. The first version skipped disk-objectstore's # `sandbox/` scratch directory on the grounds that a copy which never writes # has no use for in-flight writes. That was true and it did not matter: @@ -206,18 +473,50 @@ def copy_sqlite_storage(source: Path, target: Path) -> None: # loaded as an uninitialised container and every query raised # `UnreachableStorage`. Copy the layout as it is and let the backend decide # what it needs. - shutil.copytree(source, target) + try: + shutil.copytree( + source, target, copy_function=partial(_copy_file, progress=progress) + ) + except OSError: + # A half-written copy is worse than none: it answers queries wrongly, + # and the next attempt refuses it as something already in use. + shutil.rmtree(target, ignore_errors=True) + raise -def postgres_copy_commands( - config: dict[str, t.Any], sandbox_database: str -) -> list[tuple[str, str]]: - """The shell commands that copy a Postgres database, with explanations. +@dataclass(frozen=True) +class ShellCommand: + """A command as the pieces to run it, and as the line to paste it. - Printed rather than run for the same reason the role SQL is: creating a - database needs a privilege the profile's own user usually does not have, and - asking for a superuser connection in order to configure a safety feature is - a worse bargain than showing somebody two commands they can read. + Both from one place. ``sandbox init`` runs these where it can and prints + them where it cannot, and the two must not drift: what it runs has to be + what it would otherwise have told you to run. + """ + + explanation: str + stages: tuple[tuple[str, ...], ...] + + def as_shell(self) -> str: + """The command as a line a shell would accept, quoted by :mod:`shlex`. + + Quoted properly rather than by wrapping names in double quotes: a + database called ``my "db"`` is legal in Postgres and would otherwise + produce a line that means something else. + """ + return " | ".join(shlex.join(stage) for stage in self.stages) + + +def _connection_arguments(config: StorageConfig) -> tuple[str, ...]: + host = config.get("database_hostname") or "localhost" + port = config.get("database_port") or 5432 + user = config.get("database_username") or "" + return ("--host", str(host), "--port", str(port), "--username", str(user)) + + +def postgres_copy_commands( + config: StorageConfig, sandbox_database: str +) -> list[ShellCommand]: + """The commands that copy a Postgres database, with explanations. ``pg_dump | psql`` rather than ``CREATE DATABASE ... TEMPLATE ...``. The template form is much faster, and fails outright while any other session is @@ -225,29 +524,70 @@ def postgres_copy_commands( shell`` open, is most of the time. Returns: - ``(explanation, command)`` pairs, in the order they must be run. + The commands in the order they must be run. """ - host = config.get("database_hostname") or "localhost" - port = config.get("database_port") or 5432 - user = config.get("database_username") or "" - source = config.get("database_name") or "" - connection = f"--host {host} --port {port} --username {user}" + connection = _connection_arguments(config) + source = str(config.get("database_name") or "") return [ - ( + ShellCommand( "Create the database the copy will live in", - f'createdb {connection} "{sandbox_database}"', + (("createdb", *connection, sandbox_database),), ), - ( + ShellCommand( "Copy the data across (works while the source profile is in use)", - f'pg_dump {connection} --no-owner --no-privileges "{source}" ' - f'| psql {connection} --quiet "{sandbox_database}"', + ( + ("pg_dump", *connection, "--no-owner", "--no-privileges", source), + ("psql", *connection, "--quiet", sandbox_database), + ), ), ] +def postgres_drop_command(config: StorageConfig) -> ShellCommand: + """The command that removes a Postgres copy, with its explanation. + + ``teardown`` deletes the profile and the repository directory, and cannot + reach into the server for the database. Printed for the same reason + :func:`postgres_copy_commands` prints its two: dropping a database needs a + privilege the profile's own user usually does not have. + + Returns: + The command that removes the copied database. + """ + database = str(config.get("database_name") or "") + + return ShellCommand( + "The copied database is still in the server. To remove it as well:", + (("dropdb", *_connection_arguments(config), database),), + ) + + +def postgres_sandbox_storage( + source_config: StorageConfig, *, database: str, repository: Path +) -> ProfileStorage: + """The sandbox's own Postgres storage, cloned from the source profile's. + + Server, port and credentials come across because the copy lives in the same + server; only the database and the repository are its own. + + The repository is written as a ``file://`` URL because that is what + aiida-core requires of ``repository_uri``: given a plain path, + ``get_filepath_container`` raises ``ConfigurationError`` and the registered + profile cannot be opened at all. + """ + return ProfileStorage( + "core.psql_dos", + { + **source_config, + "database_name": database, + "repository_uri": repository.as_uri(), + }, + ) + + def sandbox_profile_dictionary( - source: dict[str, t.Any], storage: SandboxStorage + source: StorageConfig, storage: ProfileStorage ) -> dict[str, t.Any]: """A profile configuration for the sandbox, cloned from the source profile's. diff --git a/src/aiida_agents/sandbox/postgres.py b/src/aiida_agents/sandbox/postgres.py new file mode 100644 index 0000000..079022a --- /dev/null +++ b/src/aiida_agents/sandbox/postgres.py @@ -0,0 +1,140 @@ +"""Making the PostgreSQL copy, rather than telling somebody how to make it. + +``init`` used to print a ``createdb``, a ``pg_dump | psql`` pipeline and five +``GRANT`` statements, and expect a researcher to run them. They do not: they +are here to answer questions about their data, they did not ask for a safety +feature, and most of them have never written SQL. A setup step nobody performs +is a feature nobody has. + +Two things stand between us and doing it for them, and only one is real. + +**Creating the database needs a privilege the profile's user does not have.** +AiiDA creates its database users with ``CREATE USER ... WITH PASSWORD``, no +``CREATEDB`` (see ``Postgres.create_dbuser``), so connecting with the +credentials already in the profile and issuing ``CREATE DATABASE`` fails on +exactly the step that matters. This is why SQLAlchemy alone is not enough, +although it is a perfectly good way to run the statement once something has +handed you a connection that may. + +What can hand it over is :class:`aiida.manage.external.postgres.Postgres`, +which ``verdi presto`` already uses: it tries psycopg as the current user and +falls back to ``sudo su postgres``, so anyone who has a PostgreSQL profile has +been through it once already. No new dependency, and no new trust: it is the +same route that created the database being copied. + +**Copying the data needs ``pg_dump``.** SQLAlchemy cannot express it; a +row-by-row copy through the ORM is what ``verdi archive`` does, slowly. The +alternative in SQL, ``CREATE DATABASE ... TEMPLATE``, refuses to run while any +other session is connected to the source, which with a daemon running is most +of the time. + +So: a superuser connection creates the empty database, and ``pg_dump | psql`` +fills it as the profile's own user, who owns it by then. Where no superuser +connection can be found (a managed server, a remote host, no ``sudo``), the +commands are printed as before. That path is the fallback now rather than the +only road. +""" + +from __future__ import annotations + +import logging +import os +import subprocess + +from aiida_agents.sandbox.copy import ShellCommand, StorageConfig + +logger = logging.getLogger(__name__) + +__all__ = ["PostgresUnavailableError", "copy_database", "create_database"] + + +class PostgresUnavailableError(RuntimeError): + """No privileged connection, or a command that failed. + + Carries what to tell the user, because the caller's fallback is to print + the commands and this is the sentence that explains why. + """ + + +def create_database(config: StorageConfig, database: str) -> None: + """Create the sandbox's database, owned by the profile's own user. + + Raises: + PostgresUnavailableError: If no connection with the privilege to create + a database could be found, or the creation failed. + """ + from aiida.manage.external.postgres import Postgres + + owner = str(config.get("database_username") or "") + dbinfo = { + "host": config.get("database_hostname") or "localhost", + "port": config.get("database_port") or 5432, + } + + try: + # aiida's Postgres is untyped, and this is the one place we cross + # into it, which is why the ignores sit together rather than spread. + postgres = Postgres( # type: ignore[no-untyped-call] + interactive=True, quiet=False, dbinfo=dbinfo + ) + postgres.determine_setup() + except Exception as exc: + msg = f"no PostgreSQL connection that may create a database: {exc}" + raise PostgresUnavailableError(msg) from exc + + if not postgres.is_connected: + msg = "no PostgreSQL connection that may create a database" + raise PostgresUnavailableError(msg) + + try: + if postgres.db_exists(database): # type: ignore[no-untyped-call] + logger.info("database %r already exists; reusing it", database) + return + postgres.create_db(owner, database) # type: ignore[no-untyped-call] + except Exception as exc: + msg = f"could not create the database {database!r}: {exc}" + raise PostgresUnavailableError(msg) from exc + + +def copy_database(command: ShellCommand, config: StorageConfig) -> None: + """Run a copy command, with the profile's own password supplied. + + The password comes from the profile rather than a prompt, because it is + already in ``config.json`` and prompting for something the caller can read + teaches people to type their database password at anything that asks. + + Raises: + PostgresUnavailableError: If a stage failed, carrying its ``stderr``. + """ + environment = dict(os.environ) + if password := config.get("database_password"): + environment["PGPASSWORD"] = str(password) + + processes: list[subprocess.Popen[bytes]] = [] + try: + for stage in command.stages: + previous = processes[-1].stdout if processes else None + processes.append( + subprocess.Popen( # noqa: S603 + stage, + stdin=previous, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=environment, + ) + ) + if previous is not None: + # The upstream process holds the read end too; closing it here + # is what lets it see EOF if the downstream one exits early. + previous.close() + except (OSError, subprocess.SubprocessError) as exc: + msg = f"could not run {command.stages[0][0]!r}: {exc}" + raise PostgresUnavailableError(msg) from exc + + for process in processes: + _, stderr = process.communicate() + if process.returncode: + name = process.args[0] if isinstance(process.args, (list, tuple)) else "" + detail = stderr.decode(errors="replace").strip().splitlines() + msg = f"{name} failed: {detail[-1] if detail else 'no output'}" + raise PostgresUnavailableError(msg) diff --git a/src/aiida_agents/sandbox/runner.py b/src/aiida_agents/sandbox/runner.py index bf5088c..3607642 100644 --- a/src/aiida_agents/sandbox/runner.py +++ b/src/aiida_agents/sandbox/runner.py @@ -8,31 +8,29 @@ limits, its own session, and no inherited credentials. An infinite loop is a timeout rather than a hung CLI, and a crash takes nothing with it. 3. That subprocess loads a **profile the caller nominates**, which is meant to - be one whose database role cannot write. The containment is Postgres - refusing the write, not us noticing it. + be a disposable copy of the user's storage (see + :mod:`aiida_agents.sandbox.copy`). The containment is that a write which + gets past layer 1 lands somewhere nobody reads, not that it is refused. Layer 3 is the one this module cannot enforce: it loads whichever profile it is -given, and whether that profile can write is a fact about the database. +given, and whether that profile is a copy is a fact about the configuration. :func:`run_in_sandbox` therefore never guesses a profile. A caller that passes -the user's own writable profile gets exactly what it asked for, which is why -the tool layer above must not do that. - -What this does **not** yet contain, and the reason codegen should not be on by -default (`#73 `_): - -*The profile shares its storage with the user's real one.* That was a -deliberate choice --- an empty database cannot answer "which structures did I -relax last month" --- but it weighed a shared database against an *empty* one -and missed the option that a **copy** is neither. Sharing storage cost a -maintainer his real database during dogfooding, because deleting the sandbox -profile and agreeing to delete its data deletes the data underneath both. A -read-only role does not save you from that: the destructive command is run by -the user, as themselves, against a profile they were told was disposable. - -*The database role only stops database writes.* Generated code that gets past -layer 1 can still read the filesystem and reach the network. The environment -scrub and the rlimits here narrow that, but narrowing is not closing: OS-level -isolation (bwrap, nsjail, a container) is the fix, and it is not here yet. +the user's own profile gets exactly what it asked for, which is why the tool +layer above must not do that. + +What this does **not** contain: + +*Writes that layer 1 misses.* On PostgreSQL a read-only role can be put under +the copy as well, and ``sandbox check`` verifies it where it is. On SQLite +there is no such role to have, so a write the guard does not recognise +succeeds against the copy. It costs the user nothing, and costs the sandbox +its accuracy until it is rebuilt, but it means layer 1's list of forbidden +names is doing more work than a pre-check should. + +*Anything outside the database.* Generated code that gets past layer 1 can +still read the filesystem and reach the network. The environment scrub and the +rlimits here narrow that, but narrowing is not closing: OS-level isolation +(bwrap, nsjail, a container) is the fix, and it is not here yet. """ from __future__ import annotations diff --git a/src/aiida_agents/sandbox/setup.py b/src/aiida_agents/sandbox/setup.py index ef4fbf8..ebefbe7 100644 --- a/src/aiida_agents/sandbox/setup.py +++ b/src/aiida_agents/sandbox/setup.py @@ -1,9 +1,11 @@ -"""Setting up the read-only profile the sandbox runs against, and proving it is one. - -The containment in :mod:`aiida_agents.sandbox.runner` is a Postgres role that -cannot write. This module helps create that role and --- more importantly --- -lets a caller *check* it, rather than trusting that whoever set it up got the -grants right. +"""The optional read-only role over the sandbox copy, and proving it is one. + +The containment in :mod:`aiida_agents.sandbox.runner` is that the sandbox is a +copy (:mod:`aiida_agents.sandbox.copy`). On PostgreSQL a role that cannot write +can be put under that copy as a second layer, and this module helps create it +and --- more importantly --- lets a caller *check* it, rather than trusting +that whoever set it up got the grants right. SQLite has no roles, so there the +copy is the whole of it. Two deliberate choices about how far this goes. diff --git a/src/aiida_agents/tools/codegen/__init__.py b/src/aiida_agents/tools/codegen/__init__.py index d76914e..7f84cc4 100644 --- a/src/aiida_agents/tools/codegen/__init__.py +++ b/src/aiida_agents/tools/codegen/__init__.py @@ -1,9 +1,11 @@ """Tools owned by the Codegen agent: running the Python it just wrote. One tool, and the interesting thing about it is what it is *not*. It is not -approval-gated, because it cannot write: it runs against a profile whose -database role has no INSERT privilege, in a subprocess, behind a static guard -(see :mod:`aiida_agents.sandbox`). +approval-gated, because nothing it does to the profile reaches the user's own: +it runs against a disposable copy of their storage, in a subprocess, behind a +static guard (see :mod:`aiida_agents.sandbox`). That containment covers the +AiiDA storage and stops there --- the filesystem and the network the subprocess +sees are the real ones. That is the point of the whole arrangement. An approval prompt showing twenty lines of unexecuted Python asks the user to review something they cannot diff --git a/src/aiida_agents/tools/codegen/execution.py b/src/aiida_agents/tools/codegen/execution.py index 65233ec..9459dd0 100644 --- a/src/aiida_agents/tools/codegen/execution.py +++ b/src/aiida_agents/tools/codegen/execution.py @@ -12,30 +12,44 @@ _NOT_CONFIGURED = ( "No sandbox profile is configured, so this code cannot be run. Tell the " - "user to run `aiida-agents sandbox init` once to create a read-only " - "profile, then `aiida-agents sandbox check` to confirm it. Do NOT run the " + "user to run `aiida-agents sandbox init` once to copy their storage into " + "one, then `aiida-agents sandbox check` to confirm it. Do NOT run the " "code any other way, and do NOT claim to have run it: show them the " "snippet and say it is unverified." ) +_NOT_SEPARATE = ( + "The configured sandbox profile shares storage with another profile, so " + "this code was not run: anything it did would land in data that is not the " + "sandbox's. Tell the user to run `aiida-agents sandbox check`, which names " + "the profile it overlaps with. Do NOT run the code any other way, and do " + "NOT claim to have run it: show them the snippet and say it is unverified." +) + def run_aiida_code(code: str) -> str: - """Run Python against the user's AiiDA data and return what it printed. + """Run Python against a copy of the user's AiiDA data and return what it printed. Use this to **check the code you just wrote before showing it to the - user**. It runs against their real database, so the output is a real - answer --- and if the snippet is wrong you get the traceback instead of - them getting a broken snippet. + user**. It runs against a copy of their database, made when they set the + sandbox up, so the output is a real answer over their real provenance --- + and if the snippet is wrong you get the traceback instead of them getting a + broken snippet. Always ``print()`` what you want to see; a bare expression on the last line produces nothing, exactly as in a script. - The profile this runs against **cannot write**. Its database role has no - INSERT privilege, so anything that stores, deletes or submits will be - refused --- by Postgres, not by politeness. Do not attempt writes here: to - submit a workflow or import a structure, say so and let the Execution agent + **Nothing you do to this profile reaches the user's own.** Writes are + refused before the code runs, and one that slips past that check lands in + the copy, where nobody but you will ever see it. So never report having + stored, submitted, deleted or changed anything, however plainly the output + says you did: to change the user's data, say so and let the Execution agent do it behind its approval prompt. + That covers the AiiDA profile and nothing else. The filesystem and the + network are the user's real ones, so code that reads their files or reaches + out over the network is not made harmless by the copy. + If the code is refused or raises, read the reason, fix the snippet and try again. Show the user code that has run, not code you hope works. @@ -47,7 +61,10 @@ def run_aiida_code(code: str) -> str: What the code printed, or the reason it was refused, timed out or raised. """ + from aiida.manage.configuration import get_config + from aiida_agents.sandbox import run_in_sandbox + from aiida_agents.sandbox.copy import profiles_sharing_storage from aiida_agents.sandbox.setup import sandbox_profile_exists settings = SandboxSettings() @@ -60,6 +77,25 @@ def run_aiida_code(code: str) -> str: logger.warning("sandbox profile %r is not configured", profile) return _NOT_CONFIGURED + # `sandbox check` proves separation; this proves it again at the moment it + # matters. The setting is a profile name, and nothing stops it naming the + # user's own profile --- which is issue #73 with an extra step. + try: + sharing = profiles_sharing_storage(get_config(), profile) + except Exception: + # Deliberately broad, and deliberately fails closed: whatever went + # wrong, separation was not proved, and the answer to "I could not + # tell" here has to be the same as the answer to "they share". + logger.warning("could not verify sandbox profile %r", profile, exc_info=True) + return _NOT_SEPARATE + if sharing: + logger.warning( + "sandbox profile %r is not separate: %s", + profile, + ", ".join(entry.name for entry in sharing), + ) + return _NOT_SEPARATE + result = run_in_sandbox(code, profile=profile, timeout=settings.sandbox_timeout) logger.info( "sandbox run: ok=%s refused=%s timed_out=%s in %.1fs", diff --git a/tests/agents/codegen/test_codegen.py b/tests/agents/codegen/test_codegen.py index de78aa5..442738e 100644 --- a/tests/agents/codegen/test_codegen.py +++ b/tests/agents/codegen/test_codegen.py @@ -130,6 +130,65 @@ def test_the_refusal_forbids_claiming_it_ran( assert "do NOT claim to have run it" in execution.run_aiida_code("print(1)") + def test_it_refuses_a_profile_that_is_not_separate( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """`sandbox check` proves separation; this proves it again when it counts. + + The setting is a profile name and nothing stops it naming the user's + own profile, which is issue #73 with one extra step. + """ + from aiida_agents.sandbox.copy import Overlap, SharingProfile + from aiida_agents.tools.codegen import execution + + monkeypatch.setattr( + "aiida_agents.sandbox.setup.sandbox_profile_exists", + lambda profile: True, + ) + monkeypatch.setattr( + "aiida_agents.sandbox.copy.profiles_sharing_storage", + lambda config, name: [ + SharingProfile( + name="real", backend="core.sqlite_dos", overlap=Overlap.SHARED + ) + ], + ) + ran: list[str] = [] + monkeypatch.setattr( + "aiida_agents.sandbox.run_in_sandbox", + lambda *a, **k: ran.append("ran"), + ) + + result = execution.run_aiida_code("print(1)") + + assert ran == [] + assert "shares storage" in result + assert "do NOT claim to have run it" in result + + def test_a_check_that_cannot_be_made_refuses_too( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Whatever went wrong, separation was not proved.""" + from aiida_agents.tools.codegen import execution + + monkeypatch.setattr( + "aiida_agents.sandbox.setup.sandbox_profile_exists", + lambda profile: True, + ) + + def _boom(config: object, name: str) -> list[object]: + raise RuntimeError("no config here") + + monkeypatch.setattr("aiida_agents.sandbox.copy.profiles_sharing_storage", _boom) + ran: list[str] = [] + monkeypatch.setattr( + "aiida_agents.sandbox.run_in_sandbox", + lambda *a, **k: ran.append("ran"), + ) + + assert "not run" in execution.run_aiida_code("print(1)") + assert ran == [] + def test_it_runs_against_the_configured_profile( self, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -141,6 +200,10 @@ def test_it_runs_against_the_configured_profile( "aiida_agents.sandbox.setup.sandbox_profile_exists", lambda profile: True, ) + monkeypatch.setattr( + "aiida_agents.sandbox.copy.profiles_sharing_storage", + lambda config, name: [], + ) seen: dict[str, t.Any] = {} class _Result: diff --git a/tests/cli/test_commands.py b/tests/cli/test_commands.py index f086d4f..dff7f91 100644 --- a/tests/cli/test_commands.py +++ b/tests/cli/test_commands.py @@ -2,6 +2,8 @@ from __future__ import annotations +from pathlib import Path + import pytest from click.testing import CliRunner @@ -368,6 +370,9 @@ def get_profile(self, name: object = None) -> object: def delete_profile(self, name: str, delete_storage: bool = True) -> None: self.deleted = {"name": name, "delete_storage": delete_storage} + def add_profile(self, profile: object) -> None: + self.profiles.append(profile) + def store(self) -> None: return None @@ -384,6 +389,15 @@ class _Profile: profile.name = name # type: ignore[attr-defined] profile.storage_backend = backend # type: ignore[attr-defined] profile.storage_config = {"filepath": filepath or f"/data/{name}"} # type: ignore[attr-defined] + # What `sandbox_profile_dictionary` clones the sandbox's own from. + profile.dictionary = { # type: ignore[attr-defined] + "storage": {"backend": backend, "config": profile.storage_config}, # type: ignore[attr-defined] + "process_control": {"backend": None, "config": None}, + "default_user_email": "someone@example.com", + "PROFILE_UUID": "1111", + "options": {}, + "test_profile": False, + } return profile def _patch(self, monkeypatch: pytest.MonkeyPatch, config: object) -> None: @@ -421,16 +435,92 @@ def test_check_fails_on_a_backend_it_cannot_reason_about( self, monkeypatch: pytest.MonkeyPatch ) -> None: """Fails closed. Not knowing whether two profiles are separate is not - the same as knowing they are, and only one of those mistakes is safe.""" + the same as knowing they are, and only one of those mistakes is safe. + + The stand-in has to be a backend that genuinely does not exist. This + test originally used `core.sqlite_zip`, which is a real built-in + backend for imported archives -- so it asserted that a perfectly + ordinary archive profile was unreadable, which is the bug #90 fixes + rather than the behaviour this test is for. + """ config = self._config( self._profile("real", filepath="/data/real"), - self._profile("agents-sandbox", backend="core.sqlite_zip"), + self._profile("agents-sandbox", backend="thirdparty.custom_dos"), + ) + self._patch(monkeypatch, config) + + result = CliRunner().invoke(cli, ["sandbox", "check"]) + + assert result.exit_code == 1 + + def test_check_says_when_it_could_not_look_rather_than_what_it_found( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A researcher reading "deleting this destroys their data" concludes + their sandbox is pointed at their own database. Where the truth is that + a backend could not be read, the message has to say so: the first is + theirs to fix, the second is not.""" + config = self._config( + self._profile("odd", backend="thirdparty.custom_dos"), + self._profile("agents-sandbox", filepath="/data/copy"), ) self._patch(monkeypatch, config) result = CliRunner().invoke(cli, ["sandbox", "check"]) + output = " ".join(result.output.split()) assert result.exit_code == 1 + assert "cannot be told apart from 'odd'" in output + assert "thirdparty.custom_dos" in output + assert "destroys" not in output + + def test_check_reports_both_kinds_of_finding_when_both_are_present( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Deleting and rebuilding fixes the overlap and not the other half. + + `sandbox init` compares the copy against its source, so the rebuild + succeeds and the next `check` fails again on the profile that could not + be read. Telling the user only about the overlap sends them round that + loop with no idea why it did not take. + """ + config = self._config( + self._profile("real", filepath="/data/shared"), + self._profile("odd", backend="thirdparty.custom_dos"), + self._profile("agents-sandbox", filepath="/data/shared"), + ) + self._patch(monkeypatch, config) + + result = CliRunner().invoke(cli, ["sandbox", "check"]) + output = " ".join(result.output.split()) + + assert result.exit_code == 1 + assert "shares storage with 'real'" in output + assert "cannot be told apart from 'odd'" in output + assert "verdi profile delete --keep-data" in output + assert "treated as sharing until it can be ruled out" in output + + def test_check_passes_a_config_holding_an_archive_profile( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """An imported archive is a built-in backend, not an unreadable one. + + Until `storage_locations` knew it, every config holding one made this + command name profiles the sandbox had never touched, and `teardown` + refused to remove the sandbox at all. + """ + config = self._config( + self._profile("real", filepath="/data/real"), + self._profile( + "dev-archive", backend="core.sqlite_zip", filepath="/data/e.aiida" + ), + self._profile("agents-sandbox", filepath="/data/copy"), + ) + self._patch(monkeypatch, config) + + result = CliRunner().invoke(cli, ["sandbox", "check"]) + + assert result.exit_code == 0 def test_init_refuses_a_backend_it_cannot_copy( self, monkeypatch: pytest.MonkeyPatch @@ -444,17 +534,100 @@ def test_init_refuses_a_backend_it_cannot_copy( assert result.exit_code != 0 assert "cannot copy" in result.output + def _real_storage(self, tmp_path: Path) -> Path: + """A storage directory with something in it worth counting.""" + storage = tmp_path / "real-storage" + (storage / "container").mkdir(parents=True) + (storage / "database.sqlite").write_bytes(b"x" * 4096) + return storage + + def test_init_says_what_the_copy_costs_and_waits_for_an_answer( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """The copy is the whole repository, gigabytes on a real profile. + + Finding that out from `df` afterwards is the behaviour this replaces, + so the size and the destination go out first and "n" stops it before + anything is written. + """ + storage = self._real_storage(tmp_path) + config = self._config(self._profile("real", filepath=str(storage))) + self._patch(monkeypatch, config) + monkeypatch.setattr( + "aiida_agents.sandbox.copy.sandbox_storage_root", + lambda name: tmp_path / "sandbox" / name, + ) + + result = CliRunner().invoke( + cli, ["sandbox", "init", "--profile", "real"], input="n\n" + ) + output = " ".join(result.output.split()) + + assert result.exit_code == 0, "declining is an answer, not an error" + assert "4.1 kB" in output + assert str(tmp_path / "sandbox" / "agents-sandbox" / "storage") in output + assert not (tmp_path / "sandbox").exists(), "nothing may be copied on 'n'" + + def test_init_yes_copies_without_asking( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """Scripted use has to stay possible, the same way `teardown --yes` does.""" + storage = self._real_storage(tmp_path) + config = self._config(self._profile("real", filepath=str(storage))) + self._patch(monkeypatch, config) + monkeypatch.setattr( + "aiida_agents.sandbox.copy.sandbox_storage_root", + lambda name: tmp_path / "sandbox" / name, + ) + + result = CliRunner().invoke( + cli, ["sandbox", "init", "--profile", "real", "--yes"] + ) + + assert result.exit_code == 0 + copy = tmp_path / "sandbox" / "agents-sandbox" / "storage" + assert (copy / "database.sqlite").read_bytes() == b"x" * 4096 + + def test_init_refuses_an_overlapping_target_before_copying_anything( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """A source that contains the sandbox root would be copied into itself. + + The separation check used to run after the copy, so the refusal arrived + once the gigabytes were already on disk, and left them there. + """ + storage = self._real_storage(tmp_path) + config = self._config(self._profile("real", filepath=str(storage))) + self._patch(monkeypatch, config) + # The sandbox root inside the source: the shape that copies into itself. + monkeypatch.setattr( + "aiida_agents.sandbox.copy.sandbox_storage_root", + lambda name: storage / "agents-sandbox" / name, + ) + + result = CliRunner().invoke( + cli, ["sandbox", "init", "--profile", "real", "--yes"] + ) + + assert result.exit_code != 0 + assert "shares storage" in result.output + assert not (storage / "agents-sandbox").exists(), "nothing may be copied" + def test_init_refuses_to_overwrite_an_existing_sandbox( self, monkeypatch: pytest.MonkeyPatch ) -> None: - """Rebuilding is teardown then init, so nothing in use is copied over.""" + """Rebuilding is teardown then init, so nothing in use is copied over. + + Named in the message, because "already exists" without a way forward is + where a reader stops. + """ config = self._config(self._profile("real"), self._profile("agents-sandbox")) self._patch(monkeypatch, config) result = CliRunner().invoke(cli, ["sandbox", "init", "--profile", "real"]) assert result.exit_code != 0 - assert "refresh" in result.output + assert "teardown" in result.output def test_teardown_refuses_a_sandbox_that_shares_storage( self, monkeypatch: pytest.MonkeyPatch @@ -492,6 +665,159 @@ def test_teardown_never_asks_aiida_to_delete_the_storage( assert result.exit_code == 0 assert config.deleted["delete_storage"] is False # type: ignore[attr-defined] + @staticmethod + def _postgres_profile(name: str) -> object: + from tests.cli.test_commands import TestSandboxCommands as _Self + + profile = _Self._profile(name, backend="core.psql_dos") + profile.storage_config = { # type: ignore[attr-defined] + "database_name": f"aiida_{name}", + "database_hostname": "localhost", + "database_port": 5432, + "database_username": "aiida", + "database_password": "pw", + "repository_uri": f"file:///data/{name}/repository", + } + return profile + + def test_init_makes_the_postgres_copy_rather_than_printing_it( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Five commands and five GRANTs is a setup step nobody performs, and a + setup step nobody performs is a feature nobody has.""" + self._patch(monkeypatch, self._config(self._postgres_profile("real"))) + created: list[tuple[object, str]] = [] + copied: list[object] = [] + monkeypatch.setattr( + "aiida_agents.sandbox.postgres.create_database", + lambda config, database: created.append((config, database)), + ) + monkeypatch.setattr( + "aiida_agents.sandbox.postgres.copy_database", + lambda command, config: copied.append(command), + ) + monkeypatch.setattr( + "aiida_agents.cli.sandbox._database_exists", lambda config: False + ) + + result = CliRunner().invoke( + cli, ["sandbox", "init", "--profile", "real", "--yes"] + ) + + assert [database for _, database in created] == ["aiida_real_agents_sandbox"] + assert [stage[0] for stage in copied[0].stages] == ["pg_dump", "psql"] + assert "Copied into" in result.output + + def test_init_falls_back_to_printing_when_it_cannot( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A managed server, a remote host, no sudo. The printed commands were + the whole feature until now and are still the only road there.""" + from aiida_agents.sandbox.postgres import PostgresUnavailableError + + self._patch(monkeypatch, self._config(self._postgres_profile("real"))) + + def _refuse(config: object, database: str) -> None: + raise PostgresUnavailableError("no superuser connection here") + + monkeypatch.setattr("aiida_agents.sandbox.postgres.create_database", _refuse) + monkeypatch.setattr( + "aiida_agents.cli.sandbox._database_exists", lambda config: False + ) + + result = CliRunner().invoke( + cli, ["sandbox", "init", "--profile", "real", "--yes"] + ) + output = " ".join(result.output.split()) + + assert "no superuser connection here" in output + assert "createdb" in output + assert "pg_dump" in output + + def test_init_prints_postgres_commands_that_can_be_pasted( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """On the fallback path these exist only to be pasted, so a line break + is a broken command. + + Rich reflowed the `pg_dump | psql` pipeline at the terminal width and + put the database name on a line of its own, where pasting it runs the + name as a command that does not exist. + """ + from aiida_agents.sandbox.copy import postgres_copy_commands + from aiida_agents.sandbox.postgres import PostgresUnavailableError + + source = self._postgres_profile("real") + source.storage_config["database_name"] = ( # type: ignore[attr-defined] + "aiida-a-rather-long-database-name-as-they-tend-to-be" + ) + self._patch(monkeypatch, self._config(source)) + + def _refuse(config: object, database: str) -> None: + raise PostgresUnavailableError("nothing here can create a database") + + monkeypatch.setattr("aiida_agents.sandbox.postgres.create_database", _refuse) + monkeypatch.setattr( + "aiida_agents.cli.sandbox._database_exists", lambda config: False + ) + + result = CliRunner().invoke( + cli, ["sandbox", "init", "--profile", "real", "--yes"] + ) + + for command in postgres_copy_commands( + source.storage_config, # type: ignore[attr-defined] + "aiida-a-rather-long-database-name-as-they-tend-to-be_agents_sandbox", + ): + assert command.as_shell() in result.output, "reflowed, cannot be pasted" + + def test_teardown_names_the_storage_it_is_about_to_remove( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """It asked "delete profile X and its copy?", which reads as a copy of + the profile: the one thing it does not delete. Both halves are named + now, and the storage half by size and path.""" + storage = self._real_storage(tmp_path) + config = self._config( + self._profile("real", filepath="/data/real"), + self._profile("agents-sandbox", filepath=str(storage)), + ) + self._patch(monkeypatch, config) + monkeypatch.setattr( + "aiida_agents.sandbox.copy.sandbox_storage_root", lambda name: storage + ) + + result = CliRunner().invoke(cli, ["sandbox", "teardown"], input="n\n") + output = " ".join(result.output.split()) + + assert "4.1 kB" in output + assert str(storage) in output + assert storage.exists(), "nothing may be removed on 'n'" + + def test_teardown_says_what_it_cannot_reach_in_postgres( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The copied database lives in the server, so this deletes the profile + and the repository and leaves it. `init` printed the commands that made + it; leaving the one that removes it unsaid is how a database nobody + accounts for stays there.""" + sandbox = self._profile("agents-sandbox", backend="core.psql_dos") + sandbox.storage_config = { # type: ignore[attr-defined] + "database_name": "aiida_db_agents_sandbox", + "database_hostname": "localhost", + "database_port": 5432, + "database_username": "aiida", + "repository_uri": "file:///data/sandbox/repository", + } + config = self._config(self._profile("real", filepath="/data/real"), sandbox) + self._patch(monkeypatch, config) + + result = CliRunner().invoke(cli, ["sandbox", "teardown", "--yes"]) + output = " ".join(result.output.split()) + + assert "dropdb" in output + assert "aiida_db_agents_sandbox" in output + def test_teardown_on_a_missing_sandbox_is_not_an_error( self, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/cli/test_doctor.py b/tests/cli/test_doctor.py index 2ac853a..debc77c 100644 --- a/tests/cli/test_doctor.py +++ b/tests/cli/test_doctor.py @@ -281,3 +281,24 @@ def test_a_sandbox_sharing_storage_fails_the_check( assert row.ok is False assert "real" in row.detail + + +def test_a_sharing_sandbox_is_not_sent_to_a_command_that_would_refuse( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """This row used to end in "rebuild with `aiida-agents sandbox refresh`". + + Teardown refuses a sandbox in exactly this state, so that advice sent the + reader to a command that could not run, on the one row where they most need + a next step that works. There is no `refresh` at all now, which makes the + assertion cheap and worth keeping: the row must not name a remedy that + cannot be carried out. + """ + _patch_all_checks_passing(monkeypatch) + monkeypatch.setattr( + "aiida.manage.configuration.get_config", + lambda: _SandboxConfig("agents-sandbox", filepath="/data/real"), + ) + row = _rows_by_label()["Codegen sandbox (disposable copy)"] + + assert "refresh" not in row.detail diff --git a/tests/sandbox/test_copy.py b/tests/sandbox/test_copy.py index 1e20110..7fb8447 100644 --- a/tests/sandbox/test_copy.py +++ b/tests/sandbox/test_copy.py @@ -11,26 +11,70 @@ from __future__ import annotations +import shlex from pathlib import Path +from types import SimpleNamespace import pytest from aiida_agents.sandbox.copy import ( - SandboxStorage, + Location, + Overlap, + PathLocation, + ProfileStorage, + StorageConfig, + SharingProfile, copy_sqlite_storage, postgres_copy_commands, + postgres_sandbox_storage, + profiles_sharing_storage, sandbox_profile_dictionary, shares_storage, storage_locations, + storage_overlap, + storage_size, ) SQLITE = "core.sqlite_dos" POSTGRES = "core.psql_dos" +ARCHIVE = "core.sqlite_zip" +UNREADABLE = "thirdparty.custom_dos" + + +def _shares( + backend_a: str, + config_a: StorageConfig, + backend_b: str, + config_b: StorageConfig, +) -> bool: + """`shares_storage` spelled from the two halves each case is written in. + + The module compares `ProfileStorage` values, which is what stops a caller + pairing one profile's backend with another's config. These tests are about + which configurations share storage, so they name the halves and let this + do the pairing once. + """ + return shares_storage( + ProfileStorage(backend_a, config_a), ProfileStorage(backend_b, config_b) + ) + + +def _locations(backend: str, config: StorageConfig) -> frozenset[Location]: + return storage_locations(ProfileStorage(backend, config)) + + +def _overlap( + backend_a: str, + config_a: StorageConfig, + backend_b: str, + config_b: StorageConfig, +) -> Overlap: + return storage_overlap( + ProfileStorage(backend_a, config_a), ProfileStorage(backend_b, config_b) + ) -def _pg( - database: str = "aiida_db", repository: str = "/data/repo" -) -> dict[str, object]: +def _pg(database: str = "aiida_db", repository: str = "/data/repo") -> StorageConfig: return { "database_name": database, "database_hostname": "localhost", @@ -47,10 +91,10 @@ class TestSharingIsRefused: def test_the_same_sqlite_directory_is_sharing(self) -> None: config = {"filepath": "/data/storage"} - assert shares_storage(SQLITE, config, SQLITE, dict(config)) + assert _shares(SQLITE, config, SQLITE, dict(config)) def test_the_same_postgres_database_is_sharing(self) -> None: - assert shares_storage(POSTGRES, _pg(), POSTGRES, _pg()) + assert _shares(POSTGRES, _pg(), POSTGRES, _pg()) def test_a_separate_database_but_the_same_repository_is_still_sharing(self) -> None: """Half a copy is not a copy. @@ -58,7 +102,7 @@ def test_a_separate_database_but_the_same_repository_is_still_sharing(self) -> N A sandbox with its own database and the user's repository still loses them every file their nodes refer to. """ - assert shares_storage( + assert _shares( POSTGRES, _pg(database="aiida_db"), POSTGRES, @@ -66,7 +110,7 @@ def test_a_separate_database_but_the_same_repository_is_still_sharing(self) -> N ) def test_a_genuinely_separate_copy_is_not_sharing(self) -> None: - assert not shares_storage( + assert not _shares( POSTGRES, _pg(database="aiida_db", repository="/data/repo"), POSTGRES, @@ -74,7 +118,7 @@ def test_a_genuinely_separate_copy_is_not_sharing(self) -> None: ) def test_separate_sqlite_directories_are_not_sharing(self) -> None: - assert not shares_storage( + assert not _shares( SQLITE, {"filepath": "/data/real"}, SQLITE, {"filepath": "/data/copy"} ) @@ -84,7 +128,7 @@ def test_the_same_directory_written_differently_is_still_sharing(self) -> None: `/data/storage` and `/data/./sub/../storage` are one directory, and a string comparison would call them two. """ - assert shares_storage( + assert _shares( SQLITE, {"filepath": "/data/storage"}, SQLITE, @@ -103,8 +147,9 @@ class TestFailingClosed: @pytest.mark.parametrize( "backend, config", [ - pytest.param("core.sqlite_zip", {"filepath": "/x"}, id="unknown-backend"), + pytest.param(UNREADABLE, {"filepath": "/x"}, id="unknown-backend"), pytest.param(SQLITE, {}, id="sqlite-with-no-path"), + pytest.param(ARCHIVE, {}, id="archive-with-no-path"), pytest.param(POSTGRES, {}, id="postgres-with-nothing"), pytest.param( POSTGRES, {"database_name": "aiida"}, id="postgres-with-no-repository" @@ -112,16 +157,302 @@ class TestFailingClosed: pytest.param( POSTGRES, {"repository_uri": "/r"}, id="postgres-with-no-database" ), + pytest.param(SQLITE, {"filepath": "storage"}, id="sqlite-relative-path"), + pytest.param(ARCHIVE, {"filepath": "export.aiida"}, id="archive-relative"), + pytest.param( + POSTGRES, + {"database_name": "aiida", "repository_uri": "repo"}, + id="postgres-relative-repository", + ), + pytest.param(SQLITE, {"filepath": 42}, id="sqlite-with-a-number"), + pytest.param( + SQLITE, {"filepath": Path("/data/real")}, id="sqlite-with-a-path-object" + ), ], ) def test_an_unreadable_config_counts_as_sharing( self, backend: str, config: dict[str, object] ) -> None: - assert shares_storage(backend, config, SQLITE, {"filepath": "/somewhere/else"}) - assert shares_storage(SQLITE, {"filepath": "/somewhere/else"}, backend, config) + """Relative and non-string paths are in here deliberately. + + `Path.resolve` would answer a relative path against the working + directory, so the same two profiles would compare differently depending + on where the command was run, and `Path(42)` raises outright. Neither + may become "these are separate", and neither may reach the user as a + traceback out of the one check that guards their data. + + A `Path` object is refused for the same reason, even though it is a + perfectly good path: these configurations are read from `config.json`, + where every path is a string, so a value of any other type means the + configuration is not what this knows how to compare. + """ + assert _shares(backend, config, SQLITE, {"filepath": "/somewhere/else"}) + assert _shares(SQLITE, {"filepath": "/somewhere/else"}, backend, config) def test_an_unreadable_config_yields_no_locations(self) -> None: - assert storage_locations("core.sqlite_zip", {"filepath": "/x"}) == frozenset() + assert _locations(UNREADABLE, {"filepath": "/x"}) == frozenset() + + +class TestArchiveProfilesAreNotFalsePositives: + """`core.sqlite_zip` archive profiles are a common, recognised backend. + + Treating them as unknown made `sandbox check` and `sandbox teardown` fail + against any config holding an imported archive: the sandbox was reported as + sharing storage with a profile it had never touched, and teardown refused + to remove it. The archive is a single file at its own path; the sandbox is + a directory at another. They cannot overlap. + """ + + def test_an_archive_yields_its_own_location(self) -> None: + assert _locations(ARCHIVE, {"filepath": "/data/export.aiida"}) == frozenset( + {PathLocation(Path("/data/export.aiida"))} + ) + + def test_a_sqlite_sandbox_does_not_share_with_an_archive(self) -> None: + assert not _shares( + SQLITE, + {"filepath": "/data/agents-sandbox/storage"}, + ARCHIVE, + {"filepath": "/data/export.aiida"}, + ) + + def test_a_postgres_sandbox_does_not_share_with_an_archive(self) -> None: + assert not _shares(POSTGRES, _pg(), ARCHIVE, {"filepath": "/data/export.aiida"}) + + def test_two_profiles_reading_the_same_archive_still_share(self) -> None: + """An archive is read-only in every respect but the one that matters: + `SqliteZipBackend.delete` unlinks the file, so `--delete-data` on one of + these two profiles takes the other's storage with it.""" + assert _shares( + ARCHIVE, + {"filepath": "/data/export.aiida"}, + ARCHIVE, + {"filepath": "/data/./export.aiida"}, + ) + + +class TestOnePathIsOnePath: + """A location is a path, never a path plus what a profile calls it. + + The tag in front of a location used to name the kind of storage --- `dir:` + for a sqlite_dos directory, `file:` for an archive --- so two profiles + pointing at one path compared as separate whenever they disagreed about its + kind. That is the one case where the answer must be "sharing": whichever of + them is wrong about the path, deleting either destroys what the other reads. + """ + + @pytest.mark.parametrize( + "backend, config", + [ + pytest.param(SQLITE, {"filepath": "/data/thing"}, id="sqlite-dos"), + pytest.param( + POSTGRES, + {"database_name": "aiida", "repository_uri": "/data/thing"}, + id="postgres-repository", + ), + pytest.param( + POSTGRES, + {"database_name": "aiida", "repository_uri": "file:///data/thing"}, + id="postgres-repository-as-a-uri", + ), + ], + ) + def test_anything_else_at_an_archives_path_is_sharing( + self, backend: str, config: dict[str, object] + ) -> None: + assert _shares(backend, config, ARCHIVE, {"filepath": "/data/thing"}) + + def test_a_percent_encoded_repository_is_the_directory_it_names(self) -> None: + """`verdi profile setup` writes `repository_uri` with `Path.as_uri()`. + + A repository under a path with a space in it is therefore stored as + `file:///home/u/My%20Drive/...`, and stripping the scheme by hand + leaves a `%20` that no other spelling of that directory can match. The + one profile that writes the plain form is our own sandbox, which is + precisely the pair that must not compare as separate. + """ + encoded = Path("/home/u/My Drive/repository").as_uri() + + assert _shares( + POSTGRES, + {"database_name": "real", "repository_uri": encoded}, + POSTGRES, + { + "database_name": "sandbox", + "repository_uri": "/home/u/My Drive/repository", + }, + ) + + +class TestContainmentCountsAsSharing: + """One storage inside another is not two separate storages. + + `sandbox teardown` removes its root with `shutil.rmtree`, which is + recursive, so a profile whose storage sits under that root is deleted with + it. Comparing the two paths for equality called them separate right up + until the moment one destroyed the other. + """ + + SANDBOX = "/data/agents-sandbox/storage" + + def test_an_archive_inside_the_sandbox_directory_is_sharing(self) -> None: + assert _shares( + SQLITE, + {"filepath": self.SANDBOX}, + ARCHIVE, + {"filepath": f"{self.SANDBOX}/export.aiida"}, + ) + + def test_it_holds_whichever_way_round_the_two_are_given(self) -> None: + assert _shares( + ARCHIVE, + {"filepath": f"{self.SANDBOX}/export.aiida"}, + SQLITE, + {"filepath": self.SANDBOX}, + ) + + def test_a_postgres_repository_containing_a_profile_is_sharing(self) -> None: + assert _shares( + POSTGRES, + {"database_name": "aiida", "repository_uri": "/data/repo"}, + SQLITE, + {"filepath": "/data/repo/nested"}, + ) + + def test_a_name_that_merely_starts_the_same_is_not(self) -> None: + """`/data/storage-2` is not inside `/data/storage`, and a prefix + comparison on the strings would say it was.""" + assert not _shares( + SQLITE, + {"filepath": "/data/storage"}, + SQLITE, + {"filepath": "/data/storage-2"}, + ) + + +class TestSayingWhichKindOfNotSeparate: + """`shares_storage` decides; `storage_overlap` says why. + + Both readings stop the same commands, and they need different words in + front of the user: "this is the same directory as your own profile" is + their problem to fix, "this backend is one I cannot read" is not. + """ + + def test_the_same_location_is_a_proven_overlap(self) -> None: + config = {"filepath": "/data/storage"} + + assert _overlap(SQLITE, config, SQLITE, dict(config)) is Overlap.SHARED + + def test_different_locations_are_separate(self) -> None: + assert ( + _overlap( + SQLITE, {"filepath": "/data/real"}, SQLITE, {"filepath": "/data/copy"} + ) + is Overlap.SEPARATE + ) + + @pytest.mark.parametrize( + "backend, config", + [ + pytest.param(UNREADABLE, {"filepath": "/x"}, id="unknown-backend"), + pytest.param(SQLITE, {}, id="no-path-to-compare"), + ], + ) + def test_what_cannot_be_read_is_unknown_rather_than_shared( + self, backend: str, config: dict[str, object] + ) -> None: + assert ( + _overlap(SQLITE, {"filepath": "/data/real"}, backend, config) + is Overlap.UNKNOWN + ) + + def test_shares_storage_acts_on_unknown_exactly_as_on_shared(self) -> None: + """The distinction is for the message, never for the decision.""" + assert _shares(SQLITE, {"filepath": "/data/real"}, UNREADABLE, {}) + + +class _Profile: + """The three attributes `profiles_sharing_storage` reads off a profile.""" + + def __init__(self, name: str, backend: str, config: dict[str, object]) -> None: + self.name = name + self.storage_backend = backend + self.storage_config = config + + +class _Config: + def __init__(self, *profiles: _Profile) -> None: + self.profiles = list(profiles) + + +class TestWhatTheSandboxCouldNotBeClearedOf: + """The list `check`, `teardown` and `doctor` all report from.""" + + @pytest.fixture + def sandbox(self) -> _Profile: + return _Profile("agents-sandbox", SQLITE, {"filepath": "/data/copy"}) + + def test_a_separate_profile_is_not_listed(self, sandbox: _Profile) -> None: + config = _Config(sandbox, _Profile("real", SQLITE, {"filepath": "/r"})) + + assert profiles_sharing_storage(config, "agents-sandbox") == [] + + def test_an_archive_profile_is_not_listed(self, sandbox: _Profile) -> None: + """The bug this suite grew out of, at the layer that reported it.""" + config = _Config( + sandbox, + _Profile("real", SQLITE, {"filepath": "/data/real"}), + _Profile("dev-archive", ARCHIVE, {"filepath": "/data/export.aiida"}), + ) + + assert profiles_sharing_storage(config, "agents-sandbox") == [] + + def test_a_shared_profile_is_listed_with_the_overlap_proved( + self, sandbox: _Profile + ) -> None: + config = _Config(sandbox, _Profile("real", SQLITE, {"filepath": "/data/copy"})) + + assert profiles_sharing_storage(config, "agents-sandbox") == [ + SharingProfile(name="real", backend=SQLITE, overlap=Overlap.SHARED) + ] + + def test_an_unreadable_profile_is_listed_as_unknown( + self, sandbox: _Profile + ) -> None: + config = _Config(sandbox, _Profile("odd", UNREADABLE, {"filepath": "/x"})) + + assert profiles_sharing_storage(config, "agents-sandbox") == [ + SharingProfile(name="odd", backend=UNREADABLE, overlap=Overlap.UNKNOWN) + ] + + def test_a_proven_overlap_reads_as_one(self) -> None: + reason = SharingProfile( + name="real", backend=SQLITE, overlap=Overlap.SHARED + ).describe() + + assert "shares storage with 'real'" in reason + + @pytest.mark.parametrize( + "backend", + [ + pytest.param(UNREADABLE, id="a-backend-nobody-has-heard-of"), + pytest.param(POSTGRES, id="a-backend-we-know-but-a-config-we-cannot-read"), + ], + ) + def test_what_could_not_be_read_is_not_blamed_on_the_backend( + self, backend: str + ) -> None: + """The unreadable half is as often a config with a field missing, and + it can be either profile's. Saying "this backend cannot be read" would + be false for a `core.psql_dos` profile with no repository, which is the + commoner way to land here than a third-party plugin.""" + reason = SharingProfile( + name="odd", backend=backend, overlap=Overlap.UNKNOWN + ).describe() + + assert "cannot be told apart from 'odd'" in reason + assert backend in reason + assert "destroy" not in reason class TestCopyingSqliteStorage: @@ -136,6 +467,64 @@ def storage(self, tmp_path: Path) -> Path: (source / "container" / "config.json").write_text("{}") return source + def test_the_size_is_what_the_copy_will_write(self, storage: Path) -> None: + """What the caller shows before asking, and counts the bar up to.""" + expected = sum( + item.stat().st_size for item in storage.rglob("*") if item.is_file() + ) + + assert storage_size(storage) == expected + + def test_a_missing_directory_has_no_size_to_report(self, tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError): + storage_size(tmp_path / "nope") + + def test_progress_is_reported_in_bytes_as_the_copy_runs( + self, storage: Path, tmp_path: Path + ) -> None: + """A packed container is a few very large files, so a caller counting + files would show nothing until it was nearly done.""" + reported: list[int] = [] + + copy_sqlite_storage(storage, tmp_path / "copy", progress=reported.append) + + assert sum(reported) == storage_size(storage) + assert all(chunk > 0 for chunk in reported) + + def test_a_copy_with_no_room_is_refused_before_it_starts( + self, storage: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Filling the disk mid-copy leaves a partial tree that the next + attempt then refuses as something already in use, so the room is + checked while refusing still costs nothing.""" + target = tmp_path / "copy" + monkeypatch.setattr( + "shutil.disk_usage", lambda path: SimpleNamespace(total=0, used=0, free=1) + ) + + with pytest.raises(OSError, match="free"): + copy_sqlite_storage(storage, target) + + assert not target.exists() + + def test_a_failed_copy_leaves_nothing_behind( + self, storage: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Half a copy answers queries wrongly and blocks the next attempt.""" + target = tmp_path / "copy" + + def _fail(*args: object, **kwargs: object) -> None: + target.mkdir(parents=True, exist_ok=True) + (target / "half-written").write_bytes(b"...") + raise OSError("no space left on device") + + monkeypatch.setattr("shutil.copytree", _fail) + + with pytest.raises(OSError): + copy_sqlite_storage(storage, target) + + assert not target.exists() + def test_the_database_and_the_repository_both_come_across( self, storage: Path, tmp_path: Path ) -> None: @@ -199,31 +588,98 @@ def test_it_dumps_rather_than_using_a_template(self) -> None: time. `pg_dump` works against a database in use. """ commands = " ".join( - command for _, command in postgres_copy_commands(_pg(), "s") + command.as_shell() for command in postgres_copy_commands(_pg(), "s") ) assert "pg_dump" in commands assert "TEMPLATE" not in commands.upper() def test_the_database_is_created_before_it_is_filled(self) -> None: - steps = [command for _, command in postgres_copy_commands(_pg(), "sandbox_db")] + create, copy = postgres_copy_commands(_pg(), "sandbox_db") + + assert "createdb" in create.as_shell() + assert "pg_dump" in copy.as_shell() - assert "createdb" in steps[0] - assert "pg_dump" in steps[1] + def test_the_copy_is_a_pipeline_of_two_programs(self) -> None: + """Run as two processes, printed as one line with a pipe. Both come + from the same place so what `init` runs is what it would print.""" + _, copy = postgres_copy_commands(_pg(), "sandbox_db") + + assert [stage[0] for stage in copy.stages] == ["pg_dump", "psql"] + assert copy.as_shell().count(" | ") == 1 def test_every_step_explains_itself(self) -> None: """These are pasted into a terminal by hand; an unexplained one is a command somebody runs without knowing what it does.""" - assert all(explanation for explanation, _ in postgres_copy_commands(_pg(), "s")) + assert all( + command.explanation for command in postgres_copy_commands(_pg(), "s") + ) - def test_awkward_database_names_are_quoted(self) -> None: - commands = " ".join( - command - for _, command in postgres_copy_commands(_pg(database="gsoc-psql"), "s-box") + def test_awkward_database_names_survive_the_round_trip(self) -> None: + """Quoted by `shlex`, so the line means what the pieces meant. + + A name with a space or a quote in it is legal in Postgres, and hand- + rolled double quotes produced a line that ran something else. + """ + commands = postgres_copy_commands(_pg(database='we "ird" db'), "s box") + + for command in commands: + for stage in command.stages: + assert shlex.split(shlex.join(stage)) == list(stage) + assert ( + 'we "ird" db' + in postgres_copy_commands(_pg(database='we "ird" db'), "s")[1].stages[0] + ) + + +class TestThePostgresSandboxStorage: + """What `sandbox init` registers for a PostgreSQL sandbox.""" + + REPOSITORY = Path("/data/agents-sandbox/repository") + + def test_aiida_can_read_the_repository_it_writes(self) -> None: + """A plain path here made the registered profile unopenable. + + `repository_uri` is a URL to aiida-core, not a path: given + `/data/...` rather than `file:///data/...`, `get_filepath_container` + raises `ConfigurationError` and the sandbox profile cannot be loaded, + which is the one thing it exists to be. Asserted against the function + that raised rather than against the string, so it stays true if + aiida-core changes how it reads the field. + """ + from aiida.manage.configuration.profile import Profile + from aiida.storage.psql_dos.backend import get_filepath_container + + storage = postgres_sandbox_storage( + _pg(), database="aiida_db_sandbox", repository=self.REPOSITORY + ) + profile = Profile( + "agents-sandbox", + { + "storage": {"backend": storage.backend, "config": storage.config}, + "process_control": {"backend": None, "config": None}, + }, + ) + + assert get_filepath_container(profile) == self.REPOSITORY / "container" + + def test_the_server_is_the_source_profile_s_and_the_database_is_not(self) -> None: + """The copy lives in the same server, under its own name.""" + storage = postgres_sandbox_storage( + _pg(), database="aiida_db_sandbox", repository=self.REPOSITORY + ) + + assert storage.config["database_name"] == "aiida_db_sandbox" + assert storage.config["database_hostname"] == "localhost" + assert storage.config["database_username"] == "aiida" + + def test_it_does_not_share_storage_with_the_source(self) -> None: + source = _pg() + storage = postgres_sandbox_storage( + source, database="aiida_db_sandbox", repository=self.REPOSITORY ) - assert '"gsoc-psql"' in commands - assert '"s-box"' in commands + assert not _shares(POSTGRES, source, storage.backend, storage.config) class TestTheClonedProfile: @@ -239,22 +695,22 @@ def source(self) -> dict[str, object]: } @pytest.fixture - def storage(self) -> SandboxStorage: - return SandboxStorage(SQLITE, {"filepath": "/data/copy"}) + def storage(self) -> ProfileStorage: + return ProfileStorage(SQLITE, {"filepath": "/data/copy"}) def test_it_points_at_the_copy( - self, source: dict[str, object], storage: SandboxStorage + self, source: dict[str, object], storage: ProfileStorage ) -> None: result = sandbox_profile_dictionary(source, storage) assert result["storage"]["config"]["filepath"] == "/data/copy" def test_it_does_not_point_at_the_source( - self, source: dict[str, object], storage: SandboxStorage + self, source: dict[str, object], storage: ProfileStorage ) -> None: result = sandbox_profile_dictionary(source, storage) - assert not shares_storage( + assert not _shares( SQLITE, {"filepath": "/data/real"}, result["storage"]["backend"], @@ -262,7 +718,7 @@ def test_it_does_not_point_at_the_source( ) def test_the_uuid_is_regenerated( - self, source: dict[str, object], storage: SandboxStorage + self, source: dict[str, object], storage: ProfileStorage ) -> None: """Two profiles sharing a UUID are two profiles AiiDA cannot tell apart.""" result = sandbox_profile_dictionary(source, storage) @@ -270,7 +726,7 @@ def test_the_uuid_is_regenerated( assert result["PROFILE_UUID"] != "1111" def test_the_broker_is_not_carried_over( - self, source: dict[str, object], storage: SandboxStorage + self, source: dict[str, object], storage: ProfileStorage ) -> None: """The sandbox runs nothing, so it needs no queues --- and pointing it at the source profile's would let generated code reach a daemon.""" @@ -279,7 +735,7 @@ def test_the_broker_is_not_carried_over( assert result["process_control"]["backend"] is None def test_the_user_is_carried_over( - self, source: dict[str, object], storage: SandboxStorage + self, source: dict[str, object], storage: ProfileStorage ) -> None: """The copy holds the same users; a different default would not resolve.""" result = sandbox_profile_dictionary(source, storage) diff --git a/tests/sandbox/test_postgres.py b/tests/sandbox/test_postgres.py new file mode 100644 index 0000000..09755ce --- /dev/null +++ b/tests/sandbox/test_postgres.py @@ -0,0 +1,170 @@ +"""Tests for making the PostgreSQL copy rather than printing how to make it. + +The interesting half is what happens when it *cannot*: a managed server, a +remote host, no `sudo`. Every one of those has to land back on the printed +commands rather than on a traceback, because the printed commands were the +whole feature until now and are still the only road on those machines. +""" + +from __future__ import annotations + +import subprocess +import sys + +import pytest + +from aiida_agents.sandbox.copy import ShellCommand, postgres_copy_commands +from aiida_agents.sandbox.postgres import ( + PostgresUnavailableError, + copy_database, + create_database, +) + + +def _pg(database: str = "aiida_db") -> dict[str, object]: + return { + "database_name": database, + "database_hostname": "localhost", + "database_port": 5432, + "database_username": "aiida", + "database_password": "pw", + } + + +class _FakePostgres: + """`aiida.manage.external.postgres.Postgres`, as far as this uses it.""" + + def __init__(self, *, connected: bool = True, existing: bool = False) -> None: + self.is_connected = connected + self._existing = existing + self.created: list[tuple[str, str]] = [] + + def determine_setup(self) -> None: + return None + + def db_exists(self, name: str) -> bool: + return self._existing + + def create_db(self, owner: str, name: str) -> None: + self.created.append((owner, name)) + + +class TestCreatingTheDatabase: + def test_it_creates_one_owned_by_the_profile_s_user( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Owned by them, because they are who fills and reads it afterwards.""" + postgres = _FakePostgres() + monkeypatch.setattr( + "aiida.manage.external.postgres.Postgres", lambda **kwargs: postgres + ) + + create_database(_pg(), "aiida_db_agents_sandbox") + + assert postgres.created == [("aiida", "aiida_db_agents_sandbox")] + + def test_an_existing_database_is_reused_rather_than_refused( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """`init` is rerun after a half-finished setup more often than not.""" + postgres = _FakePostgres(existing=True) + monkeypatch.setattr( + "aiida.manage.external.postgres.Postgres", lambda **kwargs: postgres + ) + + create_database(_pg(), "aiida_db_agents_sandbox") + + assert postgres.created == [] + + @pytest.mark.parametrize( + "postgres, reason", + [ + pytest.param(_FakePostgres(connected=False), "no", id="no-connection"), + pytest.param(None, "boom", id="construction-raised"), + ], + ) + def test_no_privileged_connection_is_reported_not_raised_raw( + self, monkeypatch: pytest.MonkeyPatch, postgres: object, reason: str + ) -> None: + """The caller's fallback is to print the commands, and it needs a + sentence saying why it came to that.""" + + def _build(**kwargs: object) -> object: + if postgres is None: + raise RuntimeError("boom") + return postgres + + monkeypatch.setattr("aiida.manage.external.postgres.Postgres", _build) + + with pytest.raises(PostgresUnavailableError, match=reason): + create_database(_pg(), "sandbox_db") + + +class TestCopyingTheData: + def test_it_pipes_one_program_into_the_next(self) -> None: + """`pg_dump | psql` is two processes, and the second reads the first.""" + command = ShellCommand( + "copy", + ( + (sys.executable, "-c", "print('rows')"), + (sys.executable, "-c", "import sys; assert sys.stdin.read().strip()"), + ), + ) + + copy_database(command, _pg()) + + def test_a_failing_stage_carries_its_own_error(self) -> None: + """ "It didn't work" costs a support round trip; Postgres already said + what was wrong.""" + command = ShellCommand( + "copy", + ( + ( + sys.executable, + "-c", + "import sys; sys.exit('database does not exist')", + ), + ), + ) + + with pytest.raises(PostgresUnavailableError, match="database does not exist"): + copy_database(command, _pg()) + + def test_a_missing_program_is_reported_by_name(self) -> None: + """`pg_dump` absent is the common case on a machine that talks to a + remote server, and the message has to name it.""" + command = ShellCommand( + "copy", (("pg_dump_that_is_not_installed", "--version"),) + ) + + with pytest.raises(PostgresUnavailableError, match="pg_dump_that_is_not"): + copy_database(command, _pg()) + + def test_the_password_is_handed_over_rather_than_prompted_for( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """It is already in `config.json`. Prompting for it would teach people + to type their database password at whatever asks.""" + seen: dict[str, str] = {} + real = subprocess.Popen + + def _spy(*args: object, **kwargs: object) -> object: + seen.update(kwargs.get("env") or {}) # type: ignore[arg-type] + return real(*args, **kwargs) # type: ignore[arg-type] + + monkeypatch.setattr(subprocess, "Popen", _spy) + command = ShellCommand("copy", ((sys.executable, "-c", "pass"),)) + + copy_database(command, _pg()) + + assert seen["PGPASSWORD"] == "pw" + + +def test_the_commands_run_are_the_commands_printed() -> None: + """The fallback prints what the automated path runs, from one builder, so + somebody following the printed route is not following a different one.""" + create, copy = postgres_copy_commands(_pg(), "sandbox_db") + + assert create.as_shell().startswith("createdb ") + assert copy.stages[0][0] == "pg_dump" + assert "sandbox_db" in copy.stages[1]