diff --git a/.github/workflows/_test.yml b/.github/workflows/_test.yml index c8d67faf6..8d76dc72b 100644 --- a/.github/workflows/_test.yml +++ b/.github/workflows/_test.yml @@ -122,6 +122,41 @@ jobs: with: install-only: true + # The Omnigraph connector's live tests drive the real `omnigraph` CLI + # against a temporary store. They gate on OMNIGRAPH_TEST_STORE=1 and + # expect the binary at test/bin/omnigraph (git-ignored), so install + # the pinned release there on every platform it ships for; elsewhere + # the live tests skip as they do on a developer machine without it. + - name: Install Omnigraph CLI for the connector's live tests + if: runner.os != 'Windows' + shell: bash + env: + OMNIGRAPH_VERSION: "0.10.0" + run: | + case "${{ runner.os }}-${{ runner.arch }}" in + Linux-X64) asset=omnigraph-linux-x86_64 ;; + Linux-ARM64) asset=omnigraph-linux-arm64 ;; + macOS-ARM64) asset=omnigraph-macos-arm64 ;; + *) + echo "No Omnigraph release for ${{ runner.os }}-${{ runner.arch }}; live tests will skip" + exit 0 + ;; + esac + base="https://github.com/ModernRelay/omnigraph/releases/download/v${OMNIGRAPH_VERSION}" + mkdir -p test/bin + cd test/bin + curl -sSfL -o "$asset.tar.gz" "$base/$asset.tar.gz" + curl -sSfL -o "$asset.sha256" "$base/$asset.sha256" + if command -v sha256sum >/dev/null; then + sha256sum -c "$asset.sha256" + else + shasum -a 256 -c "$asset.sha256" + fi + tar xzf "$asset.tar.gz" omnigraph + rm "$asset.tar.gz" "$asset.sha256" + ./omnigraph --version + echo "OMNIGRAPH_TEST_STORE=1" >> "$GITHUB_ENV" + - name: Run build-test hooks env: UV_NO_SYNC: "1" diff --git a/dev/omnigraph-rework/CLAUDE.md b/dev/omnigraph-rework/CLAUDE.md new file mode 100644 index 000000000..198cfd016 --- /dev/null +++ b/dev/omnigraph-rework/CLAUDE.md @@ -0,0 +1,70 @@ +# CLAUDE.md + +Project instructions for Claude Code working on the **Omnigraph connector rework** (CocoIndex branch `omnigraph-connector`). + +## Overview + +The Omnigraph target connector currently drives the `omnigraph` CLI against a direct `file://` store. The rework replaces that with HTTP to `omnigraph-server` as the connector's only **data** transport: the CLI data path is deleted, and there is no Python SDK package β€” the HTTP client is private to the connector. One CLI use survives by necessity: `omnigraph cluster apply`, because the server refuses schema writes on every configuration (P1) and `managed_by="system"` must keep working (decision 32). Stage: spec reworked 2026-09-07, Epic A probes P1/P2/P12 answered, no rework code yet; the connector branch is at 37 commits over `main` with a passing live suite on Omnigraph 0.10.0. + +**`dev/omnigraph-rework/SPEC.md` is the task list and source of truth.** Start at the **Current focus** pointer near the top β€” it names the next actionable step so you don't have to scan the whole file. Work the spec: implement that step's deliverable, update its status (πŸ”² β†’ πŸ”„ β†’ βœ…) in the phase tracker, and advance the Current focus pointer. Don't skip ahead past a phase's exit guardrails β€” at a phase boundary, verify the criteria, fill in the **Actual outcome** column, and only then move on. + +The research behind the spec is `dev/omnigraph-connector-transport-analysis.md`. Read it when a step's rationale is unclear; do not re-derive its findings. Note that it predates the rework: its "Scope A: connector-private minimal HTTP client" is what is being built, and its Scope B/C sections (public SDK, TypeScript parity) are no longer the plan. + +## One repository + +Everything lives in cocoindex. Repo-wide conventions come from `AGENTS.md` and apply unchanged; this file adds only what is specific to the rework. There is no second repository, no PyPI project, and no published package β€” if a step seems to call for one, the spec is stale and you should stop and say so. + +## What the probes already settled + +Answered live on 2026-09-07 (`server-spike.md` carries every request and response): + +- **P1 β€” HTTP schema apply is refused on every server configuration.** 409 `conflict`, "server-side schema apply is disabled for cluster-backed serving". `--cluster` is the server's only boot source, and "config-free" storage-root serving still reads a cluster state ledger from the bucket, so there is no server mode that accepts it. +- **P2 β€” there is no graph-creation endpoint.** `/graphs` is GET-only; an unknown graph gives 404 `not_found`. Graphs are created by `cluster apply`. +- **P12 β€” no hot reload.** `cluster apply` succeeds and the served schema never changes until the server restarts. Crucially, `cluster apply` never contacts the server: it needs the config directory and store credentials, **not** colocation. A cocoindex host can drive it. + +Still open and worth answering before the code they affect: **P3** (does the server serialize concurrent writes and branch merges β€” decides whether `exclusive_store()` disappears), **P13** (does `cluster apply` refuse while non-main branches exist), and P4-P11. + +Probe answers go in `server-spike.md` with the exact request and response. Never summarise a probe you have not run. The author has access to the Omnigraph source, so an answer may start from reading it β€” but it is not an answer until the binary confirms it. + +## Code conventions + +- **Directory structure:** + ``` + python/cocoindex/connectors/omnigraph/ + β”œβ”€β”€ __init__.py # re-exports _target.__all__ + β”œβ”€β”€ _client.py # ConnectionFactory, _HttpClient, errors, typed results (Epic B) + β”œβ”€β”€ _cluster.py # ClusterConfig, cluster apply, restart hook, converge (Epic D) + β”œβ”€β”€ _gq.py # .pg schema editing and GQ rendering β€” transport-agnostic, untouched + └── _target.py # reconciliation; imports only _client, _cluster and _gq + python/tests/connectors/ + β”œβ”€β”€ omnigraph_server_fixture.py # rustfs + omnigraph-server, both boot modes (Epic A) + └── test_omnigraph_target.py # the suite, migrated to the fixture (Epics C, E) + ``` + There is no `_transport.py`: with one data transport, a Protocol over a single implementation is an abstraction with nothing to abstract. Typed results and a typed error hierarchy live in `_client.py`. `_cluster.py` is the *only* module allowed to spawn a subprocess, and `_target.py` reaches it through the single `apply_schema()` seam (D1.1) so a future server-side apply (D3) is a one-place change. +- **Deletions are part of the work.** Epic C1.2 removes `canonical_store`, `_lock_dir`, `_temporary_text_file`, the `fcntl`/`msvcrt` lock helpers, `OmnigraphCliError`, and `_CliClient`. Don't leave them behind "just in case" β€” the guardrail greps for them. The one lock that survives is in `_cluster.py`, guarding the read-merge-write of the `.pg` file, and the one `subprocess` use is `cluster apply`. +- **Style:** follow `AGENTS.md` (`uv run ruff format .`, `uv run mypy`, external-module underscore rules). `aiohttp` appears only in `_client.py`; `_target.py` sees typed results and typed errors, never a `ClientResponse`. +- **Dependencies:** `aiohttp>=3.9` via the optional extra `cocoindex[omnigraph]`, and `msgspec` (already core) for JSON. Both are already in this repository, including the `aiohttp` mypy override. Anything beyond these two needs a human decision. +- **Testing:** live tests need `OMNIGRAPH_TEST_SERVER=1` and the binaries at `test/bin/`: `omnigraph` and `omnigraph-server` 0.10.0, plus `rustfs` 1.0.0-rc.5 for the object store. omnigraph reads standard AWS SDK env vars β€” `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_ENDPOINT_URL`, `AWS_REGION`, and `AWS_ALLOW_HTTP=true` for a plaintext rustfs endpoint. The whole test module must run with the sandbox disabled: `create_test_env` opens LMDB at import and gets EPERM inside it, and `omnigraph-server` cannot bind a port there either. Use `.venv/bin/python -m pytest python/tests/connectors/test_omnigraph_target.py`. With the sandbox off, `$TMPDIR` is the macOS default rather than the scratchpad, so pass absolute paths. Every step with a code deliverable ships a test; live behaviour claims cite the request and the response. +- **Commits:** short, human-looking messages, no co-author trailers (user rule). Prefix with the area, e.g. `omnigraph: rustfs + server fixture (A1.1)`. One step per commit where practical. + +## Engine facts to keep in mind + +Verified live on Omnigraph 0.10.0; re-verify before relying on them against a newer release. + +- `omnigraph-server` is **cluster-only at boot**: `--cluster` takes either a config directory (storage resolved through `cluster.yaml`) or a storage-root URI directly (`s3://`, `az://` β€” config-free serving). There is no `file://` server mode, which is why the connector can no longer offer a "point it at a directory" option. +- **Every** server answers `POST /graphs/{id}/schema/apply` with 409 `conflict` β€” storage-root boots included (P1). Schema changes go through `omnigraph cluster apply --config `, which refuses `--server`, and take effect only after a server restart (P12). `cluster.yaml` takes a `storage:` key naming the root, and `cluster import` seeds the ledger before the first `apply`. +- `GET /graphs` is forbidden unless a policy bundle opts in, even with `--unauthenticated`. Always configure `graph_id`. +- Errors are `{"error": str, "code": ErrorCode}` plus optional structured fields (`merge_conflicts`, `precondition_failure`, `resource_limit`, `key_conflict`). Body deserialization failures return 422 plain text. +- The server serves its own OpenAPI document at `/openapi.json` β€” useful as reference while writing the calls, though nothing is generated from it. `/healthz` returns `{"status","version","internal_schema_version"}`. +- Engine rules that held on the direct store and are expected to hold over HTTP, each with a probe to confirm it: no mixing upsert and delete in one mutation (P7), an 8,192-entity cap per mutation (P6), `cluster apply` may refuse while non-main branches exist, as direct-store `schema apply` did (P13), `branch merge` has no precondition (P11), destructive branch operations on a non-local scope needed CLI `--yes` consent (P8). + +## Workflow + +- **Autonomous:** any step whose deliverable is code plus tests β€” the Epic A fixture, all of Epics B and C, the chosen Epic D path, and Epic E β€” including running the live suite and recording guardrail outcomes. +- **Needs human input:** the D3 upstream ask (it is a ModernRelay product change, not a cocoindex one); the restart-hook API shape if D2 finds the spec's signature insufficient; anything that adds a row to the Decision Log or resolves an Open Question; dependency additions beyond `aiohttp` and `msgspec`; anything touching CI secrets; and any finding that makes a spec step wrong rather than merely harder. +- **The loop:** pick the next πŸ”² step β†’ implement the deliverable β†’ run tests and linters β†’ update the status table β†’ commit referencing the step. At a phase boundary, stop, verify each guardrail, fill in **Actual outcome**, and only then continue. +- **When blocked or ambiguous:** mark the step ⏸️ with a note, add an Open Question to the spec, and surface it rather than guessing. + +## Goals + +Near term: turn the manual rustfs + `omnigraph-server` setup from the P1 probe into the reusable fixture (A1.1-A1.3), then answer P3 and P13. MVP is reached when the migrated live suite passes over HTTP against a pinned 0.10.0 server in CI for both `managed_by="user"` and `managed_by="system"`, the connector's data path contains no subprocess or file-lock code, and the docs state plainly that schema changes require a server restart. diff --git a/dev/omnigraph-rework/SPEC.md b/dev/omnigraph-rework/SPEC.md new file mode 100644 index 000000000..adbb3c86d --- /dev/null +++ b/dev/omnigraph-rework/SPEC.md @@ -0,0 +1,569 @@ +# Omnigraph Connector Rework β€” Technical Spec + +**Author:** Roman Pronskiy Β· **Created:** 2026-09-06 + +> πŸ“„ **This is a living document.** Status markers, decisions, and guardrail outcomes are meant to be updated as the work happens. See [How to Update This Document](#how-to-update-this-document) before editing. + +### Changelog + +| Date | Change | Author | +|------|--------|--------| +| 2026-09-06 | Initial spec created from `dev/omnigraph-connector-transport-analysis.md` plus a live probe of `omnigraph-server` 0.10.0 | Roman Pronskiy | +| 2026-09-06 | SDK is ModernRelay's official Python client; renamed the package from `omnigraph-sdk` to `omnigraph` (decisions 17, 18) | Roman Pronskiy | +| 2026-09-06 | SDK repository starts at `pronskiy/omnigraph-python`, moving to the ModernRelay organisation in H1.5 (decisions 19, 20) | Roman Pronskiy | +| 2026-09-06 | Added step C2.7: claim the PyPI name with a functional `0.10.0a1` (decisions 21-23) | Roman Pronskiy | +| 2026-09-07 | **Rework.** The Python SDK is cut and the CLI transport is removed: HTTP is the connector's only transport (decisions 24-30). Old Epics C (SDK package), G (direct S3), and H (SDK polish) are cut; old Epic E (cluster control plane) is replaced by a spike-gated schema epic. Epics renumbered A-E | Roman Pronskiy || 2026-09-07 | Probes P1, P2, P12 answered live: HTTP schema apply is refused on every server configuration, there is no graph-creation endpoint, and no server hot-reloads. `managed_by="system"` is kept β€” every other connector supports it β€” through a `cluster apply` control plane behind an `apply_schema()` seam. Epic D rewritten; decisions 31-33 | Roman Pronskiy | + +### Status legend + +πŸ”² Not started Β· πŸ”„ In progress Β· βœ… Done Β· ⏸️ Blocked Β· ❌ Cut + +### Current focus + +**Now on:** Epic A β†’ Phase A2 β€” probes P1, P2, and P12 are answered (see `server-spike.md`); P3 through P11 and the new P13 remain. Phase A1's fixture exists only as a manual setup and still needs to become the reusable pytest fixture described in A1.1-A1.3. + +--- + +## 1. Executive summary + +The Omnigraph target connector on branch `omnigraph-connector` drives the `omnigraph` CLI against a direct `file://` store. That works for local and single-host use, but it cannot serve a deployed `omnigraph-server`: every operation is a subprocess, errors are text, locks are host-local, and workers need storage credentials plus a pinned 211 MB binary. + +The previous version of this spec answered that by building an official Python SDK and keeping the CLI alongside it. Both are now cut. A published SDK is one more repository, PyPI project, release cadence, and version lock to maintain, and its only consumer would be this connector. Keeping the CLI alongside HTTP means maintaining two transports, two sets of semantics, and two test modes forever. + +**The connector talks HTTP to `omnigraph-server`, and that is the only transport.** Roughly nine endpoints, `aiohttp` (already an optional dependency in this repository) and `msgspec` (already a core one), inside `python/cocoindex/connectors/omnigraph/_client.py`. No second repository, no PyPI project, no generated wire models, no drift check, no version-sync automation, and no binary on data-plane workers. + +Three costs are accepted and tracked here rather than discovered later: + +1. There is no "point it at a directory" mode any more. `omnigraph-server` is cluster-only at boot, so every user needs a running server. cocoindex documents how to start one and ships no process management (decision 27). +2. The existing live suite (~331 tests as of 2026-09-06) runs against a `file://` store through the CLI. All of it migrates to a server + rustfs fixture (Epic E1), and a representative slice migrates early, during Epic C, so server-side semantic differences surface while the client is still soft. +3. `managed_by="system"` β€” automatic type creation and evolution β€” cannot be done over HTTP at all. Probed live on 2026-09-07: `POST /schema/apply` returns 409 on **every** server configuration, because `--cluster` is the server's only boot source and the refusal is scoped to cluster-backed serving. Since all fourteen connectors that use `ManagedBy` default to `SYSTEM` and none refuses it, dropping the feature is not acceptable (decision 32). Epic D therefore keeps it through a control plane β€” edit the cluster config's `.pg`, `omnigraph cluster apply`, restart hook, poll until the served schema matches β€” behind an `apply_schema()` seam so a future server-side apply replaces it in one place. This is the one place the connector still needs the `omnigraph` binary, and only on the host that runs schema changes. + +Everything is open source; nothing here is commercial. + +--- + +## 2. Technical decisions + +| Area | Decision | Rationale | +|------|----------|-----------| +| Transport | HTTP to `omnigraph-server` is the connector's only transport. The CLI subprocess path is deleted, not kept alongside. | One code path, one set of semantics, one test mode. The connector is not on `main`, so there is no compatibility obligation. | +| No SDK | The HTTP client is private to the connector (`_client.py`), not a published package. | A public SDK is a repository, a PyPI project, a release cadence, and a version lock whose only consumer is this connector. | +| HTTP stack | `aiohttp>=3.9` behind the optional extra `cocoindex[omnigraph]`; JSON through `msgspec`. | `aiohttp` is already an optional extra here (doris) with a mypy override in place, and `msgspec` is already a core dependency. No new dependency enters the ecosystem. | +| Client lifecycle | `ConnectionFactory` lazily creates and caches one `aiohttp.ClientSession`; the environment lifespan closes it. | Matches the Neo4j convenience API; users provide connection details, not a session. | +| Local development | Documented, not managed. cocoindex ships no server-bootstrap API; the docs give a copy-paste `omnigraph-server` command and the test suite has a private fixture. | Process management and binary discovery are exactly what removing the CLI deleted; re-adding them as public API would undo the change. | +| Object store | rustfs is the S3-compatible backend for the test fixture and CI, in both server boot modes. | Standard for Omnigraph. It also means object-store use is verified by the default suite rather than by a deferred epic. | +| Schema | `managed_by="system"` is supported through a control plane: edit the type's block in the cluster config's `.pg`, run `omnigraph cluster apply`, restart via a user hook, poll `GET /schema` until it matches. It sits behind an `apply_schema()` seam so a future server-side apply replaces it in one place. | Answered live 2026-09-07 (P1, P12): HTTP `schema/apply` returns 409 on **every** server configuration β€” `--cluster` is the server's only boot source β€” and no server reloads schema without a restart. Every other cocoindex connector supports `SYSTEM` and defaults to it, so user-managed-only is not an acceptable end state (decision 32). | +| Graph provisioning | An unknown graph fails with a message naming the graph, the base URL, and the `cluster apply` procedure. Where a `ClusterConfig` is present the connector can create the graph itself, since `cluster apply` creates graphs (verified: `graph.probe create applied`). | P2: there is no graph-creation endpoint β€” `/graphs` is GET-only, unknown graph gives 404 `not_found`. | +| Atomicity | No host-local locks. Scratch branches carry the worker's identity and a UTC timestamp; the reaper deletes only branches past a TTL and never one this process created. `exclusive_store()` is removed once probe P3 confirms the server serializes writes and merges per graph. | `flock` coordinates one host; workers are many hosts. Server-side serialization is the only ordering that generalises. | +| Retries | Never auto-retry writes. Reads may retry with backoff behind an opt-in. Ambiguous write outcomes surface as errors; reconciliation re-reads on the next run. | A lost response after a durable commit must not duplicate work. | +| Errors | Mapped from HTTP status plus the `code` field and structured sub-objects (`merge_conflicts`, `precondition_failure`, `resource_limit`, `key_conflict`). Non-JSON bodies are tolerated verbatim. | The server returns 422 plain text for body deserialization errors (verified live). | +| Version coupling | A `/healthz` check at first use. An unexpected server minor logs one warning and proceeds. | The server routes new capabilities on new paths, so an old server returns 404 rather than a wrong result. Hard-failing on a minor bump would be stricter than the server's own contract. | +| Contract drift | The live suite against a pinned `omnigraph-server` 0.10.0 is the contract test. There are no generated models to drift. | The honest cost of hand-rolling: keep the surface at roughly nine calls so the suite can cover all of it. | +| Baseline | Omnigraph 0.10.0 only; no 0.9 compatibility. | Decided 2026-09-04; the live suite passes on 0.10.0 unchanged. | +| De-risking | The fixture and the probe list come first. No client code is written until every probe has a recorded answer. | Three of the eleven probes change the shape of what gets built. | + +--- + +## 3. Architecture overview + +``` + CocoIndex engine: declared target states, change detection + β”‚ actions per component + β–Ό + _target.py reconciliation Β· identity (coco_key) Β· ownership + scratch-branch atomicity Β· endpoint stubs + β”‚ β”‚ apply_schema() seam (D1.1) + β”‚ typed calls, typed errors β–Ό + β”‚ _cluster.py ClusterConfig(config_dir, restart) + β–Ό merge the type's block into .pg + _client.py ConnectionFactory(base_url, `omnigraph cluster apply --config` + graph_id, branch, token) restart hook β†’ poll GET /schema + _HttpClient β€” aiohttp, β”‚ + JSON via msgspec β”‚ writes cluster state + β”‚ POST /query Β· /mutate β”‚ (never contacts the server) + β”‚ /branches/* β”‚ + β”‚ GET /schema Β· /healthz β”‚ + β–Ό β–Ό + omnigraph-server ──────────────────── rustfs / S3 / cluster dir + β–² β”‚ + └───── restart, then serves the new schema (P12) +``` + +**The data plane is pure HTTP.** The only exception is schema: the server refuses `POST +/schema/apply` on every configuration (P1), so `managed_by="system"` goes around it through the +cluster config. `cluster apply` writes to the store and never contacts the server, so the control +plane needs the config directory and store credentials but not colocation with the server β€” only +the restart has to happen where the server runs, which is what the user-supplied hook is for. + +`_gq.py` (GQ rendering and `.pg` schema editing) is transport-agnostic today and is untouched by +this rework β€” the control plane reuses `merge_type_into_schema` and `remove_type_from_schema` +verbatim. A component's actions arrive at `_target.py`, which plans commits and asks the client +resolved from the target's `ContextKey` to execute them. + +--- + +## 4. Epics + +| Epic | Name | MVP | Depends on | +|------|------|-----|------------| +| A | Server + rustfs fixture and probe spike | Yes | β€” | +| B | HTTP client | Yes | A | +| C | Connector adaptation | Yes | B | +| D | Schema management (`managed_by="system"`) | Yes | A, B | +| E | Suite migration, docs, CI | Yes | B, C, D | + +B and C do not touch schema, so they proceed in parallel with D. + +--- + +### Epic A β€” Server + rustfs fixture and probe spike Β· MVP + +**Goal:** A reusable fixture that boots `omnigraph-server` 0.10.0 over rustfs in both boot modes, and a recorded, reproducible answer to every server-behaviour question the connector depends on β€” before any client code exists. Output is a findings document and decision rows, not shipped code (the fixture itself is shipped and reused by Epics C and E). +**Success metrics:** Every probe in A2 has a command and a response in `dev/omnigraph-rework/server-spike.md`; probes P1, P2, and P3 each have a decision row in Β§6. + +The author has access to the Omnigraph source, so answers may come from reading it, but each one is still confirmed against the running binary and recorded with the command and the response. + +#### Phase A1 β€” Fixture + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| A1.1 | rustfs fixture: start an S3-compatible endpoint, create a bucket, expose credentials and endpoint URL | πŸ”² | | +| A1.2 | Server fixture booting `omnigraph-server --unauthenticated` in **both** modes: `--cluster ` and `--cluster ` | πŸ”² | | +| A1.3 | Probe harness: a small raw-HTTP helper (`aiohttp`, no connector code) plus `restart()` / `stop()` on the fixture | πŸ”² | | + +**Steps (detail):** + +- **A1.1 β€” rustfs.** Deliverable: `python/tests/connectors/omnigraph_server_fixture.py` gains a fixture that starts rustfs, creates a bucket, and yields the endpoint URL plus credentials. Use a testcontainers-managed container following the postgres pattern in `AGENTS.md` β€” module-scoped sync fixture for the backend, function-scoped async fixture for per-test resources. If no usable image is available, fall back to the binary with a documented download step, and record which one in a decision row. +- **A1.2 β€” Server boot.** Deliverable: the same module boots `test/bin/omnigraph-server` on a free `127.0.0.1` port and polls `/healthz` until `{"status":"ok"}`. Two parametrisations, because they are different products for schema purposes: a `cluster.yaml` config directory (`version: 1`, `graphs..schema: `, applied with `omnigraph cluster apply --config`), and a bare storage-root URI pointing at the rustfs bucket. `omnigraph-server` cannot bind a port inside the Claude Code sandbox, so these tests run with the sandbox disabled. +- **A1.3 β€” Harness.** Deliverable: an async helper that issues raw requests against the booted server and returns status, headers, and body without interpretation, so probe answers record what the server actually said. `restart()` stops and re-boots with the same storage, for the probes that need it. + +**Exit guardrails β€” Phase A1 β†’ A2** + +| Guardrail | Criteria (pass/fail) | Status | Actual outcome | +|-----------|----------------------|--------|----------------| +| Both modes boot | A test boots the server over rustfs in cluster-dir mode and in storage-root mode and gets `{"status":"ok"}` from `/healthz` in each | πŸ”² | | +| Reproducible | The fixture runs twice in a row from a clean state with no manual step | πŸ”² | | + +#### Phase A2 β€” Probes + +Answers go to `dev/omnigraph-rework/server-spike.md`, each with the exact request and response. + +| Probe | Question | Decides | Status | Answer | +|-------|----------|---------|--------|--------| +| P1 | Does a server booted from a bare storage-root URI accept `POST /graphs/{id}/schema/apply`, or return 409 as the cluster-dir boot does? | **Epic D's path** | βœ… | **No β€” 409 `conflict`, "server-side schema apply is disabled for cluster-backed serving". `--cluster` is the server's only boot source, so this holds for every server. Epic D keeps `managed_by="system"` through the control plane instead.** | +| P2 | With no CLI, is there any HTTP path that creates a graph that does not yet exist? If not, what does the server answer for an unknown graph id? | **Graph provisioning; the uninitialized-store branch in `_target.py`** | βœ… | **No creation route β€” `/graphs` is GET-only (405 on POST). Unknown graph β†’ 404 `{"error":"graph 'X' not found","code":"not_found"}`. Provisioning is permanently an operator task.** | +| P12 | Does a running server pick up an applied revision without a restart? | **Whether the control plane needs a restart hook** | βœ… | **No β€” `cluster apply` succeeds, the served schema never changes; a restart picks it up. `cluster apply` never contacts the server, so the control plane needs the config dir and store credentials, not colocation.** | +| P13 | Does `cluster apply` refuse while non-main branches exist, as direct-store `schema apply` does? | D2.6 β€” control plane vs. scratch branches | πŸ”² | | +| P3 | Does the server serialize concurrent writes and concurrent branch merges per graph? | **Whether `exclusive_store()` can be removed outright** | πŸ”² | | +| P4 | Merge conflict over HTTP: status, `code`, and the `merge_conflicts` shape | Error hierarchy, failure matrix | πŸ”² | | +| P5 | Endpoint-not-found wording over HTTP β€” is it the same engine message the CLI printed (`(src\|dst) '…' not found in \w+`)? | Endpoint-stub recovery | πŸ”² | | +| P6 | Is the 8,192-entity mutation cap enforced server-side, and as 413 with `resource_limit` or as 400? | Chunking, failure matrix | πŸ”² | | +| P7 | Mixed upsert and delete in one mutation β€” same refusal as the direct store? | Scratch-branch flow | πŸ”² | | +| P8 | Does `branch delete` over HTTP need a consent equivalent of the CLI's `--yes` on a non-local store? | Branch teardown | πŸ”² | | +| P9 | Is `merge` with `delete_branch: true` reliable enough to drop the explicit delete? | Scratch-branch teardown | πŸ”² | | +| P10 | Bearer token setup and the minimal Cedar action set for query, mutate, branch create/merge/delete, schema get | Epic E2 least-privilege test | πŸ”² | | +| P11 | `mutate/if-graph-commit`: the 412 body, and whether a read's `graph_commit_id` is usable as the precondition | Optional single-commit fast path | πŸ”² | | + +**Exit guardrails β€” Epic A β†’ Epics B, C, D** + +| Guardrail | Criteria (pass/fail) | Status | Actual outcome | +|-----------|----------------------|--------|----------------| +| Probes answered | All eleven probes have a verified answer with command and response in `server-spike.md` | πŸ”² | | +| Path chosen | P1's answer is a decision row in Β§6 | βœ… | Decisions 31-33: HTTP apply impossible everywhere; SYSTEM supported via the control plane behind an `apply_schema()` seam | +| Provisioning settled | P2's answer is a decision row stating how a missing graph is created or reported | βœ… | Decision 31: no creation endpoint; created via `cluster apply` when a `ClusterConfig` is present, else reported | +| Locking settled | P3's answer is a decision row stating whether `exclusive_store()` is removed or replaced | πŸ”² | | + +--- + +### Epic B β€” HTTP client Β· MVP + +**Goal:** `_client.py` is an `aiohttp`-based client for the roughly nine endpoints the connector needs, with typed results and a typed error hierarchy, and no subprocess, temp file, or file lock anywhere in it. +**Success metrics:** Every method has a mocked-transport unit test for its success path and every error class; a live smoke test passes against the Epic A fixture; `uv run mypy` and `uv run ruff format --check .` clean. + +#### Phase B1 β€” Factory, session, packaging + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| B1.1 | `ConnectionFactory` with `base_url`, `graph_id`, `branch`, `token`, `timeout`; lazy cached `aiohttp.ClientSession`; `aclose()` from the lifespan | πŸ”² | | +| B1.2 | Token handling: `OMNIGRAPH_TOKEN` fallback, never in `repr`, logs, or exception text; TLS verification on with no public switch | πŸ”² | | +| B1.3 | Optional extra `cocoindex[omnigraph] = ["aiohttp>=3.9.0"]`, `ci-enabled-optional-deps` entry, import guard with a clear message | πŸ”² | | +| B1.4 | `/healthz` version check at first use; one warning on an unexpected server minor | πŸ”² | | + +**Steps (detail):** + +- **B1.1 β€” Factory.** Deliverable in `_client.py`: + ```python + @dataclasses.dataclass(frozen=True) + class ConnectionFactory: + base_url: str + graph_id: str + branch: str = "main" + token: str | None = None + timeout: float = 30.0 + + def client(self) -> _HttpClient: ... # creates the session once, caches it + async def aclose(self) -> None: ... + ``` + The factory is provided through the same `ContextKey` mechanism as today and resolved at action time, never captured at declare time. The old `store` and `cli` fields are deleted, not aliased. +- **B1.3 β€” Packaging.** Deliverable: `pyproject.toml` extra plus the mypy override (`aiohttp`, `aiohttp.*` are already listed). The import guard raises a message naming the extra when `aiohttp` is missing. + +#### Phase B2 β€” Calls and typed results + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| B2.1 | `read_schema()`, `query()`, `mutate()` with typed results | πŸ”² | | +| B2.2 | `branch_create/list/merge/delete` | πŸ”² | | +| B2.3 | ~~`apply_schema()` on the HTTP client~~ | ❌ | Cut: P1 proved the endpoint is refused on every server. Schema writes live in `_cluster.py` (Epic D); revisit only if D3 lands | +| B2.4 | `mutate_if_graph_commit()` β€” only if P11 says it is usable | πŸ”² | Optional | + +**Steps (detail):** + +- **B2.1 β€” Reads and writes.** Deliverable: + ```python + class MutationResult(NamedTuple): + affected_nodes: int + affected_edges: int + commit_id: str | None + + class MergeOutcome(enum.Enum): + ALREADY_UP_TO_DATE = "already_up_to_date" + FAST_FORWARD = "fast_forward" + MERGED = "merged" + + class MergeResult(NamedTuple): + outcome: MergeOutcome + source_deleted: bool + ``` + `query()` returns rows untouched β€” no key-case conversion is applied to user data. `Query` (positional `$?` slots, from `_gq.py`) is rendered into the request body with its params; the temp-file plumbing the CLI needed is gone. + +#### Phase B3 β€” Errors and retries + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| B3.1 | `OmnigraphError` hierarchy keyed on status, `code`, and structured fields; 422 plain text tolerated | πŸ”² | | +| B3.2 | Classification methods: `uninitialized_graph()`, `missing_endpoint()`, `blocked_by_non_main_branches()` | πŸ”² | Wording from P2, P5 | +| B3.3 | Retry policy: none on writes; opt-in backoff on reads for 429, 503, and network errors | πŸ”² | | + +**Steps (detail):** + +- **B3.1 β€” Errors.** Deliverable: `OmnigraphError` β†’ `HttpError(status, code, message, request_id)` β†’ `BadRequestError` (400, 422), `UnauthorizedError` (401), `ForbiddenError` (403), `NotFoundError` (404), `ConflictError` (409, carrying `merge_conflicts` and `key_conflict`), `PreconditionFailedError` (412), `PayloadTooLargeError` (413, carrying `resource_limit`), `RateLimitedError` (429), `ServerUnavailableError` (503), plus `NetworkError` and `ConfigurationError`. A body that is not the JSON envelope becomes the message verbatim. +- **B3.2 β€” Classification.** Deliverable: the three predicates `_target.py` needs, as methods on the error rather than regexes in the reconciliation code. `missing_endpoint()` uses the wording confirmed in P5 and lives in exactly one place. + +**Exit guardrails β€” Epic B β†’ Epic C** + +| Guardrail | Criteria (pass/fail) | Status | Actual outcome | +|-----------|----------------------|--------|----------------| +| Mocked coverage | Every B2 method's success path and every B3.1 error class has a unit test against a mocked transport | πŸ”² | | +| Secrets | A test asserts the token is absent from `repr(factory)`, `repr(client)`, and every raised exception's `str()` | πŸ”² | | +| Live smoke | `read_schema`, `query`, `mutate`, and the four branch calls pass against the Epic A fixture | πŸ”² | | +| No CLI residue | `_client.py` contains no `subprocess`, `tempfile`, `fcntl`, or `msvcrt` import | πŸ”² | | + +--- + +### Epic C β€” Connector adaptation Β· MVP + +**Goal:** `_target.py` reconciles over HTTP with the same data semantics it has today, with no host-local locking and no CLI error text. +**Success metrics:** No occurrence of `_CliClient`, `OmnigraphCliError`, or `_ENDPOINT_NOT_FOUND_RE` anywhere in the connector; the migrated slice of the live suite (C1.4) passes against the Epic A fixture. + +#### Phase C1 β€” Swap the client + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| C1.1 | `_apply_entity_actions`, `_apply_type_actions`, `_scratch_branch`, `_reap_abandoned_scratch_branches`, `_keep_referenced_nodes`, `_mutate_with_endpoint_retry` take `_HttpClient` | πŸ”² | 26 CLI references today | +| C1.2 | Delete the CLI machinery: `canonical_store`, `_lock_dir`, `_temporary_text_file`, the flock/msvcrt helpers, `OmnigraphCliError`, `_CliClient` | πŸ”² | | +| C1.3 | Endpoint-stub recovery and the uninitialized-graph branch use B3.2's classification | πŸ”² | | +| C1.4 | Migrate a representative slice of the live suite now, not at the end: entity lifecycle, mixed upsert/delete, endpoint stubs, app drop | πŸ”² | Surfaces server-semantic differences while the client is still soft | + +#### Phase C2 β€” Atomicity without host locks + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| C2.1 | Scratch branch names `coco_scratch___` | πŸ”² | | +| C2.2 | Reaper deletes only branches past a TTL (default 1 h) and never one this process created | πŸ”² | | +| C2.3 | Remove or replace `exclusive_store()` per probe P3 | πŸ”² | Blocked on P3 | +| C2.4 | Cancellation before send, during upload, and while awaiting a response leaves no scratch branch | πŸ”² | | + +#### Phase C3 β€” Failure matrix + +| Condition | Signal | Connector behaviour | +|---|---|---| +| 401 / 403 | `UnauthorizedError` / `ForbiddenError` | Fail the action immediately, naming the operation and graph; no retry | +| 404 graph | `NotFoundError` | Fail with "graph `` is not served by ``"; if P2 gave a provisioning path, take it instead | +| Endpoint missing | classified engine message | Existing endpoint-stub recovery on the scratch path | +| 409 merge conflict | `ConflictError.merge_conflicts` | Fail the component sync, delete the scratch branch, list conflicts in the error | +| 409 other | `ConflictError` | Fail; the next run re-reads state | +| 413 | `PayloadTooLargeError.resource_limit` | Fail with the limit and the actual; the connector already chunks at 8,192 | +| 422 | `BadRequestError` with raw text | Fail; indicates a client/server mismatch | +| 429 / 503 on a read | `RateLimitedError` / `ServerUnavailableError` | Retry with backoff up to the policy limit | +| 429 / 503 on a write | same | Fail; no retry | +| Timeout after a write was sent | `NetworkError` | Fail as ambiguous; tracking is not advanced, so the next run reconciles | + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| C3.1 | One test per row of the matrix above | πŸ”² | | + +**Exit guardrails β€” Epic C β†’ Epic E** + +| Guardrail | Criteria (pass/fail) | Status | Actual outcome | +|-----------|----------------------|--------|----------------| +| No transport leak | Grepping the connector for `_CliClient`, `OmnigraphCliError`, `_ENDPOINT_NOT_FOUND_RE`, and `subprocess` finds nothing | πŸ”² | | +| Slice passes | The C1.4 tests pass against the Epic A fixture | πŸ”² | | +| No leaked branches | After every HTTP-mode test, the graph's branch list contains only `main` | πŸ”² | | +| Matrix covered | One passing test per C3 row | πŸ”² | | + +--- + +### Epic D β€” Schema management (`managed_by="system"`) Β· MVP + +**Context β€” every other connector supports `SYSTEM`, and so must this one.** + +`ManagedBy` is a shared enum (`connectorkits/target.py`: `SYSTEM` / `USER`). Fourteen connectors +use it, **every one defaults to `ManagedBy.SYSTEM`, and none refuses it** β€” including this +connector today, in six public signatures. + +| Connector | Reaches store via | How it applies schema under `SYSTEM` | +|---|---|---| +| postgres | asyncpg wire | `CREATE TABLE`, `ALTER TABLE`, `CREATE INDEX` | +| neo4j | bolt | `CREATE CONSTRAINT`, index creation | +| bigquery | HTTP API | `CREATE TABLE`, `ALTER TABLE`, `DROP TABLE` | +| qdrant | HTTP/gRPC client | `_create_collection` | +| snowflake | driver | `CREATE DATABASE`, `CREATE TABLE`, `ALTER TABLE` | +| doris | aiohttp + MySQL wire | `CREATE TABLE`, `ALTER TABLE` | +| surrealdb | client | `DEFINE TABLE` | +| falkordb | redis protocol | `CREATE INDEX` | +| turbopuffer | HTTP API | namespace implicit on first write; schema rides in the write payload | +| valkey, zvec | protocol / client | schemaless β€” no DDL needed | +| lancedb, sqlite | embedded / local | `create_table` | +| **omnigraph over HTTP** | **omnigraph-server** | **refused: 409, unconditionally (P1)** | + +The structural pattern is that **schema travels the same channel as the data**. DDL is just another +authorized operation on the wire, including for the three that are pure remote HTTP APIs +(bigquery, qdrant, turbopuffer). `omnigraph-server` is the only store here that deliberately closes +that channel: schema is cluster configuration, applied out of band and picked up on restart. + +Dropping the SDK did not cause this β€” dropping **direct store access** did. The published-SDK plan +would have hit the identical 409, which is why the original spec invented a control plane. The +real fork was always *server versus direct store*, and the CLI was silently carrying schema +management. Shipping `managed_by="user"`-only would make omnigraph the single connector whose +default value does not work, so it is rejected as an end state (decision 32). + +**Goal:** `managed_by="system"` works against a server-backed graph: the connector owns its type's +block in the graph's `.pg` inside the cluster config, applies it with `omnigraph cluster apply`, +triggers a restart through a user-supplied hook, and proceeds only once the served schema matches. +The operation sits behind one seam so a future server-side apply (D3) is a drop-in replacement. +**Success metrics:** the type lifecycle tests (create, property add, encoder change, release to +user, removal) pass in HTTP mode using the fixture's restart hook; a factory without a restart hook +fails with an actionable message and converges on re-run after a manual restart. + +What P1, P2, and P12 established, and what the design must therefore accept: HTTP schema apply is +refused on every server configuration; there is no graph-creation endpoint; and a running server +never reloads schema without a restart. But `cluster apply` needs only the CLI binary, the config +directory, and object-store credentials β€” **not colocation with the server** β€” so the control plane +can run from a cocoindex host. + +#### Phase D1 β€” The seam and the config model + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| D1.1 | `apply_schema(schema_pg)` as the one schema-write operation `_target.py` calls, with a pluggable implementation | πŸ”² | Makes D3 a drop-in | +| D1.2 | `ClusterConfig` dataclass; validate `cluster.yaml`, including a `storage:` root | πŸ”² | | +| D1.3 | Locate or create the graph's `.pg` from `graphs..schema` | πŸ”² | The HTTP-mode equivalent of `init` | +| D1.4 | Reconcile-time refusal of `managed_by="system"` without a `ClusterConfig`, naming the target | πŸ”² | Fail early, not mid-sync | + +**Steps (detail):** + +- **D1.2 β€” Config.** Deliverable: + ```python + @dataclasses.dataclass(frozen=True) + class ClusterConfig: + config_dir: pathlib.Path # contains cluster.yaml + cli: str = "omnigraph" # runs `cluster apply` + restart: Callable[[], Awaitable[None]] | None = None + settle_timeout: float = 120.0 # wait for GET /schema to match + ``` + The CLI is required only on the host running the control plane, and object-store credentials only + when the cluster's `storage:` is an object-store root. Data-plane workers need neither. + +#### Phase D2 β€” Control plane + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| D2.1 | Read-merge-write of the type's block with `_gq.merge_type_into_schema` / `remove_type_from_schema`, under a lock keyed on the config directory | πŸ”² | `_gq.py` unchanged | +| D2.2 | Run `omnigraph cluster apply --config --json`; parse `ok`, `changes`, `errors` | πŸ”² | Verified shape: `changes[].{resource, operation, disposition}` | +| D2.3 | Invoke the restart hook; with none, raise `SchemaRestartRequired` naming the graph and config directory | πŸ”² | | +| D2.4 | Poll `GET /schema` until `schema_source` matches, else fail after `settle_timeout` | πŸ”² | Exponential from 0.5 s | +| D2.5 | Removal semantics: soft drops only; graph deletion never automated | πŸ”² | | +| D2.6 | Scratch-branch interaction: confirm whether `cluster apply` refuses while non-main branches exist | πŸ”² | New probe P13 | +| D2.7 | Ownership marker handling (`coco_managed_by_`) unchanged across the transport swap | πŸ”² | | + +**Steps (detail):** + +- **D2.3 β€” Restart.** Deliverable: the hook is awaited once per apply; exceptions propagate and fail + the sync. `SchemaRestartRequired` is raised after a *successful* apply when no hook is configured, + so a re-run after the operator's restart finds the served schema already matching and continues. +- **D2.4 β€” Converge.** Deliverable: on timeout, an error carrying the digest of the expected and the + served source. P12 confirmed the served schema never changes on its own, so this poll is strictly + waiting on the restart, not on eventual consistency. + +#### Phase D3 β€” Native server-side apply Β· Post-MVP Β· blocked on upstream + +The only path that removes both the CLI dependency and the restart, and the only one that puts +omnigraph on equal footing with the other thirteen connectors. + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| D3.1 | Upstream ask to ModernRelay: a per-graph opt-in that permits `POST /schema/apply` and reloads that graph in place | πŸ”² | The 409 is a deliberate guard, not an oversight | +| D3.2 | When it lands: a second implementation behind the D1.1 seam; `ClusterConfig` becomes optional | πŸ”² | | +| D3.3 | Retire the restart hook wherever the server supports in-place reload | πŸ”² | | + +**Exit guardrails β€” Epic D β†’ Epic E** + +| Guardrail | Criteria (pass/fail) | Status | Actual outcome | +|-----------|----------------------|--------|----------------| +| Lifecycle | Type create, property add, encoder change, release to user, and removal pass in HTTP mode with the fixture's restart hook | πŸ”² | | +| No hook | With `restart=None`, the first sync raises `SchemaRestartRequired`; after the fixture restarts the server, the second sync converges | πŸ”² | | +| Seam honoured | `_target.py` calls `apply_schema()` only; no `subprocess` or `cluster apply` string appears outside the control-plane module | πŸ”² | | +| Data isolation | No data-plane test invokes `cluster apply` | πŸ”² | | + +--- + +### Epic E β€” Suite migration, docs, CI Β· MVP + +**Goal:** The whole live suite runs over HTTP against the Epic A fixture, the docs describe the server-only connection model honestly, and CI runs it. +**Success metrics:** `test_omnigraph_target.py` passes in full against the fixture with a recorded count; `omnigraph.mdx` no longer mentions a `store` URI or a CLI. + +#### Phase E1 β€” Suite migration + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| E1.1 | Re-point the remaining tests from the `file://` CLI fixture to the server fixture | πŸ”² | ~331 tests as of 2026-09-06 | +| E1.2 | Classify every failure as "server semantics differ" or "fixture gap"; the first kind gets a decision row | πŸ”² | | +| E1.3 | Record the final passing count in this spec | πŸ”² | | + +#### Phase E2 β€” Live matrix + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| E2.1 | Unauthenticated server | πŸ”² | | +| E2.2 | Bearer-authenticated server with a least-privilege policy bundle (P10) | πŸ”² | | +| E2.3 | Two graph ids on one server; wrong graph id | πŸ”² | | +| E2.4 | Missing permission for each of query, mutate, branch create, merge, delete | πŸ”² | | +| E2.5 | Writes above one request's entity limit | πŸ”² | | +| E2.6 | Merge conflict from a concurrent write on `main` | πŸ”² | | +| E2.7 | Cancellation at three points | πŸ”² | | +| E2.8 | Connection loss after a successful write (transport fault injection) | πŸ”² | | +| E2.9 | Two processes writing to one graph concurrently | πŸ”² | | +| E2.10 | Server minor-version mismatch produces one warning and proceeds | πŸ”² | | + +#### Phase E3 β€” Documentation + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| E3.1 | Rewrite `omnigraph.mdx` around `base_url` / `graph_id`; delete the `store` URI model | πŸ”² | | +| E3.2 | Local development section: copy-paste `omnigraph-server --cluster ./og --unauthenticated`, and the rustfs variant the tests use | πŸ”² | | +| E3.3 | Schema section per Epic D's path β€” automatic, or the operator runbook | πŸ”² | | +| E3.4 | Failure and retry behaviour table; security notes (token sourcing, least privilege, no `/graphs` enumeration) | πŸ”² | | + +#### Phase E4 β€” CI + +| Step | Description | Status | Notes | +|------|-------------|--------|-------| +| E4.1 | Install `omnigraph-server` 0.10.0 and rustfs in `.github/workflows/_test.yml`, cached by version | πŸ”² | 211 MB per platform | +| E4.2 | Gate live tests on `OMNIGRAPH_TEST_SERVER=1` | πŸ”² | `OMNIGRAPH_TEST_STORE` retires with the CLI | +| E4.3 | `aiohttp` in `ci-enabled-optional-deps` | πŸ”² | | + +**Exit guardrails β€” Epic E β†’ MVP done** + +| Guardrail | Criteria (pass/fail) | Status | Actual outcome | +|-----------|----------------------|--------|----------------| +| Suite green | The full live suite passes in CI on Linux x86_64 and macOS arm64 | πŸ”² | | +| Pre-submission | `prek run --all-files` passes | πŸ”² | | +| Docs honest | `omnigraph.mdx` states that a running `omnigraph-server` is required and what schema management the connector does | πŸ”² | | + +--- + +## 5. Risk register + +| Risk | Likelihood | Impact | Mitigation | +|------|------------|--------|------------| +| ~~Probe P1 says no~~ β€” **fired 2026-09-07.** HTTP schema apply is refused on every server configuration | β€” | High | Resolved by Epic D's control plane rather than by dropping the feature: every other connector supports `SYSTEM` and defaults to it (decision 32) | +| The restart hook is operationally awkward β€” rolling deploys, managed hosting, multi-replica clusters | High | Med | `SchemaRestartRequired` converges on re-run after a manual restart; `managed_by="user"` documented as the zero-restart option; D3 removes it entirely if upstream lands | +| D3 never lands upstream, so the CLI dependency for schema is permanent | Med | Med | The control plane is MVP and ships without D3. The binary is needed only on the host running schema changes, and P12 showed that host need not be the server's | +| No HTTP graph-creation path (P2), so a pure-HTTP deployment cannot bootstrap its own graph | Med | Med | With a `ClusterConfig` the connector creates it via `cluster apply` (verified: `graph.probe create applied`). Without one, fail naming the graph, the base URL, and the procedure; document it in the quickstart ahead of the connector code | +| Suite migration surfaces server semantics that differ from the direct store, forcing reconciliation rework late | Med | High | C1.4 migrates a representative slice during Epic C rather than at the end | +| Host locks are gone and the server does not serialize branch merges as assumed β†’ lost updates between workers | Med | High | Probe P3 before C2.3; the two-process concurrent-writer test in E2.9 | +| The hand-rolled client drifts from the server contract, with no generated models and no drift check | Med | Med | The honest cost of dropping the SDK. Keep the surface at roughly nine calls; the live suite against a pinned server is the contract test | +| Losing the local `file://` mode hurts the five-minute story | High | Med | Accepted (decision 27); docs snippet in E3.2 | +| `omnigraph-server` is a 211 MB asset per platform, plus rustfs, in CI | High | Low | Cache the download by version; run live tests on one Linux runner plus macOS arm64 | +| Server minor release changes the wire contract mid-project | Med | Med | Pin 0.10.0 in fixtures and CI; `/healthz` check at first use | +| Ambiguous write after a timeout leads to duplicate or missing data | Low | High | No write retries; tracking advances only on a confirmed response; E2.8 fault injection | +| `GET /graphs` enumeration is forbidden by default, so a misconfigured graph id fails late | Med | Low | Validate with `GET /graphs/{id}/schema` at first use; the error names the graph and the base URL | + +--- + +## 6. Decision log + +Append-only. Rows 1-23 are kept for history; rows 24-30 supersede the ones they name. + +| # | Date | Decision | Context | Decided by | +|---|------|----------|---------|------------| +| 1 | 2026-09-04 | Target Omnigraph 0.10.0 only; no 0.9 compatibility | Live suite passes on 0.10.0 unchanged; every 0.9 constraint the connector works around still holds | Roman Pronskiy | +| 2 | 2026-09-06 | Keep the CLI transport as the local/direct-store mode | It needs no server, no extra package, and is fully verified | Roman Pronskiy | +| 3 | 2026-09-06 | Build a public Python SDK (`omnigraph`) rather than a connector-private client | Scope B chosen over Scope A in the interview | Roman Pronskiy | +| 4 | 2026-09-06 | SDK lives in a separate repository, not as a workspace member | Clean ownership and release cadence | Roman Pronskiy | +| 5 | 2026-09-06 | PyPI name `omnigraph-sdk`, import `omnigraph_sdk` | Free on PyPI as of 2026-09-06; mirrors the TypeScript package's positioning without squatting the product name | Roman Pronskiy | +| 6 | 2026-09-06 | The connection factory owns the client; the lifespan closes it | Neo4j parity; users provide connection details, not a client | Roman Pronskiy | +| 7 | 2026-09-06 | `managed_by="system"` in cluster mode via a control plane (`ClusterConfig`, `cluster apply`, restart hook, poll) | Verified live: HTTP schema apply β†’ 409, `cluster apply` refuses `--server`, no hot reload, restart required | Roman Pronskiy | +| 8 | 2026-09-06 | Without a `ClusterConfig`, HTTP targets are `managed_by="user"` only and refuse `system` at reconcile time | Fail early with a clear message rather than mid-sync | Roman Pronskiy | +| 9 | 2026-09-06 | Run a CLI `--server --graph` spike before SDK work; not a supported user mode | Cheapest way to verify server semantics | Roman Pronskiy | +| 10 | 2026-09-06 | Separate factory classes; rename `ConnectionFactory` β†’ `CliConnectionFactory` with no alias | Connector not on `main`; no compatibility obligation | Roman Pronskiy | +| 11 | 2026-09-06 | Generate wire models from the server's served `/openapi.json` with `datamodel-code-generator` β†’ `msgspec.Struct`; handwritten facade | Server publishes its spec; avoids a literal TypeScript port | Roman Pronskiy | +| 12 | 2026-09-06 | No automatic retries on writes; opt-in backoff on reads | Lost responses after durable commits must not be replayed | Roman Pronskiy | +| 13 | 2026-09-06 | SDK minor version tracks server minor version | Same rule as the TypeScript SDK; server fails closed on new paths | Roman Pronskiy | +| 14 | 2026-09-06 | Direct S3 stays unverified until Epic G | Host-local locks and `--yes` consent untested against non-local stores | Roman Pronskiy | +| 15 | 2026-09-06 | The connector never calls `GET /graphs` | Forbidden by default even on an unauthenticated server (verified live) | Roman Pronskiy | +| 16 | 2026-09-06 | Everything is open source; no private spec | No commercial plans in scope | Roman Pronskiy | +| 17 | 2026-09-06 | `omnigraph-sdk` is ModernRelay's official Python SDK, under the ModernRelay organisation | The author is a ModernRelay employee | Roman Pronskiy | +| 18 | 2026-09-06 | PyPI name `omnigraph`, import `omnigraph`, client class `omnigraph.Omnigraph`; supersedes 5 and 17 | Mirrors the npm package `@modernrelay/omnigraph` | Roman Pronskiy | +| 19 | 2026-09-06 | The SDK repository starts at `github.com/pronskiy/omnigraph-python`; supersedes 17 | Development can start without org provisioning | Roman Pronskiy | +| 20 | 2026-09-06 | Step H1.5 is reinstated as the org transfer rather than a handover offer | Decision 19 defers the move | Roman Pronskiy | +| 21 | 2026-09-06 | Claim the PyPI name with a functional `0.10.0a1` at the end of Phase C2 | A pending publisher reserves nothing; an empty package is squatting under PEP 541 | Roman Pronskiy | +| 22 | 2026-09-06 | Trusted publishing from the start rather than a long-lived API token | A long-lived token is a standing secret | Roman Pronskiy | +| 23 | 2026-09-06 | The PyPI project is created personally and moves to a ModernRelay organisation in H1.5 | Follows decision 19 | Roman Pronskiy | +| 24 | 2026-09-07 | **HTTP is the connector's only transport; the CLI transport is removed rather than kept alongside.** Supersedes 2, 9, 10 | Two transports means two sets of semantics, two test modes, and subprocess plumbing maintained forever. The connector is not on `main`, so nothing is owed to existing configurations | Roman Pronskiy | +| 25 | 2026-09-07 | **No separate Python SDK package.** The HTTP client is private to the connector in `_client.py`. Supersedes 3, 4, 5, 11, 17, 18, 19, 20, 21, 22, 23; cuts old Epics C and H | A published SDK is a repository, a PyPI project, a release cadence, and a version lock whose only consumer is this connector | Roman Pronskiy | +| 26 | 2026-09-07 | `aiohttp` behind the optional extra `cocoindex[omnigraph]`; JSON through `msgspec` | Both are already present in this repository β€” `aiohttp` as the doris extra with a mypy override, `msgspec` as a core dependency. No new dependency enters the ecosystem | Roman Pronskiy | +| 27 | 2026-09-07 | Local development is documented, not managed: cocoindex ships no server-bootstrap API | Process management and binary discovery are what removing the CLI deleted; re-adding them as public API would undo the change | Roman Pronskiy | +| 28 | 2026-09-07 | rustfs is the S3-compatible backend for fixtures and CI, so object-store use is verified by the default suite. **Old Epic G is cut.** Supersedes 14 | rustfs is the standard local backend for Omnigraph, and with the CLI gone there is no direct-store mode left to verify separately | Roman Pronskiy | +| 29 | 2026-09-07 | `managed_by="system"` support is decided by probe P1; the fallback is user-managed-only with an actionable diff. Supersedes 7, 8; replaces the old Epic E with Epic D | The 409 was measured on a `cluster.yaml`-backed server; a config-free storage-root boot is untested and may behave differently | Roman Pronskiy | +| 30 | 2026-09-07 | Version coupling becomes a `/healthz` check at first use; an unexpected minor logs one warning and proceeds. Supersedes 13 | With no published package there is no version to lock. The server routes new capabilities on new paths, so it fails closed on its own | Roman Pronskiy | +| 31 | 2026-09-07 | **P1, P2, and P12 answered live.** HTTP `POST /schema/apply` returns 409 on every server configuration; there is no graph-creation endpoint; a running server never reloads schema without a restart. Supersedes the branch in 29 | `--cluster` is the server's only boot source (RFC-011), and the 409 is scoped to "cluster-backed serving", so it applies to every served graph. Evidence in `server-spike.md` | Roman Pronskiy | +| 32 | 2026-09-07 | **`managed_by="system"` must be supported.** User-managed-only is rejected as an end state | `ManagedBy` is a shared enum used by 14 connectors; every one defaults to `SYSTEM` and none refuses it. Shipping otherwise would make omnigraph the only connector whose default value does not work | Roman Pronskiy | +| 33 | 2026-09-07 | Support it through a `cluster apply` control plane behind an `apply_schema()` seam, and pursue a server-side apply upstream as the eventual replacement (Epic D3) | The control plane works today with no upstream dependency, and P12 showed `cluster apply` needs the config directory and store credentials but not colocation with the server. The seam keeps the eventual native path a one-place change | Roman Pronskiy | + +--- + +## 7. Open questions + +The eleven Epic A probes (P1-P11) are the live open questions; they are tracked in the A2 table rather than duplicated here. Beyond them: + +- [ ] How should an ambiguous write be surfaced in `App.update()` reporting beyond a failed action? +- [ ] Should `load/ndjson` replace generated GQ upserts for large batches, and can it preserve delete semantics? Deferred past MVP unless E2.5 shows a need. +- [ ] Does the migrated suite need a slower CI lane, given every test now needs a booted server and rustfs? +- [x] ~~SDK repository name, PyPI account ownership, and the org transfer.~~ Moot as of 2026-09-07: no SDK (decision 25). +- [x] ~~Does ModernRelay plan an official Python SDK?~~ Moot; the connector no longer needs one. +- [x] ~~Is the served `/openapi.json` equivalent to the TypeScript repository's copy?~~ Moot; no models are generated. The served document remains useful as reference while writing the nine calls. + +--- + +## How to Update This Document + +This spec is the source of truth for the build. Keep it current as work happens: + +- **Status markers.** Update a step's status in its tracker table as you go: πŸ”² β†’ πŸ”„ β†’ βœ…. Use ⏸️ for blocked (note why in Notes) and ❌ for cut (leave the row; the strikethrough of history is useful). +- **Current focus.** Keep the pointer at the top aimed at the next actionable πŸ”² step. Update it the moment you finish a step or cross a phase boundary β€” a stale pointer is worse than none, since it sends the next reader to the wrong place. +- **Guardrails.** When you hit a phase boundary, fill the **Actual outcome** column with what really happened and set the guardrail status. Don't advance to the next phase until its entry guardrails pass β€” or log a decision explaining why you're proceeding anyway. +- **Probes.** Epic A's probe answers go in `dev/omnigraph-rework/server-spike.md` with the exact request and response, and the one-line answer goes in the A2 table. Never summarise a probe you have not run. +- **Decisions.** Any non-trivial choice made during the build gets a new row in the Decision Log (Β§6). It's append-only β€” reversals are new rows naming what they supersede, not edits. If the choice changes the architecture, also update the Technical Decisions snapshot (Β§2). +- **Spec changes.** Structural changes (new epic, re-scoped phase) get a Changelog row at the top. Keep the executive summary honest if the project's shape shifts. +- **Open questions.** When one resolves, strike it from Β§7 and log the decision in Β§6. diff --git a/dev/omnigraph-rework/server-spike.md b/dev/omnigraph-rework/server-spike.md new file mode 100644 index 000000000..370ba4da0 --- /dev/null +++ b/dev/omnigraph-rework/server-spike.md @@ -0,0 +1,174 @@ +# Omnigraph server spike β€” probe findings + +Answers to the Epic A probes in `SPEC.md` Β§4. Every answer records the exact request and the +exact response. An answer sourced from reading Omnigraph source is not an answer until the +binary confirms it. + +## Environment + +- `omnigraph-server` / `omnigraph` 0.10.0 from `test/bin/` (macOS arm64). +- rustfs `1.0.0-rc.5` (`rustfs-macos-aarch64-v1.0.0-rc.5.zip`, sha256 verified against the + release `SHA256SUMS`), at `test/bin/rustfs`. No Homebrew formula exists; all rustfs releases + are prereleases, so `releases/latest` returns 404 β€” fetch a specific tag. +- rustfs: `rustfs server --address 127.0.0.1:9100 --access-key cocotest --secret-key cocotestsecret ./scratch/rustfs-data` +- Bucket `omnigraph` created with boto3, path-style addressing. +- omnigraph reads standard AWS SDK / `object_store` env vars. Working set: + `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_ENDPOINT_URL`, `AWS_REGION`, `AWS_ALLOW_HTTP=true`. +- Both servers must run with the Claude Code sandbox disabled β€” neither can bind a port inside it. + +## P1 β€” Does a storage-root-booted server accept HTTP schema apply? + +**Answer: no. 409, identical refusal to the cluster-dir boot. There is no server mode that accepts it.** + +First, a finding that reframes the question: **"config-free serving" is not config-free.** Booting +against a bare storage root with no prior cluster state fails outright: + +``` +$ omnigraph-server --cluster s3://omnigraph/probe --unauthenticated --bind 127.0.0.1:9110 +Error: + 0: the cluster at 's3://omnigraph/probe' is not ready to serve: + [cluster_state_missing] __cluster/state.json: no cluster state ledger; + run `cluster import` and `cluster apply` first +``` + +The storage-root boot is not an alternative to cluster configuration β€” it is the *same* cluster, +with the state ledger living in the object store instead of a local directory. `cluster.yaml` +accepts a `storage:` key naming the root: + +```yaml +version: 1 +storage: s3://omnigraph/probe +graphs: + probe: + schema: graphs/probe.pg +``` + +`omnigraph cluster import --config ` then `cluster apply --config ` write +`s3://omnigraph/probe/__cluster/state.json` and create the graph (`applied_count: 2`, +`converged: true`). The server then boots and serves it: + +``` +$ curl http://127.0.0.1:9110/healthz +{"status":"ok","version":"0.10.0","internal_schema_version":6} +$ curl http://127.0.0.1:9110/graphs/probe/schema +{"schema_source":"node Person {\n coco_key: String @key\n name: String\n}\n"} +``` + +The probe itself, adding a nullable property to the served schema: + +``` +$ curl -X POST http://127.0.0.1:9110/graphs/probe/schema/apply \ + -H 'Content-Type: application/json' \ + -d '{"schema_source":"node Person {\n coco_key: String @key\n name: String\n nickname: String?\n}\n"}' + +{"error":"server-side schema apply is disabled for cluster-backed serving; update the cluster config, run `omnigraph cluster apply`, and restart the server.","code":"conflict"} +HTTP 409 +``` + +**Why this generalises.** The refusal is scoped to "cluster-backed serving", and `--cluster` is the +server's *only* boot source (`omnigraph-server --help`: "The server's only boot source (RFC-011 +cluster-only)"). There is no `--store` flag on the server. Every served graph is therefore +cluster-backed, so `POST /schema/apply` is disabled unconditionally. The endpoint exists in the +OpenAPI document, but no server configuration accepts it. + +**Consequence:** `managed_by="system"` cannot be implemented over HTTP by any means available to +the connector. Since every other cocoindex connector supports `SYSTEM` and defaults to it, dropping +the feature was rejected (decision 32); Epic D keeps it through the `cluster apply` control plane +instead, behind an `apply_schema()` seam so a future server-side apply (D3) replaces it in one +place. + +## P2 β€” Is there any HTTP path that creates a graph? + +**Answer: no.** + +``` +$ curl -X POST http://127.0.0.1:9110/graphs -d '{"graph_id":"newgraph"}' +HTTP 405 (empty body) + +$ curl http://127.0.0.1:9110/graphs/nosuchgraph/schema +{"error":"graph 'nosuchgraph' not found","code":"not_found"} +HTTP 404 +``` + +The full served path list contains no graph-creation route β€” `/graphs` is `GET` only: + +``` +GET /graphs +GET,HEAD /graphs/{graph_id}/blob +GET,POST /graphs/{graph_id}/branches +POST /graphs/{graph_id}/branches/merge +DELETE /graphs/{graph_id}/branches/{branch} +POST /graphs/{graph_id}/change +GET /graphs/{graph_id}/changes +POST /graphs/{graph_id}/changes/baseline +GET /graphs/{graph_id}/commits +GET /graphs/{graph_id}/commits/{commit_id} +GET /graphs/{graph_id}/commits/{commit_id}/changes +POST /graphs/{graph_id}/export +POST /graphs/{graph_id}/ingest +POST /graphs/{graph_id}/load +POST /graphs/{graph_id}/load/ndjson +POST /graphs/{graph_id}/mutate +POST /graphs/{graph_id}/mutate/if-graph-commit +GET /graphs/{graph_id}/queries +POST /graphs/{graph_id}/queries/{name} +POST /graphs/{graph_id}/queries/{name}/if-graph-commit +POST /graphs/{graph_id}/query +POST /graphs/{graph_id}/read +GET /graphs/{graph_id}/schema +POST /graphs/{graph_id}/schema/apply (409 β€” see P1) +GET /graphs/{graph_id}/snapshot +GET /healthz +``` + +**Consequence:** graph provisioning is permanently an operator task (`cluster.yaml` β†’ +`cluster apply` β†’ restart). The connector's uninitialized-graph branch becomes a clear failure +naming the graph, the base URL, and the provisioning procedure β€” never an attempt to create one. +The 404 body above is the signal to key on. + +`GET /graphs` is served here because the server runs `--unauthenticated` with no policy bundle; +decision 15 (never call it) still stands, since it is forbidden by default under a policy. + +## P12 β€” Does a storage-root server pick up an applied revision without a restart? (added) + +**Answer: no. `cluster apply` succeeds and the store is updated, but a running server keeps serving +the old schema indefinitely. A restart picks it up.** + +Served schema before: + +``` +{"schema_source":"node Person {\n coco_key: String @key\n name: String\n}\n"} +``` + +Added a nullable property to `graphs/probe.pg`, then: + +``` +$ omnigraph cluster apply --config ../scratch/p1probe --json +ok: True | converged: True | applied: 1 + graph.probe update derived + schema.probe update applied +``` + +Polled `GET /graphs/probe/schema` every 3s for 60s β€” the served source never changed. After +`pkill omnigraph-server` and an identical re-boot, the same request returns: + +``` +{"schema_source":"node Person {\n coco_key: String @key\n name: String\n nickname: String?\n}\n"} +``` + +**Two consequences.** + +1. A restart is unavoidable for any connector-driven schema change, on either boot mode. There is + no polling, no signal, no in-place reload. +2. **`cluster apply` never contacted the server.** It ran from a plain config directory on a + different host from the one serving, wrote `__cluster/state.json` into the bucket, and the + server only learned about it at boot. So the control plane needs the CLI binary, the config + directory, and object-store credentials β€” but *not* colocation with the server. That is a + materially weaker requirement than the original spec assumed, and it means a cocoindex worker + can drive schema itself as long as something else can restart the server. + +## Probes not yet run + +P3 (write/merge serialization), P4 (merge conflict shape), P5 (endpoint-not-found wording over +HTTP), P6 (entity cap status), P7 (mixed upsert/delete), P8 (branch delete consent), P9 +(`delete_branch: true`), P10 (Cedar actions), P11 (`if-graph-commit` 412 body). diff --git a/docs/src/content/docs/connectors/index.mdx b/docs/src/content/docs/connectors/index.mdx index eefd70c14..66b52d02a 100644 --- a/docs/src/content/docs/connectors/index.mdx +++ b/docs/src/content/docs/connectors/index.mdx @@ -58,6 +58,7 @@ Connectors that export data out of a flow. They're grouped by the kind of store, | [Neo4j](/docs/connectors/neo4j/) | Property graph over Bolt β€” node and relationship tables, vector indexes | | [FalkorDB](/docs/connectors/falkordb/) | Redis-backed Cypher graph, per-graph multitenancy, vector indexes | | [SurrealDB](/docs/connectors/surrealdb/) | Multi-model β€” normal and relation (edge) tables, vector indexes | +| [Omnigraph](/docs/connectors/omnigraph/) | Lakehouse-native property graph with `.pg` schema and git-style branching | ### Message streams diff --git a/docs/src/content/docs/connectors/omnigraph.mdx b/docs/src/content/docs/connectors/omnigraph.mdx new file mode 100644 index 000000000..36b7dcca8 --- /dev/null +++ b/docs/src/content/docs/connectors/omnigraph.mdx @@ -0,0 +1,420 @@ +--- +title: "*Omnigraph* connector" +toc_max_heading_level: 4 +description: > + Write to Omnigraph β€” a lakehouse-native property-graph database β€” with + support for node types, edge types, `.pg`-schema evolution, git-style + branching, and a synthetic `coco_key` property that gives every managed + node and edge a stable, filterable identity. +--- + +The `omnigraph` connector writes records to [Omnigraph](https://github.com/ModernRelay/omnigraph), a lakehouse-native property-graph database: typed nodes and edges declared in a `.pg` schema, Lance-backed storage, and git-style branching, all driven through a CLI. It supports node types, edge types, additive `.pg` schema evolution, and per-`ConnectionFactory` branch targeting. + +```python +from cocoindex.connectors import omnigraph +``` + +:::note[Requirements] +Unlike the other graph connectors, this one needs **no extra Python package** β€” it shells out to the `omnigraph` CLI binary rather than linking a driver. Install the binary and make sure it's on `PATH` (or point `ConnectionFactory(cli=...)` at it directly). Everything else β€” `pip install cocoindex` β€” is unchanged. + +Verified against Omnigraph 0.10.0 (earlier releases are not supported). The connector only invokes CLI subcommands present in that release; HTTP transport is a future increment (the `omnigraph-server` HTTP surface exists but is not yet used here). +::: + +## Connection setup + +Create a `ConnectionFactory` and provide it via a `ContextKey`. `store` is a graph URI β€” typically a local, git-ignored `file://` path β€” and `branch` names the branch every write in this connection targets. + +:::note +The key name is load-bearing across runs β€” it's the stable identity CocoIndex uses to track managed nodes and edges. See [ContextKey as stable identity](../programming_guide/context#contextkey-as-stable-identity) before renaming. +::: + +```python +from collections.abc import AsyncIterator +from cocoindex.connectors import omnigraph +import cocoindex as coco + +OG: coco.ContextKey[omnigraph.ConnectionFactory] = coco.ContextKey("og") + +@coco.lifespan +async def coco_lifespan(builder: coco.EnvironmentBuilder) -> AsyncIterator[None]: + builder.provide( + OG, + omnigraph.ConnectionFactory( + store="file:///var/lib/cocoindex/graph.omni", + branch="main", + ), + ) + yield +``` + +```python +class ConnectionFactory: + store: str # a graph URI, e.g. "file:///abs/path/g.omni" + branch: str = "main" # the branch every write through this connection targets + cli: str = "omnigraph" # override if the binary isn't on PATH +``` + +If the graph at `store` hasn't been initialized yet, the connector runs `omnigraph init` for you on the first schema write β€” there's no separate provisioning step. + +### Multiple graphs or branches + +Pair each graph (or branch) with its own `ContextKey` and `ConnectionFactory`: + +```python +OG: coco.ContextKey[omnigraph.ConnectionFactory] = coco.ContextKey("og") +STAGING_OG: coco.ContextKey[omnigraph.ConnectionFactory] = coco.ContextKey("staging_og") + +@coco.lifespan +async def coco_lifespan(builder: coco.EnvironmentBuilder) -> AsyncIterator[None]: + builder.provide(OG, omnigraph.ConnectionFactory(store="file:///data/g.omni", branch="main")) + builder.provide(STAGING_OG, omnigraph.ConnectionFactory(store="file:///data/g.omni", branch="staging")) + yield +``` + +:::caution[A non-`main` branch blocks every schema change] +Omnigraph applies schema to the whole graph, not per branch, and refuses to do it at all while any other branch exists: `schema apply requires a graph with only main; found non-main branches: staging`. So pointing a `ConnectionFactory` at a branch other than `main` works **only** if the connector never needs to write schema β€” that is, if every type it touches is `managed_by="user"`. With any `managed_by="system"` type, the first type creation or schema change on that store fails. + +(This does not affect the scratch branches the connector opens for oversized syncs: those are transient and deleted before the next schema write. If a process is killed mid-sync and leaves one behind, the next schema change on the store deletes any `coco_scratch_*` branch whose owner is provably gone and retries β€” only the connector's own prefix is ever reaped, and only after its per-branch lock has been taken.) +::: + +A type's tracking identity is `(ContextKey, type name)` β€” the store URI and branch live on the `ConnectionFactory` and are resolved at write time, not baked into the tracking records. Two apps that should not share tracking must therefore use two different `ContextKey`s, as above; repointing one `ContextKey`'s `ConnectionFactory` at a different branch reuses the same records, which now describe another branch's contents. + +## As target + +The `omnigraph` connector provides target state APIs for writing nodes to node types and edges to edge types. CocoIndex tracks what nodes and edges should exist and automatically handles upserts and deletions. + +### Declaring target states + +#### Node types (parent state) + +Declares a node type as a target state. Returns a `NodeTarget` for declaring nodes. + +```python +def mount_node_target( + db: coco.ContextKey[ConnectionFactory], + type_name: str, + schema: NodeSchema, + *, + managed_by: Literal["system", "user"] = "system", +) -> NodeTarget[Any] +``` + +**Parameters:** + +- `db` β€” A `ContextKey[omnigraph.ConnectionFactory]` for the Omnigraph connection. +- `type_name` β€” The `.pg` node type name (e.g. `"Source"`). +- `schema` β€” A `NodeSchema` describing the type's properties and key. Unlike the sibling connectors, this is required, not optional β€” a node type is always self-describing. +- `managed_by` β€” Whether CocoIndex manages the type's `.pg` schema (`"system"`) or assumes it exists (`"user"` β€” never written, never validated; see [Schema evolution](#schema-declaration-and-evolution) below). + +**Returns:** A resolved `NodeTarget`. `declare_node_target` (sync, for endpoint-only types) and `node_target` (the raw `TargetState`, for composing with `coco.mount`) are also available. + +#### Nodes (child states) + +Once a `NodeTarget` is resolved, declare nodes to be upserted (translated to a keyed `insert`, which Omnigraph treats as an upsert by key): + +```python +def NodeTarget.declare_node( + self, + *, + node: RowT, +) -> None +``` + +**Parameters:** + +- `node` β€” A row object (dict, dataclass, NamedTuple, or Pydantic model) carrying every property named in `schema`, including the key field(s). + +#### Edge types (parent state) + +Declares an edge type as a target state. `from_target`/`to_target` are passed **positionally**, so a mistyped endpoint is a mount-time Python error rather than a write-time one; the connector uses them to render `edge {Name}: {From} -> {To} { ... }` and to build endpoint stubs (see [Cross-component ordering](#cross-component-ordering) below). + +```python +def mount_edge_target( + db: coco.ContextKey[ConnectionFactory], + type_name: str, + from_target: NodeTarget[Any], + to_target: NodeTarget[Any], + schema: EdgeSchema | None = None, + *, + managed_by: Literal["system", "user"] = "system", +) -> EdgeTarget[Any] +``` + +**Parameters:** + +- `db`, `type_name`, `managed_by` β€” as above. `type_name` must not already name a **node** type in the same graph: Omnigraph accepts both in a schema, but a mutation resolves the bare name to the node type, leaving the edge unwritable. Declaring one raises `ValueError` on the sync that would have created it. +- `from_target` / `to_target` β€” the `NodeTarget`s this edge type connects. +- `schema` β€” An `EdgeSchema` with the edge's own properties. Required if you intend to pass `record=` to `declare_edge`: Omnigraph is schema-first, so an edge type's properties live in its `.pg` block and an insert naming an undeclared one is refused β€” passing a record to a schema-less edge type raises `TypeError` at declare time. (Neo4j's `MERGE` creates relationship properties on the fly, so a ported app that mounted a relation bare and still passed a record needs a schema added here.) An edge has no key of its own β€” its identity is always `(from_id, to_id)` (see below) β€” so `EdgeSchema.from_class` takes no `key=`, and a `NodeSchema` is refused here. + +**Returns:** A resolved `EdgeTarget`. `declare_edge_target` and `edge_target` are also available, matching the node-type trio. + +Mounting validates both endpoints β€” see [Endpoint requirements](#endpoint-requirements). + +#### Edges (child states) + +Once an `EdgeTarget` is resolved, declare edges: + +```python +def EdgeTarget.declare_edge( + self, + *, + from_id: Any, + to_id: Any, + record: RowT | None = None, +) -> None +``` + +**Parameters:** + +- `from_id` / `to_id` β€” the source and target nodes' key values. Each must be a single scalar the endpoint type's key can be rendered from β€” see [Endpoint requirements](#endpoint-requirements). +- `record` β€” Optional row object whose fields populate the edge's own properties. + + +Omnigraph assigns every edge its own internal id (a ULID) that isn't addressable β€” there is no way to `insert`-then-upsert an edge by a caller-supplied id the way a keyed node works. The connector's identity for an edge is always the `(from_id, to_id)` pair: a second `declare_edge` call for the same pair updates that one edge (delete-then-insert under the hood) rather than creating a second, parallel edge. + +### The synthetic `coco_key` property + +Every type this connector manages β€” node or edge β€” gets one extra, non-nullable property: `coco_key: String`. You never set it yourself; it isn't a field on your dataclass, and declaring a property literally named `coco_key` raises `ValueError`. Every type the connector created also declares a second synthetic property, `coco_managed_by_: Bool?` β€” the app being your `AppConfig.name`, made an identifier if it isn't one β€” which is never written and stays `null` on every row: it is how the connector tells the types this app currently owns from ones you declared and from another app's. Don't declare it on a `managed_by="user"` block. Every property name starting with `coco_` is reserved for these; a dataclass field named that way raises `ValueError`. + +It exists because Omnigraph gives the connector no other way to address one specific entity for an update or delete: + +- Edges have no caller-suppliable id (see above), and their auto-assigned ULID isn't a filterable GQ property. +- `where` accepts exactly **one** equality predicate β€” `where from = $a and to = $b` is a parse error, so a node's or edge's own key fields can't be combined into a compound filter either. + +So the connector derives a stable value (a fingerprint of the node's key tuple, or of `(from_id, to_id)` for an edge) and writes it into `coco_key` on every insert, then deletes and updates by `where coco_key = $value`. This is purely internal bookkeeping β€” it shows up in `omnigraph export`/`schema show` output, but nothing in the public API asks you to touch it. + +### Node schema: from Python class + +Build a `NodeSchema` by introspecting a record type: + +```python +@classmethod +async def NodeSchema.from_class( + cls, + target: type, + *, + key: str | Sequence[str], +) -> NodeSchema +``` + +**Parameters:** + +- `target` β€” A dataclass. Unlike the sibling connectors' `TableSchema.from_class`, this does not also accept a NamedTuple or Pydantic model β€” `NodeSchema.from_class` introspects via `dataclasses.fields()` directly. It must not have a field named **`id`** β€” that's Omnigraph's own node identity column, materialized from the `@key` property, and declaring one alongside it makes the schema invalid. On an edge type, `src`, `dst`, `from` and `to` are reserved too. All of these raise `ValueError` naming the field. +- `key` β€” The field name forming the type's key. Required β€” an unkeyed `insert` is a strict insert in Omnigraph, so every re-run would duplicate every node. Exactly **one** field: Omnigraph node types support a single `@key`, and a two-`@key` block is rejected outright (`node type X has multiple @key constraints; only one is supported`). A `Sequence[str]` is accepted for signature parity with the sibling connectors' `primary_key=`, but it must hold exactly one name. If your entity is identified by a combination of fields, derive a single id from them and key on that. A `Date` key is tracked by its ISO form, and a `DateTime` key by the instant it denotes in epoch milliseconds β€” the graph's own identity for it, so two offsets spelling one instant are one node, sub-millisecond precision is dropped, and a naive `datetime` is read as UTC. Such keys must be passed as `datetime.date` / `datetime.datetime` values: a string or an integer for a `Date` or `DateTime` key is refused at `declare_node`, since only the Python value yields the graph's identity exactly. + +**Returns:** A `NodeSchema` populated from the class's fields. + +An edge type's own properties come from `EdgeSchema` instead β€” the same introspection, minus the key an edge never has: + +```python +@classmethod +async def EdgeSchema.from_class(cls, target: type) -> EdgeSchema +``` + +Or construct either directly: `NodeSchema(properties={...}, key=("slug",))`, `EdgeSchema(properties={...})`, with each value a `PropertyDef`. + +#### Default Python β†’ `.pg` type mapping + +| Python type | `.pg` type | Notes | +|---|---|---| +| `bool` | `Bool` | Checked before `int` β€” `bool` is a subclass of `int`. | +| `int` | `I64` | | +| `float` | `F64` | | +| `str` | `String` | | +| `datetime.date` | `Date` | Encoded via `.isoformat()`. | +| `datetime.datetime` | `DateTime` | Encoded via `.isoformat()`; the engine accepts ISO with or without a trailing `Z`. | +| `X \| None` | `X?` | Any of the above, made nullable. | +| `list[X]` for scalar `X` | `[X]` | See the list-element constraint below. | +| `bytes` | *(no mapping)* | See below β€” raises `TypeError`. | + +`.pg` also has `I32`, `U32`, `U64`, `F32`, `Vector(N)`, and `enum(a, b, ...)` keywords, plus `Blob`, but none of these has a default Python mapping β€” reach them with `OmnigraphType` (below) and your own encoder. Vector-index attachments (the equivalent of the neo4j connector's `declare_vector_index`) aren't implemented yet. + +**`bytes` has no mapping, on purpose.** Omnigraph's `Blob` type is an external URI reference the engine *fetches* β€” a `file://` value β€” not inline bytes. A Python `bytes` value can never legally become a `Blob`, so `bytes` fields raise `TypeError` at schema-build time rather than being silently coerced into something the engine would reject (or worse, quietly accept as the wrong thing). If you need blob-like content, store the URI as a `str` property and manage the referenced file yourself. + +**List elements must be non-nullable scalars.** `list[str]` maps to `[String]`, but `list[str | None]` (`[String?]`) and `list[list[str]]` (`[[String]]`) both raise `TypeError` at schema-build time β€” the engine rejects both forms outright (`"expected core_type"` / `"expected base_type"`). + +#### `OmnigraphType` + +Override the default mapping for a single property with `OmnigraphType`, via `typing.Annotated`: + +```python +from typing import Annotated +from dataclasses import dataclass +from cocoindex.connectors.omnigraph import OmnigraphType + +@dataclass +class Reading: + sensor: str + value: Annotated[int, OmnigraphType("I32")] # narrower than the default I64 +``` + +```python +class OmnigraphType(NamedTuple): + pg_type: str +``` + +The `pg_type` string is validated against Omnigraph's type grammar (`validate_pg_type`) at schema-build time β€” an invalid expression fails fast in Python, before it would otherwise reach the engine as generated `.pg` text. + +### Node schema: explicit property definitions + +Build a `NodeSchema` directly when the row shape is dynamic: + +```python +from cocoindex.connectors.omnigraph import NodeSchema, PropertyDef + +schema = NodeSchema( + properties={ + "slug": PropertyDef("slug", "String"), + "title": PropertyDef("title", "String"), + "note": PropertyDef("note", "String?"), + }, + key=("slug",), +) +``` + +`PropertyDef` fields: + +- `name` β€” The property name (must match its key in `properties`). +- `pg_type` β€” The `.pg` type string (see the table above; `?` suffix for nullable). +- `encoder` β€” Optional `Callable[[Any], Any]` applied to non-`None` values before they're sent to Omnigraph. Change detection fingerprints the encoded value, so changing an encoder rewrites every node or edge it applies to on the next update β€” including rows declared from a memoized component: the type's tracking record remembers each property's encoder (by qualified name, plus a digest of the code for a Python function), and a change there re-runs every component that declared into the type. `Date`/`DateTime` properties get an ISO-format encoding by default, and a list of either gets it mapped over the elements β€” whether the `PropertyDef` came from `from_class` or was written by hand with no encoder. Passing `datetime.date.isoformat` on a `Date` key or `datetime.datetime.isoformat` on a `DateTime` key yourself is the same encoding and is accepted; each is accepted only on the type it encodes, since `datetime.date.isoformat` applied to a `datetime` drops the time and would key two instants on one node. A key property may not carry a custom encoder β€” mounting raises `ValueError` β€” because the key is the node's identity and edges address it by the raw key value: an encoder there would make the graph key on something edges never reference, and changing it would leave the old node beside the new one. Normalize key values before declaring the node. (The built-in `Date`/`DateTime` encoders are allowed on a key; they never change, and such types cannot be edge endpoints.) + + +## Schema declaration and evolution + +For each managed type, the connector maintains a `.pg` fragment and reconciles it against the graph's actual schema on change. Because Omnigraph's `schema apply`/`init` each take the graph's **complete** desired schema β€” not a per-type diff β€” the connector reads the current whole-graph `.pg` source back, replaces (or appends) just this type's block, and writes the merged whole back. A store-scoped file lock, keyed by the canonical store identity (so `file:///a/store` and `file:///a/./store` share it), serializes that read-merge-write sequence across components and processes on the same host. It also covers scratch branches, because Omnigraph rejects schema changes while any non-main branch exists; each live scratch branch additionally holds a lock of its own, and a leftover `coco_scratch_*` branch is only ever reaped after that lock has been taken, so a branch another process is still using is never deleted. Processes on different hosts are not coordinated by these local locks and must not apply schema changes to the same store concurrently. + +| Schema change | What happens | Effect on existing nodes/edges | +|---|---|---| +| Type not yet in the graph | `.pg` fragment applied (`init` on the very first type, `schema apply` after) | β€” | +| Property added | Additive `schema apply` | None re-upserted β€” existing rows just gain the column | +| Property dropped or retyped | `schema apply` | Every existing node/edge of that type is forced to re-upsert | +| `@key` changed, or a node/edge switched shape | Type dropped and recreated (two `schema apply` calls β€” see below) | Every child record is re-declared from scratch | +| Type no longer declared | Its block is removed from the schema (`schema apply`) β€” **unless** it is `managed_by="user"`, which is left alone | Every node/edge of that type is removed with it; nothing is removed for a `managed_by="user"` type | +| `managed_by="user"` | Nothing β€” never mutates `.pg`; migrate the type yourself with `omnigraph schema apply`, then declare the new shape | The declaration is tracked as-is; a write that names a property the graph lacks fails with the engine's own error | + +**A non-nullable property can only be declared at type creation, never added by migration.** Omnigraph's `schema apply` rejects adding a non-nullable property to an existing type outright (`"requires a backfill and is not supported in schema migration v1"`). The connector checks for this in Python and raises a clear `ValueError` naming the property β€” make it optional (`X | None`), or accept that adding it means dropping and recreating the type. + +**A `@key` change can't be applied in place.** The engine flatly rejects an in-place key change as an unsupported migration step. The connector works around this with two separate `schema apply` calls: drop the type's block from the schema and apply that, then apply again with the type re-added under its new definition. This is exactly the `"destructive"` row above β€” every child node or edge of that type is re-declared from scratch, since their prior tracking is invalidated along with the old key. + +When `managed_by="user"` is set, the connector never mutates `.pg` for that type. Node/edge-level upserts and deletes still work normally β€” **provided your own `.pg` declares `coco_key: String` on that type.** + +The one exception is the handoff itself. When a type the connector created is re-declared with `managed_by="user"`, the next update writes once to drop the `coco_managed_by_` property from its block and touches nothing else; from then on the block is yours, and a later drop of a node type it references refuses rather than removing it. Declaring it `managed_by="system"` again re-renders the block from your dataclass, `coco_managed_by_` included, and the connector owns it once more. + +:::caution[`managed_by="user"` types must declare `coco_key` yourself] +The connector addresses every entity through the synthetic `coco_key` property, and under `managed_by="user"` it never writes the schema β€” so it cannot add that column for you. Every insert names `coco_key`, and every delete filters on it. Against a type whose `.pg` lacks it, *all* row writes fail with `type 'X' has no property 'coco_key'`. + +Add it to the type's definition when you create the graph: + +``` +node Source { + slug: String @key + title: String? + coco_key: String +} +``` +::: + +The connector does not compare your declaration against the graph, on the first run or later: the type is adopted as declared. If your dataclass and your `.pg` drift apart, the first write that names a property the graph lacks fails with the engine's own error, which names the property. + +Removing a `managed_by="user"` type from your app is also never destructive: the connector didn't create the type and won't drop it. + +**Dropping a node type that an edge type still references.** Omnigraph rejects a schema with an unresolved endpoint, and each mounted type is torn down by its own component β€” so when you drop an app, or stop declaring a node type and its edge types together, the node type's removal may reach the graph before the edge type's. The connector removes the referencing edge types *this app* owns (those whose block declares its `coco_managed_by_` property) in the same schema write; the edge type's own removal then finds nothing left to do. Any other edge type is never removed on your behalf β€” one another app of this connector owns (its block names that app), one you declared with `managed_by="user"`, or another tool's: the drop fails naming it and, where there is one, the app it belongs to, and you either have that app stop declaring it, remove it from the schema yourself, or keep the node type. A `@key` change on a referenced node type fails the same way, since the edge stays declared. + +### Cross-component ordering + +An edge's endpoint node is routinely owned by a *different* component than the edge itself β€” in a per-file pipeline, a corpus-wide "attended" edge might reference a `Meeting` node a per-file component declares independently, and the two commit on their own schedules. The connector handles an edge arriving before its endpoint reactively: it applies the edge insert optimistically first, and only on an engine "not found" failure does it insert a key-only stub for the reported-missing endpoint and retry β€” one stub-and-retry round per distinct endpoint the component's edges reference, since the engine names only the first missing endpoint per attempt. The endpoint's own owning component later upserts its real row over the stub (a keyed `insert` is an upsert), so nothing is duplicated. A stub is not tracked as a target state of its own: if no component ever declares that node, the key-only stub stays in the graph after the edge is removed, until its node type is dropped β€” the same as an endpoint the Neo4j connector's `MERGE` creates. Declare every endpoint node somewhere if you need them cleaned up with their edges. Stubbing reactively rather than unconditionally matters: a keyed `insert` is a full-record *replace* in Omnigraph, not a partial merge, so stubbing an endpoint that already has real data would silently null out its other nullable properties. + +The same shape covers the reverse order. Deleting a node in Omnigraph cascades to every edge at it, silently as far as the component that declared those edges is concerned β€” its tracking would still say they exist, and they would never be written again. So when a component stops declaring a node that an edge outside that same sync still references, the connector reduces the node to the key-only stub instead of deleting it: the edge survives, and the node's next declaration fills the row back in. A node undeclared together with all of its edges is deleted outright. To know which is which, a sync that deletes nodes reads the schema once and, for each node whose type some edge type has as an endpoint, reads that node's edges β€” one read per edge type and end. A stub left this way is not tracked, like any other stub. + +### Endpoint requirements + +`mount_edge_target` validates both endpoint types at **mount time**, not mid-sync β€” so a mistake raises `ValueError` naming the type immediately, rather than failing later from whatever unrelated component happens to declare the edge first. Two requirements, both consequences of how Omnigraph stores an edge: + +**Every non-key property must be optional.** The stub above can only ever carry the endpoint's key property, so a non-nullable property outside the key would make the stub insert impossible (`insert for 'Meeting' must provide non-nullable property 'note_file'`). Declare them as `X | None`. + +**The key must be a `String` or an integer.** An edge's `from`/`to` doesn't hold the key property β€” it holds the endpoint node's `id`, which the engine renders as a String from the key value whatever that key's declared type is. CocoIndex reproduces that rendering with `str(value)`, which is exact for strings and integers and for nothing else: a `Date` key of `2026-01-05` has the id `"20458"` (days since the epoch) and a `DateTime` id is epoch milliseconds. Key such a type on a generated string or integer id instead. + +## Branching and the sink + +Within one component's sync, the connector partitions its writes by target graph, and orders node upserts before edge inserts (so endpoints exist before edges reference them) and edge deletes before node deletes (so no edge dangles). If everything fits in one Omnigraph commit β€” a single kind of change (all upserts, or all deletes) for one type, no more than 8,192 entities of any one type and 32,768 in total β€” it's sent as one direct write to the branch configured on `ConnectionFactory`. + +When a component's writes don't fit β€” more than 8,192 entities of one type, or a mix of upserts *and* deletes in the same sync (Omnigraph refuses to combine `insert`/`update` with `delete` in a single mutation) β€” the connector opens a scratch branch off the target branch, applies every chunk there, and merges once the whole component is written. A single-commit sync that turns out to need an endpoint stub escalates to the same path: the recovery is two commits, and running it against the live branch could leave an orphan stub node behind if the retry then failed. Nothing is lost by the attempt β€” a failed Omnigraph mutation applies nothing at all. This keeps the component's target-state sync atomic from a reader's perspective, even though it lands as several underlying Omnigraph commits. The scratch branch is cleaned up (deleted) whether the merge succeeds or fails; a failed merge doesn't leave orphan branches behind. A process killed between creating and deleting the branch does leave one, and Omnigraph then refuses schema changes on the store until it is gone β€” so the next schema change reaps any `coco_scratch_*` branch it finds (under the same store lock a live scratch branch is held in, which is what makes a visible one provably abandoned) and retries. + +Every mutation reaches the CLI as two files β€” `--query` for the GQ source and `--params-file` for the bound values β€” never as inline arguments. A full-size commit is over 2 MB across the two, which exceeds the operating system's argument-list limit; the file forms have no size limit, so the 8,192-entity cap is a real, reachable number rather than one the transport would fail long before. + +## Example: Node and edge types + +```python +from collections.abc import AsyncIterator +from dataclasses import dataclass +from typing import Any +import cocoindex as coco +from cocoindex.connectors import omnigraph + +OG: coco.ContextKey[omnigraph.ConnectionFactory] = coco.ContextKey("og") + + +# Both types are used as edge endpoints below, so their non-key properties +# are optional -- see "Endpoint requirements". +@dataclass +class Source: + slug: str + title: str | None + + +@dataclass +class Claim: + slug: str + statement: str | None + + +@coco.lifespan +async def coco_lifespan(builder: coco.EnvironmentBuilder) -> AsyncIterator[None]: + builder.provide(OG, omnigraph.ConnectionFactory(store="file:///tmp/demo.omni")) + yield + + +@coco.fn +async def app_main() -> None: + sources = await omnigraph.mount_node_target( + OG, "Source", await omnigraph.NodeSchema.from_class(Source, key="slug"), + ) + claims = await omnigraph.mount_node_target( + OG, "Claim", await omnigraph.NodeSchema.from_class(Claim, key="slug"), + ) + supports = await omnigraph.mount_edge_target(OG, "Supports", sources, claims) + + sources.declare_node(node=Source(slug="overview", title="Overview")) + claims.declare_node(node=Claim(slug="claim-1", statement="CocoIndex is declarative.")) + supports.declare_edge(from_id="overview", to_id="claim-1") + + +app = coco.App(coco.AppConfig(name="demo_to_omnigraph"), app_main) +``` + +Generated schema: + +``` +node Source { + slug: String @key + title: String? + coco_key: String + coco_managed_by_demo_to_omnigraph: Bool? +} + +node Claim { + slug: String @key + statement: String? + coco_key: String + coco_managed_by_demo_to_omnigraph: Bool? +} + +edge Supports: Source -> Claim { + coco_key: String + coco_managed_by_demo_to_omnigraph: Bool? +} +``` + +## Example + +See [`examples/meeting_notes_graph_omnigraph`](https://github.com/cocoindex-io/cocoindex/tree/main/examples/meeting_notes_graph_omnigraph) for the port, side-by-side with [`examples/meeting_notes_graph_neo4j`](https://github.com/cocoindex-io/cocoindex/tree/main/examples/meeting_notes_graph_neo4j). diff --git a/docs/src/data/docs-meta.json b/docs/src/data/docs-meta.json index acf2ca413..124dff918 100644 --- a/docs/src/data/docs-meta.json +++ b/docs/src/data/docs-meta.json @@ -86,6 +86,9 @@ "connectors/oci_object_storage": { "reviewedTs": 1783166400 }, + "connectors/omnigraph": { + "reviewedTs": 1783166400 + }, "connectors/postgres": { "reviewedTs": 1783166400 }, diff --git a/docs/src/data/docs-sidebar.ts b/docs/src/data/docs-sidebar.ts index 4cea40299..913e69a84 100644 --- a/docs/src/data/docs-sidebar.ts +++ b/docs/src/data/docs-sidebar.ts @@ -75,6 +75,7 @@ export const sidebar: SidebarItem[] = [ { type: 'doc', slug: 'connectors/localfs', label: 'Local filesystem' }, { type: 'doc', slug: 'connectors/neo4j', label: 'Neo4j' }, { type: 'doc', slug: 'connectors/oci_object_storage', label: 'OCI Object Storage' }, + { type: 'doc', slug: 'connectors/omnigraph', label: 'Omnigraph' }, { type: 'doc', slug: 'connectors/postgres', label: 'Postgres' }, { type: 'doc', slug: 'connectors/qdrant', label: 'Qdrant' }, { type: 'doc', slug: 'connectors/snowflake', label: 'Snowflake' }, diff --git a/docs/src/data/examples.ts b/docs/src/data/examples.ts index 469039a7e..0e32ae90b 100644 --- a/docs/src/data/examples.ts +++ b/docs/src/data/examples.ts @@ -695,6 +695,7 @@ export const EXAMPLE_CATALOG_GROUPS: ExampleCatalogGroup[] = [ blurb: 'Extract entities and relationships into graph databases that stay in sync.', entries: [ { dir: 'meeting_notes_graph_falkordb', title: 'Meeting Notes β†’ Knowledge Graph Β· FalkorDB', description: 'Extract structured info from Google Drive meeting notes into a FalkorDB knowledge graph.' }, + { dir: 'meeting_notes_graph_omnigraph', title: 'Meeting Notes β†’ Knowledge Graph Β· Omnigraph', description: 'Extract structured info from Google Drive meeting notes into an Omnigraph knowledge graph.' }, { dir: 'product_recommendation', docs: 'product-recommendation', title: 'Product Recommendation Graph', description: 'LLM-extract what each product is and what pairs with it from product docs, into a Neo4j graph of products and taxonomies that powers recommendations.', run: RUN_MAIN }, ], }, diff --git a/examples/meeting_notes_graph_omnigraph/.env.example b/examples/meeting_notes_graph_omnigraph/.env.example new file mode 100644 index 000000000..df975ba2a --- /dev/null +++ b/examples/meeting_notes_graph_omnigraph/.env.example @@ -0,0 +1,28 @@ +# Example environment variables for this example +# Copy this to .env and fill in your actual values + +COCOINDEX_DB=./cocoindex.db + +# OpenAI API key (used via LiteLLM) +#! PLEASE FILL IN +OPENAI_API_KEY= + +# Google Drive service account credential path +#! PLEASE FILL IN +GOOGLE_SERVICE_ACCOUNT_CREDENTIAL=/path/to/service_account_credential.json + +# Google Drive root folder IDs, comma separated +#! PLEASE FILL IN +GOOGLE_DRIVE_ROOT_FOLDER_IDS=id1,id2 + +# Omnigraph store: a graph URI (local file:// path by default) and branch. +# Requires the `omnigraph` CLI binary on PATH. +OMNIGRAPH_STORE=file:///tmp/meeting_notes.omni +# Keep this as `main`. Omnigraph refuses `schema apply` while any other +# branch exists, so a non-main branch blocks every schema change this +# example needs to create its node and edge types. +OMNIGRAPH_BRANCH=main + +# LLM models (LiteLLM-prefixed; e.g. openai/gpt-5-mini, anthropic/claude-..., ...) +LLM_MODEL=openai/gpt-5-mini +RESOLUTION_LLM_MODEL=openai/gpt-5-mini diff --git a/examples/meeting_notes_graph_omnigraph/.gitignore b/examples/meeting_notes_graph_omnigraph/.gitignore new file mode 100644 index 000000000..cfccfee0f --- /dev/null +++ b/examples/meeting_notes_graph_omnigraph/.gitignore @@ -0,0 +1,5 @@ +.env +*.db +*.omni +.venv/ +*.egg-info/ diff --git a/examples/meeting_notes_graph_omnigraph/README.md b/examples/meeting_notes_graph_omnigraph/README.md new file mode 100644 index 000000000..e9eff1e36 --- /dev/null +++ b/examples/meeting_notes_graph_omnigraph/README.md @@ -0,0 +1,122 @@ +

Turn meeting notes into a self-updating graph in Omnigraph.

+ +

+ An LLM pulls the organizer, participants, and tasks out of each meeting; an embedding + LLM pass collapses "Alice", "Alice Chen", and "alice c." into one Person node β€” into Omnigraph, in plain async Python.
+ Point it at a Drive folder of Markdown notes, and it re-extracts only the note you edited, then reconciles the graph. +

+ +

+ Star us β€οΈ β†’ Star CocoIndex on GitHub  Β·  + CocoIndex documentation  Β·  + Join the CocoIndex Discord +

+ +
+ +[![stars](https://img.shields.io/github/stars/cocoindex-io/cocoindex?style=flat-square&label=stars&color=FB6A76)](https://github.com/cocoindex-io/cocoindex) +[![pypi](https://img.shields.io/pypi/v/cocoindex?style=flat-square&label=pypi&color=E59A63)](https://pypi.org/project/cocoindex/) +[![discord](https://img.shields.io/discord/1314801574169673738?style=flat-square&logo=discord&logoColor=white&label=discord&color=5865F2)](https://discord.com/invite/zpA9S2DR7s) +[![license](https://img.shields.io/badge/license-Apache--2.0-5B5BD6?style=flat-square)](https://opensource.org/licenses/Apache-2.0) + +
+ +
+ +This is the meeting-notes knowledge graph, targeting [Omnigraph](https://github.com/ModernRelay/omnigraph) instead of Neo4j β€” a lakehouse-native property-graph database with a `.pg` schema, Lance-backed storage, and git-style branching, that you talk to through a CLI. Meeting notes are a graph pretending to be a folder of documents: every note records who ran the meeting, who showed up, what got decided, and who owns each task. But it's prose, scattered across a shared drive, so you can full-text search it and not much else. You declare the transformation in native Python and your own types β€” `target_state = transformation(source_state)` β€” and the heavy lifting (incremental processing, change tracking, managed graph targets) runs in a Rust engine underneath, so editing one note re-extracts one note, and the graph reconciles itself: no orphaned people, no stale edges, no cleanup scripts. + +## How it works + +Three node types, three edge types, and "who is on the hook for what" becomes an edge you traverse: + +- **`Meeting`** nodes β€” one per meeting section, keyed by a stable integer id derived from `(note_file, date)`. +- **`Person`** nodes β€” canonical organizers, participants, and assignees, deduplicated by an embedding + LLM entity-resolution pass. +- **`Task`** nodes β€” tasks decided in meetings, keyed by description. +- **`ATTENDED`** edges β€” `Person -> Meeting`, carrying an `is_organizer` flag. **`DECIDED`** edges β€” `Meeting -> Task`. **`ASSIGNED_TO`** edges β€” `Person -> Task`. + +Because people are shared across notes, the pipeline runs in three phases β€” read it top-to-bottom in [`main.py`](main.py): + +```python +@coco.fn(memo=True) # Phase 1 β€” per note: split into meetings, declare Meeting/Task + DECIDED, carry raw names forward +async def process_file(file, meeting_table, task_table, decided_rel) -> list[MeetingExtraction]: + for section in _split_meetings(await file.read_text()): + extracted = await extract_meeting(section) + meeting_id = await id_generator.next_id(extracted.time) + meeting_table.declare_node(node=Meeting(meeting_id=meeting_id, ...)) + for task in extracted.tasks: + task_table.declare_node(node=Task(description=task.description)) + decided_rel.declare_edge(from_id=meeting_id, to_id=task.description) + ... + +@coco.fn(memo=True) # Phase 2 β€” collapse "Alice" / "Alice Chen" / "alice c." into canonical names +async def _resolve_persons(raw_persons: set[str]) -> ResolvedEntities: + return await resolve_entities(entities=raw_persons, embedder=coco.use_context(EMBEDDER), + resolve_pair=LlmPairResolver(model=coco.use_context(RESOLUTION_LLM_MODEL))) + +@coco.fn # Phase 3 β€” declare canonical Person nodes + ATTENDED / ASSIGNED_TO using resolved names +async def create_person_relations(meetings, persons, person_table, attended_rel, assigned_rel) -> None: + for canonical_name in persons.canonicals(): person_table.declare_node(node=Person(name=canonical_name)) + ... +``` + +Extraction is [instructor](https://github.com/instructor-ai/instructor) over [LiteLLM](https://docs.litellm.ai/) with your own Pydantic models; `DECIDED` and `ASSIGNED_TO` carry no payload, so the Omnigraph connector derives their identity from the endpoints β€” one edge per pair. + +

+ πŸ“˜ Full Tutorial β†’
+ The closest walkthrough is the Neo4j version β€” same extraction, resolution, and three-phase flow; only the graph store differs. Step-by-step coverage of the property-graph schema, entity resolution, and exactly what happens on each kind of change. +

+ +## Why it's worth a star ⭐ + +- **Entity resolution built in.** CocoIndex's [`entity_resolution`](https://cocoindex.io/docs/ops/entity_resolution/) op embeds every raw name, filters by vector similarity, and asks the LLM to confirm *only* the close pairs β€” so the same person written five ways collapses to one node, cheaply. +- **Cross-file nodes, owned in one place.** People are shared across notes, so no single note's component can own a `Person` node. The two cross-file phases own the canonical set and the person-touching edges, exactly once. +- **Incremental by default.** `@coco.fn(memo=True)` caches each extraction by content; edit one note and only that note re-extracts, then resolution and the graph diff. A no-change re-run makes zero LLM calls. +- **Two models on purpose.** A stronger `LLM_MODEL` does the structured extraction; a cheaper `RESOLUTION_LLM_MODEL` confirms resolution pairs β€” both are [LiteLLM provider strings](https://docs.litellm.ai/docs/providers) you can swap. +- **Honest cache busting.** The model ids and embedder are declared with `detect_change=True`, so swapping any of them re-extracts against it with no cache to clear by hand. +- **No database to run.** Omnigraph's `file://` store is a local, git-ignored path β€” no container, no server. Just the `omnigraph` CLI binary on `PATH`. + +## Run it + +**1. Install the `omnigraph` CLI** and make sure it's on `PATH`. The graph itself needs no separate setup: the connector runs `omnigraph init` for you on the first write. + +**2. Configure & install** β€” this source reads notes from one or more Google Drive folders shared with a service account (see [Setting up a service account](https://cocoindex.io/docs/connectors/google_drive/#setting-up-a-service-account)): + +```sh +cp .env.example .env # set OPENAI_API_KEY, GOOGLE_SERVICE_ACCOUNT_CREDENTIAL, GOOGLE_DRIVE_ROOT_FOLDER_IDS +pip install -e . +``` + +`OMNIGRAPH_STORE` defaults to a local `file:///tmp/meeting_notes.omni`; `OMNIGRAPH_BRANCH` defaults to `main`. + +**3. Build the graph:** + +```sh +cocoindex update main +``` + +**4. Explore the graph:** + +```sh +# Every row on the branch, one JSON object per line +omnigraph export --store "$OMNIGRAPH_STORE" --branch "$OMNIGRAPH_BRANCH" + +# Who attended which meetings, including organizer status +omnigraph export --store "$OMNIGRAPH_STORE" --branch "$OMNIGRAPH_BRANCH" \ + | jq -sc '.[] | select(.edge == "ATTENDED")' + +# Everything one person is on the hook for +omnigraph export --store "$OMNIGRAPH_STORE" --branch "$OMNIGRAPH_BRANCH" \ + | jq -sc '.[] | select(.edge == "ASSIGNED_TO" and .from == "Alice Chen")' +``` + +Or query directly with GQ (`insert`/`update`/`delete` only β€” reads go through `export`): `omnigraph mutate --store "$OMNIGRAPH_STORE" --json -e '...'`. See the [connector docs](https://cocoindex.io/docs/connectors/omnigraph/#the-synthetic-coco_key-property) for why Omnigraph addresses a specific node or edge by the synthetic `coco_key` property rather than a caller-chosen id. + +This pipeline is the [docs knowledge graph](https://cocoindex.io/docs/examples/docs-to-knowledge-graph/) plus an entity-resolution pass β€” the natural next step when the LLM names the same thing two ways. Prefer another graph store? See the [Neo4j variant](https://github.com/cocoindex-io/cocoindex/tree/main/examples/meeting_notes_graph_neo4j). + +--- + +

+ If this turned your shared drive into a graph, give CocoIndex a star ⭐ β€” it helps a lot.
+ Docs Β· Walkthrough Β· Discord Β· See all examples β†’ +

+ + diff --git a/examples/meeting_notes_graph_omnigraph/main.py b/examples/meeting_notes_graph_omnigraph/main.py new file mode 100644 index 000000000..b5f030496 --- /dev/null +++ b/examples/meeting_notes_graph_omnigraph/main.py @@ -0,0 +1,462 @@ +""" +Meeting Notes Graph (v1) β€” CocoIndex pipeline example, Omnigraph flavor. + +Ingest Markdown meeting notes from Google Drive, split each note into +per-meeting sections at heading boundaries, extract structured information +with LiteLLM + instructor, deduplicate person names with embedding-based +entity resolution, and build a knowledge graph in Omnigraph: + + Meeting nodes β€” one per meeting section + Person nodes β€” canonical organizers, participants, and task assignees + Task nodes β€” tasks decided in meetings + + ATTENDED Person -> Meeting (with is_organizer flag) + DECIDED Meeting -> Task + ASSIGNED_TO Person -> Task + +The pipeline runs in three phases: + 1. Per-file extraction declares Meeting and Task nodes plus DECIDED edges, + and emits raw (un-resolved) person names for downstream resolution. + 2. Person entity resolution maps raw names to canonical names. + 3. A final pass declares canonical Person nodes and the person-touching + edges (ATTENDED, ASSIGNED_TO) using resolved names. + +Side-by-side diff with examples/meeting_notes_graph_neo4j/main.py: the flow +is identical, and the API is graph-native β€” node/edge targets, declare_node +and declare_edge β€” where the neo4j version speaks in tables and relations. +What else differs, and why: + + * The connector import, the ConnectionFactory arguments (uri + auth + + database vs store + branch), the AppConfig name, and the schema + classes: NodeSchema.from_class(..., key=...) for node types and + EdgeSchema.from_class(...) for an edge type's own properties. + * Meeting.id is `meeting_id` here. `id` is Omnigraph's own node identity + column; declaring a property by that name makes the schema invalid. + * Meeting's non-key properties are optional. An edge may be written before + the component owning its endpoint node has run, and the connector + recovers by inserting a key-only stub β€” which Omnigraph rejects if the + type declares a non-nullable property outside its key. Neo4j has no such + constraint, so its Meeting can require them. + * ATTENDED is mounted WITH a schema. Omnigraph is schema-first: an edge + type's properties are declared in `.pg`, and an insert naming an + undeclared one is refused. Neo4j's MERGE creates them on the fly, so the + original can mount ATTENDED bare and still pass a record. +""" + +from __future__ import annotations + +import asyncio +import datetime +import os +import re +from collections.abc import AsyncIterator +from dataclasses import dataclass +from typing import Any, cast + +import cocoindex as coco +import instructor +import litellm +import pydantic +from cocoindex.connectors import google_drive, omnigraph +from cocoindex.ops.entity_resolution import ResolvedEntities, resolve_entities +from cocoindex.ops.entity_resolution.llm_resolver import LlmPairResolver +from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder +from cocoindex.resources.id import IdGenerator + +litellm.drop_params = True + + +# --------------------------------------------------------------------------- +# Context keys +# --------------------------------------------------------------------------- + +KG_DB = coco.ContextKey[omnigraph.ConnectionFactory]("kg_db") +LLM_MODEL = coco.ContextKey[str]("llm_model", detect_change=True) +RESOLUTION_LLM_MODEL = coco.ContextKey[str]("resolution_llm_model", detect_change=True) +EMBEDDER = coco.ContextKey[SentenceTransformerEmbedder]("embedder", detect_change=True) + + +# --------------------------------------------------------------------------- +# Lifespan +# --------------------------------------------------------------------------- + + +@coco.lifespan +async def coco_lifespan( + builder: coco.EnvironmentBuilder, +) -> AsyncIterator[None]: + builder.provide( + KG_DB, + omnigraph.ConnectionFactory( + store=os.environ.get("OMNIGRAPH_STORE", "file:///tmp/meeting_notes.omni"), + branch=os.environ.get("OMNIGRAPH_BRANCH", "main"), + ), + ) + builder.provide(LLM_MODEL, os.environ.get("LLM_MODEL", "openai/gpt-5-mini")) + builder.provide( + RESOLUTION_LLM_MODEL, + os.environ.get("RESOLUTION_LLM_MODEL", "openai/gpt-5-mini"), + ) + builder.provide( + EMBEDDER, + SentenceTransformerEmbedder("Snowflake/snowflake-arctic-embed-xs"), + ) + yield + + +# --------------------------------------------------------------------------- +# Omnigraph record schemas (dataclasses for declare_node / declare_edge) +# --------------------------------------------------------------------------- + + +@dataclass +class Meeting: + meeting_id: int # Generated via generate_id((note_file, time_iso)) + # Optional so Meeting stays usable as an edge endpoint: the key-only stub + # the connector inserts when an edge races ahead of the component that + # owns the node cannot populate a non-nullable property outside the key. + note_file: str | None + time: datetime.date | None + note: str | None + + +@dataclass +class Person: + name: str # canonical + + +@dataclass +class Task: + description: str + + +@dataclass +class AttendedRel: + """ATTENDED edge payload. The edge's identity is always + (from_id=person, to_id=meeting_id), giving exactly one edge per + (person, meeting); these are only its properties. + """ + + is_organizer: bool + + +# DECIDED and ASSIGNED_TO carry no payload β€” declared without schema or +# record, with the connector deriving PKs from (from_id, to_id). + + +# --------------------------------------------------------------------------- +# LLM extraction schemas (Pydantic, for instructor) +# --------------------------------------------------------------------------- + + +class ExtractedPerson(pydantic.BaseModel): + name: str = pydantic.Field( + description="Full name of the person, as written in the note." + ) + + +class ExtractedTask(pydantic.BaseModel): + description: str = pydantic.Field( + description="Concise, standalone description of the task or action item." + ) + assigned_to: list[ExtractedPerson] = pydantic.Field( + default_factory=list, + description="People the task is assigned to.", + ) + + +class ExtractedMeeting(pydantic.BaseModel): + time: datetime.date = pydantic.Field( + description="Date of the meeting in ISO format (YYYY-MM-DD)." + ) + note: str = pydantic.Field( + description="A brief summary or notes from the meeting section.", + ) + organizer: ExtractedPerson = pydantic.Field( + description="The person who organized or led the meeting." + ) + participants: list[ExtractedPerson] = pydantic.Field( + default_factory=list, + description=( + "People who attended the meeting other than the organizer. " + "Do not include the organizer here." + ), + ) + tasks: list[ExtractedTask] = pydantic.Field( + default_factory=list, + description="Action items or tasks decided in the meeting.", + ) + + +EXTRACT_PROMPT = """\ +You are an expert at reading meeting notes and extracting structured information. + +Given a single meeting section (Markdown), extract: +- The meeting date (look for a date in the heading or body; required). +- A brief note summarizing what the meeting was about. +- The organizer (the person who ran the meeting). If unclear, pick the person + who appears most central to the meeting. +- Participants other than the organizer. +- Tasks or action items decided, including who they are assigned to. + +Return only what is supported by the text. Use full names where available. +""" + + +# --------------------------------------------------------------------------- +# LLM extraction +# --------------------------------------------------------------------------- + + +@coco.fn(memo=True) +async def extract_meeting(section_text: str) -> ExtractedMeeting: + """Extract a structured Meeting from a Markdown section via LiteLLM + instructor.""" + client = cast( + instructor.AsyncInstructor, + instructor.from_litellm(litellm.acompletion, mode=instructor.Mode.JSON), + ) + result = await client.chat.completions.create( + model=coco.use_context(LLM_MODEL), + response_model=ExtractedMeeting, + messages=[ + {"role": "system", "content": EXTRACT_PROMPT}, + {"role": "user", "content": section_text}, + ], + ) + # Re-validate to restore class identity for pickling. + return ExtractedMeeting.model_validate(result.model_dump()) + + +# --------------------------------------------------------------------------- +# Splitting β€” match v0's `\n\n##? ` heading regex +# --------------------------------------------------------------------------- + +_HEADING_RE = re.compile(r"\n\n##?\s+") + + +def _split_meetings(text: str) -> list[str]: + parts = _HEADING_RE.split("\n\n" + text) + return [p.strip() for p in parts if p.strip()] + + +# --------------------------------------------------------------------------- +# Internal transfer types (Phase 1 β†’ Phase 3) +# --------------------------------------------------------------------------- + + +@dataclass +class MeetingExtraction: + """Raw per-meeting data carried forward to entity resolution + relation declaration.""" + + meeting_id: int + organizer: str # raw name + participants: list[str] # raw names + task_assignees: list[ + tuple[str, list[str]] + ] # (task_description, [raw assignee names]) + + +# --------------------------------------------------------------------------- +# Phase 1: per-meeting and per-file processing +# --------------------------------------------------------------------------- + + +@coco.fn(memo=True) +async def process_file( + file: google_drive.DriveFile, + meeting_table: omnigraph.NodeTarget[Meeting], + task_table: omnigraph.NodeTarget[Task], + decided_rel: omnigraph.EdgeTarget[Any], +) -> list[MeetingExtraction]: + text = await file.read_text() + note_file = file.file_path.path.as_posix() + id_generator = IdGenerator() + extractions = [] + for section in _split_meetings(text): + extracted = await extract_meeting(section) + meeting_id = await id_generator.next_id(extracted.time) + + meeting_table.declare_node( + node=Meeting( + meeting_id=meeting_id, + note_file=note_file, + time=extracted.time, + note=extracted.note, + ) + ) + + for task in extracted.tasks: + task_table.declare_node(node=Task(description=task.description)) + decided_rel.declare_edge(from_id=meeting_id, to_id=task.description) + + extractions.append( + MeetingExtraction( + meeting_id=meeting_id, + organizer=extracted.organizer.name, + participants=[p.name for p in extracted.participants], + task_assignees=[ + (t.description, [a.name for a in t.assigned_to]) + for t in extracted.tasks + ], + ) + ) + return extractions + + +# --------------------------------------------------------------------------- +# Phase 2: Person entity resolution +# --------------------------------------------------------------------------- + + +@coco.fn(memo=True) +async def _resolve_persons(raw_persons: set[str]) -> ResolvedEntities: + return await resolve_entities( + entities=raw_persons, + embedder=coco.use_context(EMBEDDER), + resolve_pair=LlmPairResolver(model=coco.use_context(RESOLUTION_LLM_MODEL)), + ) + + +# --------------------------------------------------------------------------- +# Phase 3: declare canonical Person nodes + person-touching relations +# --------------------------------------------------------------------------- + + +@coco.fn +async def create_person_relations( + meetings: list[MeetingExtraction], + persons: ResolvedEntities, + person_table: omnigraph.NodeTarget[Person], + attended_rel: omnigraph.EdgeTarget[Any], + assigned_rel: omnigraph.EdgeTarget[Any], +) -> None: + # Declare canonical Person nodes. + for canonical_name in persons.canonicals(): + person_table.declare_node(node=Person(name=canonical_name)) + + for m in meetings: + # ATTENDED β€” aggregate organizer + participants. Organizer flag wins + # on collision so a person listed as both gets a single edge with + # is_organizer=true. Resolution happens before aggregation so two + # raw names that resolve to the same person also collapse. + attendees: dict[str, bool] = {persons.canonical_of(m.organizer): True} + for p in m.participants: + attendees.setdefault(persons.canonical_of(p), False) + + for canonical, is_organizer in attendees.items(): + attended_rel.declare_edge( + from_id=canonical, + to_id=m.meeting_id, + record=AttendedRel(is_organizer=is_organizer), + ) + + # ASSIGNED_TO β€” dedup per (canonical person, task description). + for task_desc, assignees in m.task_assignees: + seen: set[str] = set() + for raw in assignees: + canonical = persons.canonical_of(raw) + if canonical in seen: + continue + seen.add(canonical) + assigned_rel.declare_edge(from_id=canonical, to_id=task_desc) + + +# --------------------------------------------------------------------------- +# App main +# --------------------------------------------------------------------------- + + +@coco.fn +async def app_main() -> None: + # --- Mount node types --- + meeting_table = await omnigraph.mount_node_target( + KG_DB, + "Meeting", + await omnigraph.NodeSchema.from_class(Meeting, key="meeting_id"), + ) + person_table = await omnigraph.mount_node_target( + KG_DB, + "Person", + await omnigraph.NodeSchema.from_class(Person, key="name"), + ) + task_table = await omnigraph.mount_node_target( + KG_DB, + "Task", + await omnigraph.NodeSchema.from_class(Task, key="description"), + ) + + # --- Mount edge types --- + # ATTENDED carries is_organizer, so it needs a schema: Omnigraph declares + # an edge type's properties in `.pg` and refuses an insert naming one it + # doesn't have. + attended_rel = await omnigraph.mount_edge_target( + KG_DB, + "ATTENDED", + person_table, + meeting_table, + await omnigraph.EdgeSchema.from_class(AttendedRel), + ) + decided_rel = await omnigraph.mount_edge_target( + KG_DB, "DECIDED", meeting_table, task_table + ) + assigned_rel = await omnigraph.mount_edge_target( + KG_DB, "ASSIGNED_TO", person_table, task_table + ) + + # --- Phase 1: per-file extraction --- + credential_path = os.environ["GOOGLE_SERVICE_ACCOUNT_CREDENTIAL"] + root_folder_ids = [ + folder.strip() + for folder in os.environ["GOOGLE_DRIVE_ROOT_FOLDER_IDS"].split(",") + if folder.strip() + ] + source = google_drive.GoogleDriveSource( + service_account_credential_path=credential_path, + root_folder_ids=root_folder_ids, + ) + + file_coros = [] + async for path_key, file in source.items(): + file_coros.append( + coco.use_mount( + coco.component_subpath("file", path_key), + process_file, + file, + meeting_table, + task_table, + decided_rel, + ) + ) + per_file: list[list[MeetingExtraction]] = list(await asyncio.gather(*file_coros)) + all_meetings: list[MeetingExtraction] = [m for ms in per_file for m in ms] + + # --- Phase 2: Person entity resolution --- + raw_persons: set[str] = set() + for m in all_meetings: + raw_persons.add(m.organizer) + raw_persons.update(m.participants) + for _task_desc, assignees in m.task_assignees: + raw_persons.update(assignees) + + persons = await coco.use_mount( + coco.component_subpath("resolve_persons"), + _resolve_persons, + raw_persons, + ) + + # --- Phase 3: declare Person nodes + person-touching relations --- + await coco.mount( + coco.component_subpath("person_relations"), + create_person_relations, + all_meetings, + persons, + person_table, + attended_rel, + assigned_rel, + ) + + +app = coco.App( + coco.AppConfig(name="MeetingNotesGraphOmnigraph"), + app_main, +) diff --git a/examples/meeting_notes_graph_omnigraph/pyproject.toml b/examples/meeting_notes_graph_omnigraph/pyproject.toml new file mode 100644 index 000000000..de9241a97 --- /dev/null +++ b/examples/meeting_notes_graph_omnigraph/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "meeting-notes-graph-omnigraph" +version = "0.1.0" +description = "CocoIndex example: build a knowledge graph from meeting notes in Google Drive, stored in Omnigraph." +requires-python = ">=3.11" +dependencies = [ + "cocoindex[google_drive,sentence_transformers,entity_resolution_llm]>=1.0.7", + "instructor", + "litellm", + "pydantic", +] + +[tool.setuptools] +py-modules = ["main"] diff --git a/python/cocoindex/_internal/component_ctx.py b/python/cocoindex/_internal/component_ctx.py index 898795062..6123fdbb3 100644 --- a/python/cocoindex/_internal/component_ctx.py +++ b/python/cocoindex/_internal/component_ctx.py @@ -401,6 +401,14 @@ def app_main() -> None: return value +def current_app_name() -> str: + """Name of the app whose component is running, as given to `AppConfig`. + + Raises `RuntimeError` outside a component context. + """ + return get_context_from_ctx()._core_processor_ctx.app_name + + def get_component_context() -> ComponentContext: """ Get the current ComponentContext explicitly. diff --git a/python/cocoindex/_internal/core.pyi b/python/cocoindex/_internal/core.pyi index 6bbc8b33b..a21abb3a9 100644 --- a/python/cocoindex/_internal/core.pyi +++ b/python/cocoindex/_internal/core.pyi @@ -99,6 +99,8 @@ class ComponentProcessorContext: @property def stable_path(self) -> StablePath: ... @property + def app_name(self) -> str: ... + @property def live(self) -> bool: ... def join_fn_call(self, child_fn_ctx: FnCallContext) -> None: ... def initial_context_memo_states( diff --git a/python/cocoindex/connectors/omnigraph/__init__.py b/python/cocoindex/connectors/omnigraph/__init__.py new file mode 100644 index 000000000..30139a1a6 --- /dev/null +++ b/python/cocoindex/connectors/omnigraph/__init__.py @@ -0,0 +1,4 @@ +from . import _target +from ._target import * + +__all__ = _target.__all__ diff --git a/python/cocoindex/connectors/omnigraph/_client.py b/python/cocoindex/connectors/omnigraph/_client.py new file mode 100644 index 000000000..a8667574b --- /dev/null +++ b/python/cocoindex/connectors/omnigraph/_client.py @@ -0,0 +1,446 @@ +"""Transport for the Omnigraph connector β€” the only module that does I/O.""" + +from __future__ import annotations + +import asyncio +import contextlib +import dataclasses +import hashlib +import json +import pathlib +import posixpath +import sys +import tempfile +import urllib.parse +import urllib.request +from collections.abc import AsyncIterator, Iterator +from typing import BinaryIO + +if sys.platform == "win32": + import msvcrt +else: + import fcntl + +from cocoindex.connectors.omnigraph._gq import Query + +if sys.platform == "win32": + + def _lock_file(lock_file: BinaryIO) -> None: + # ``msvcrt.locking`` locks a byte range, so make sure byte zero exists + # before asking for an exclusive lock on it. + lock_file.seek(0, 2) + if lock_file.tell() == 0: + lock_file.write(b"\0") + lock_file.flush() + lock_file.seek(0) + msvcrt.locking(lock_file.fileno(), msvcrt.LK_LOCK, 1) + + def _try_lock_file(lock_file: BinaryIO) -> bool: + lock_file.seek(0, 2) + if lock_file.tell() == 0: + lock_file.write(b"\0") + lock_file.flush() + lock_file.seek(0) + try: + msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1) + except OSError: + return False + return True + + def _unlock_file(lock_file: BinaryIO) -> None: + lock_file.seek(0) + msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1) + +else: + + def _lock_file(lock_file: BinaryIO) -> None: + fcntl.flock(lock_file, fcntl.LOCK_EX) + + def _try_lock_file(lock_file: BinaryIO) -> bool: + try: + fcntl.flock(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError: + return False + return True + + def _unlock_file(lock_file: BinaryIO) -> None: + fcntl.flock(lock_file, fcntl.LOCK_UN) + + +def canonical_store(uri: str) -> str: + """One identity for every spelling of the same store. + + The store lock is keyed by this. Keyed by the raw URI text, two + processes addressing one store as `file:///a/store` and + `file:///a/./store` took different locks β€” and one of them reaped the + scratch branch the other was still using (verified live). A `file://` + URI resolves to its real path; any other scheme gets a lower-cased + scheme and host and a normalised path. + """ + parts = urllib.parse.urlsplit(uri) + scheme = parts.scheme.lower() + if scheme == "file": + return pathlib.Path(urllib.request.url2pathname(parts.path)).resolve().as_uri() + path = posixpath.normpath(parts.path) if parts.path else "" + return urllib.parse.urlunsplit((scheme, parts.netloc.lower(), path, "", "")) + + +def _lock_dir() -> pathlib.Path: + return pathlib.Path(tempfile.gettempdir()) + + +@contextlib.contextmanager +def _temporary_text_file(content: str, *, suffix: str) -> Iterator[str]: + """Write a closed, short-lived file that another process can reopen. + + Windows denies reopening a ``NamedTemporaryFile`` while its original + handle is still open. Creating it with ``delete=False`` lets us close the + handle before invoking the CLI and still remove the file deterministically + afterward on every platform. + """ + path: pathlib.Path | None = None + try: + with tempfile.NamedTemporaryFile( + "w", suffix=suffix, encoding="utf-8", delete=False + ) as f: + f.write(content) + path = pathlib.Path(f.name) + yield str(path) + finally: + if path is not None: + path.unlink(missing_ok=True) + + +@dataclasses.dataclass(frozen=True) +class ConnectionFactory: + """Identifies an Omnigraph store and the branch to write to. + + Provided once in the app lifespan and resolved at action time via the + ContextProvider β€” never captured at declare time, since delete actions run + in a process where the declaring code never executes. + + ``store`` is a URI (``file:///abs/path/g.omni``), not a filesystem path. + """ + + store: str + branch: str = "main" + cli: str = "omnigraph" + + +class OmnigraphCliError(RuntimeError): + """Non-zero exit from the omnigraph CLI, carrying its stderr.""" + + +class _CliClient: + def __init__(self, conn: ConnectionFactory) -> None: + self._conn = conn + + @property + def store_lock_path(self) -> pathlib.Path: + digest = hashlib.sha256(canonical_store(self._conn.store).encode()).hexdigest() + return _lock_dir() / f"cocoindex-omnigraph-{digest}.lock" + + @staticmethod + def _scratch_branch_lock_path(name: str) -> pathlib.Path: + return _lock_dir() / f"cocoindex-omnigraph-scratch-{name}.lock" + + @contextlib.asynccontextmanager + async def store_lock(self) -> AsyncIterator[None]: + """Serialize operations that require exclusive access to this store, + across components and processes on this host.""" + lock_file = await asyncio.to_thread(self.store_lock_path.open, "a+b") + try: + await asyncio.to_thread(_lock_file, lock_file) + yield + finally: + await asyncio.to_thread(_unlock_file, lock_file) + await asyncio.to_thread(lock_file.close) + + @contextlib.asynccontextmanager + async def hold_scratch_branch(self, name: str) -> AsyncIterator[None]: + """Hold the liveness lock of scratch branch `name` for the block β€” + from before the branch is created until after it is deleted β€” so a + reaper anywhere on this host can tell the branch is in use.""" + path = self._scratch_branch_lock_path(name) + lock_file = await asyncio.to_thread(path.open, "a+b") + try: + await asyncio.to_thread(_lock_file, lock_file) + yield + finally: + await asyncio.to_thread(_unlock_file, lock_file) + await asyncio.to_thread(lock_file.close) + with contextlib.suppress(OSError): + await asyncio.to_thread(path.unlink) + + @contextlib.asynccontextmanager + async def claim_scratch_branch(self, name: str) -> AsyncIterator[bool]: + """Try to take the liveness lock of scratch branch `name` without + waiting. Yields True if it was free β€” its owner is gone, so the + branch is abandoned and the block may delete it β€” and False while + another process still holds it.""" + path = self._scratch_branch_lock_path(name) + lock_file = await asyncio.to_thread(path.open, "a+b") + claimed = False + try: + claimed = await asyncio.to_thread(_try_lock_file, lock_file) + yield claimed + finally: + if claimed: + await asyncio.to_thread(_unlock_file, lock_file) + await asyncio.to_thread(lock_file.close) + if claimed: + with contextlib.suppress(OSError): + await asyncio.to_thread(path.unlink) + + # --- argv builders (pure, unit-tested) --- + + def _mutate_argv( + self, query_path: str, params_path: str, *, branch: str + ) -> list[str]: + """Both the GQ source and the bound params go by FILE, never inline. + + The inline forms (`-e `, `--params `) put the whole commit + into argv, and a commit is a whole component's writes: at the 8,192 + entity cap that is ~2.3 MB across the two arguments, well past + darwin's 1 MB `ARG_MAX` β€” `create_subprocess_exec` raises `OSError: + [Errno 7] Argument list too long`, which isn't an `OmnigraphCliError` + and so isn't caught anywhere. Linux binds tighter still: its + 128 KiB-per-argument `MAX_ARG_STRLEN` caps the expression alone at + roughly 800 entities whatever `ARG_MAX` allows. + + `--query ` and `--params-file ` take exactly the same + input with no size limit at all (verified against the binary, both + with and without the positional query name, and at the full 8,192 + cap). So the transport simply doesn't put payloads in argv. + """ + return [ + self._conn.cli, + "mutate", + "--store", + self._conn.store, + "--branch", + branch, + "--query", + query_path, + "--params-file", + params_path, + "--json", + "--quiet", + ] + + def _query_argv( + self, query_path: str, params_path: str, *, branch: str + ) -> list[str]: + """The read side of `_mutate_argv`: same file-borne payloads.""" + return [ + self._conn.cli, + "query", + "--store", + self._conn.store, + "--branch", + branch, + "--query", + query_path, + "--params-file", + params_path, + "--json", + "--quiet", + ] + + def _merge_argv(self, name: str, *, into: str) -> list[str]: + # `branch merge` has no compare-and-swap precondition (as of 0.10.0 + # only `mutate` takes `--if-commit`). A conflict is a non-zero exit. + return [ + self._conn.cli, + "branch", + "merge", + name, + "--into", + into, + "--store", + self._conn.store, + "--json", + "--quiet", + ] + + def _init_argv(self, schema_path: str) -> list[str]: + # `init` takes the graph URI POSITIONALLY, not via --store, and has + # no --json flag at all β€” it prints a plain "initialized " line + # to stdout instead of JSON. + return [ + self._conn.cli, + "init", + "--schema", + schema_path, + "--quiet", + self._conn.store, + ] + + def _apply_schema_argv(self, schema_path: str) -> list[str]: + return [ + self._conn.cli, + "schema", + "apply", + "--schema", + schema_path, + "--store", + self._conn.store, + "--json", + "--quiet", + ] + + def _schema_show_argv(self) -> list[str]: + return [ + self._conn.cli, + "schema", + "show", + "--store", + self._conn.store, + "--json", + "--quiet", + ] + + def _branch_create_argv(self, name: str, *, frm: str) -> list[str]: + return [ + self._conn.cli, + "branch", + "create", + name, + "--from", + frm, + "--store", + self._conn.store, + "--json", + "--quiet", + ] + + def _branch_delete_argv(self, name: str) -> list[str]: + return [ + self._conn.cli, + "branch", + "delete", + name, + "--store", + self._conn.store, + "--json", + "--quiet", + ] + + def _branch_list_argv(self) -> list[str]: + return [ + self._conn.cli, + "branch", + "list", + "--store", + self._conn.store, + "--json", + "--quiet", + ] + + # --- execution --- + + async def _run(self, argv: list[str]) -> dict[str, object]: + proc = await asyncio.create_subprocess_exec( + *argv, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE + ) + try: + out, err = await proc.communicate() + except BaseException: + # Cancellation must not leave the CLI running: `communicate()` + # does nothing to the child when the awaiting task is cancelled, + # so a cancelled `mutate` kept writing to the store after the + # connector had given up on it. Kill it and reap it, then let + # the cancellation propagate. + if proc.returncode is None: + with contextlib.suppress(ProcessLookupError): + proc.kill() + await proc.wait() + raise + if proc.returncode != 0: + raise OmnigraphCliError( + f"{' '.join(argv[:3])} exited {proc.returncode}: " + f"{err.decode(errors='replace').strip()}" + ) + if "--json" not in argv: + # `init` has no --json flag; its stdout is a plain diagnostic + # line, not JSON, so there is nothing to parse. + return {} + text = out.decode(errors="replace").strip() + return json.loads(text) if text else {} + + async def init_graph(self, schema_pg: str) -> None: + with _temporary_text_file(schema_pg, suffix=".pg") as path: + await self._run(self._init_argv(path)) + + async def apply_schema(self, schema_pg: str) -> None: + with _temporary_text_file(schema_pg, suffix=".pg") as path: + await self._run(self._apply_schema_argv(path)) + + async def read_schema(self) -> str | None: + """Return the graph's current, complete `.pg` schema source, or + `None` if the graph hasn't been `init`'d yet. + + Omnigraph's schema is applied whole-graph, not per type: `schema + apply`/`init` treat their input as the complete desired schema, so + callers that reconcile one type at a time must read this back and + merge before writing, rather than applying a single type's + fragment directly (see `_gq.merge_type_into_schema`). + + Not-yet-initialized is detected on the CLI's stderr text β€” the raw + Lance "Dataset ... not found" (verified against the binary) β€” since + the exit code alone doesn't distinguish it from any other failure. + """ + try: + result = await self._run(self._schema_show_argv()) + except OmnigraphCliError as e: + msg = str(e).lower() + # Match the engine's actual phrasing, not the two words + # separately: the store URI is echoed into every error, so a store + # under `~/datasets/` would otherwise turn any unrelated + # "not found" failure into a bogus "graph not initialized". + if "dataset at path" in msg and "was not found" in msg: + return None + raise + source = result["schema_source"] + assert isinstance(source, str) + return source + + async def mutate(self, mutation: Query, *, branch: str) -> None: + with ( + _temporary_text_file(mutation.expr, suffix=".gq") as query_path, + _temporary_text_file( + json.dumps(mutation.params), suffix=".json" + ) as params_path, + ): + await self._run(self._mutate_argv(query_path, params_path, branch=branch)) + + async def query(self, query: Query, *, branch: str) -> list[dict[str, object]]: + """Run a read query and return its rows, one dict per row keyed by + the `return` clause's column aliases.""" + with ( + _temporary_text_file(query.expr, suffix=".gq") as query_path, + _temporary_text_file( + json.dumps(query.params), suffix=".json" + ) as params_path, + ): + result = await self._run( + self._query_argv(query_path, params_path, branch=branch) + ) + rows = result["rows"] + assert isinstance(rows, list) + return rows + + async def branch_create(self, name: str, *, frm: str) -> None: + await self._run(self._branch_create_argv(name, frm=frm)) + + async def branch_merge(self, name: str, *, into: str) -> None: + await self._run(self._merge_argv(name, into=into)) + + async def branch_delete(self, name: str) -> None: + await self._run(self._branch_delete_argv(name)) + + async def branch_list(self) -> list[str]: + result = await self._run(self._branch_list_argv()) + branches = result["branches"] + assert isinstance(branches, list) + return [str(name) for name in branches] diff --git a/python/cocoindex/connectors/omnigraph/_gq.py b/python/cocoindex/connectors/omnigraph/_gq.py new file mode 100644 index 000000000..b5409435e --- /dev/null +++ b/python/cocoindex/connectors/omnigraph/_gq.py @@ -0,0 +1,725 @@ +"""Pure builders for Omnigraph `.pg` schema fragments and `.gq` mutations. + +Every function here is a pure function over strings and primitives β€” no I/O, +no connector types β€” so the whole module is unit-testable without a running +Omnigraph. Mirrors the role `neo4j/_cypher.py` plays for the Neo4j connector. +""" + +from __future__ import annotations + +import hashlib +import re +from collections.abc import Sequence +from typing import NamedTuple + +#: Synthetic property carrying the connector's derived key, declared on every +#: type we manage. Omnigraph gives edges no settable id and its auto-assigned +#: ULID is not filterable, while `where` accepts exactly one equality +#: predicate β€” so a declared property is the only way to address one specific +#: entity for deletion. Non-nullable, which is legal at init time. +COCO_KEY = "coco_key" + +#: Prefix of every synthetic property this connector declares β€” `coco_key` +#: and the ownership property below. A user's own property may not start +#: with it. +COCO_PREFIX = "coco_" + +#: Prefix of the synthetic, nullable, never-written property declared on +#: every block this connector renders, and on nothing else; the rest of the +#: name is the app that owns the block (see `ownership_property`). It is how +#: the schema sink tells this app's types from a user's and from another +#: app's: a `managed_by=user` type must declare `coco_key` as well, so that +#: property alone cannot, and a marker that only said "some app of this +#: connector" let one app's drop take another app's edge types along. It is +#: a property rather than a comment because the engine stores an applied +#: schema's source only when the apply changes something structural β€” a +#: comment-only change reports `applied: false` and keeps the old source +#: (verified against the binary), so a comment could never be released. And +#: the app is in the property's *name* because the engine refuses to retype +#: a property (`enum(a)` to `enum(b)` fails with "changing property type ... +#: not supported"), while adding one nullable property and dropping another +#: in a single apply is a migration it performs without flags β€” both +#: verified against the binary. +COCO_MANAGED_PREFIX = "coco_managed_by_" + +_IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def ownership_property(app_name: str) -> str: + """The property marking a block as owned by the app named `app_name`. + + An app name that is an identifier is used as is. Any other name is made + one, with a digest of the original appended so that two names which + sanitize alike (`my-app`, `my_app`) still get distinct properties. + """ + if _IDENTIFIER_RE.fullmatch(app_name): + return f"{COCO_MANAGED_PREFIX}{app_name}" + sanitized = re.sub(r"[^A-Za-z0-9_]", "_", app_name) + digest = hashlib.sha256(app_name.encode()).hexdigest()[:8] + return f"{COCO_MANAGED_PREFIX}{sanitized}_{digest}" + + +#: Scalar type keywords `.pg`/`.gq` accept. `DateTime` must precede `Date` in +#: the alternation below so a caller can't rely on `Date` matching first and +#: leaving `Time...` as unconsumed trailing text (Python's `re` backtracks +#: across alternatives, so match order doesn't actually matter for +#: correctness β€” but keeping the longer name first avoids relying on that). +_PG_BASE_TYPES = ( + "String", + "Bool", + "I32", + "I64", + "U32", + "U64", + "F32", + "F64", + "DateTime", + "Date", + "Blob", +) +_PG_BASE_ALT = "|".join(_PG_BASE_TYPES) +_PG_IDENT = r"[A-Za-z_][A-Za-z0-9_]*" + +#: `Vector(N)`, `[T]` for scalar `T`, and `enum(a, b, ...)`, plus an optional +#: trailing `?` for nullability β€” the full `.pg`/`.gq` type-expression +#: grammar. Anything else is rejected. +_PG_TYPE_RE = re.compile( + rf"^(?:{_PG_BASE_ALT}" + rf"|Vector\([1-9][0-9]*\)" + rf"|\[(?:{_PG_BASE_ALT})\]" + rf"|enum\({_PG_IDENT}(?:,\s*{_PG_IDENT})*\)" + rf")\??$" +) + + +def validate_identifier(name: str, what: str) -> None: + """Reject anything that could break out of a `.pg` or `.gq` fragment. + + Omnigraph identifiers have no quoting form we can rely on, so the only + safe policy is to refuse non-conforming names outright. + """ + if not _IDENTIFIER_RE.fullmatch(name): + raise ValueError(f"Invalid Omnigraph {what}: {name!r}") + + +def validate_pg_type(pg_type: str) -> None: + """Reject anything that isn't a legal Omnigraph type expression. + + A mutation's `pg_type` is never bound as a parameter β€” it is spliced + directly into the query *signature* (`$p_email: `), so this is + as much an injection boundary as `validate_identifier` and needs the + same strictness. + """ + if not _PG_TYPE_RE.fullmatch(pg_type): + raise ValueError(f"Invalid Omnigraph property type: {pg_type!r}") + + +def render_property(name: str, pg_type: str, *, is_key: bool) -> str: + validate_identifier(name, "property name") + validate_pg_type(pg_type) + return f"{name}: {pg_type} @key" if is_key else f"{name}: {pg_type}" + + +#: Property names Omnigraph itself reserves on a node type. `id` is the +#: engine's own identity column, materialized from the `@key` property: +#: declaring one alongside it fails at `init`/`schema apply` with "physical +#: schema for 'node:X' must contain exactly one top-level `id` field; found +#: 2" (verified against the binary), and using `id` *as* the key fails even +#: more confusingly with "@key must reference declared properties". +_RESERVED_NODE_PROPERTIES = ("id",) + +#: The same for an edge type, plus the two endpoint columns. `src`/`dst` +#: collide with the engine's stored endpoint references ("Duplicate field +#: name \"src\" in schema"); `from`/`to` parse fine in `.pg` but are the +#: literal field names an edge `insert` assigns its endpoints to, so a +#: property by either name could never be written (`build_edge_insert` +#: refuses it) β€” rejecting it here means the failure lands at schema +#: declaration rather than at the first write. +_RESERVED_EDGE_PROPERTIES = ("id", "src", "dst", "from", "to") + + +def _check_not_reserved( + properties: Sequence[tuple[str, str]], type_name: str, reserved: Sequence[str] +) -> None: + synthetic = sorted(name for name, _ in properties if name.startswith(COCO_PREFIX)) + if synthetic: + raise ValueError( + f"{synthetic!r} {'is' if len(synthetic) == 1 else 'are'} reserved by " + f"the CocoIndex connector (every name starting with {COCO_PREFIX!r} " + f"is) and cannot be declared on {type_name!r}" + ) + clashes = sorted({name for name, _ in properties} & set(reserved)) + if clashes: + raise ValueError( + f"{clashes!r} {'is' if len(clashes) == 1 else 'are'} reserved by " + f"Omnigraph and cannot be declared on {type_name!r}; rename the " + f"field (e.g. {clashes[0]!r} -> {type_name.lower()}_{clashes[0]})" + ) + + +def render_node_type( + type_name: str, + properties: Sequence[tuple[str, str]], + key: tuple[str, ...], + *, + owner: str, +) -> str: + """Render a node type's block, marked as owned by the app `owner`.""" + validate_identifier(type_name, "node type") + _check_not_reserved(properties, type_name, _RESERVED_NODE_PROPERTIES) + if len(key) > 1: + raise ValueError( + f"Omnigraph node types support exactly one @key property, but " + f"{type_name} declares {key!r}. The engine rejects the schema " + f'outright ("node type {type_name} has multiple @key constraints; ' + f'only one is supported"), so a composite key cannot be expressed ' + f"at all β€” derive a single key field instead." + ) + by_name = dict(properties) + for k in key: + if k not in by_name: + raise ValueError(f"key property {k!r} is not declared on {type_name}") + if by_name[k].endswith("?"): + raise ValueError(f"key property {k!r} must not be nullable on {type_name}") + lines = [ + f" {render_property(name, pg_type, is_key=name in key)}" + for name, pg_type in properties + ] + lines.extend(_render_synthetic_properties(owner)) + body = "\n".join(lines) + return f"node {type_name} {{\n{body}\n}}" + + +def _render_synthetic_properties(owner: str) -> list[str]: + return [ + f" {render_property(COCO_KEY, 'String', is_key=False)}", + f" {render_property(ownership_property(owner), 'Bool?', is_key=False)}", + ] + + +def render_edge_type( + type_name: str, + from_type: str, + to_type: str, + properties: Sequence[tuple[str, str]], + *, + owner: str, +) -> str: + """Render an edge type's block, marked as owned by the app `owner`.""" + validate_identifier(type_name, "edge type") + validate_identifier(from_type, "node type") + validate_identifier(to_type, "node type") + _check_not_reserved(properties, type_name, _RESERVED_EDGE_PROPERTIES) + lines = [ + f" {render_property(name, pg_type, is_key=False)}" + for name, pg_type in properties + ] + lines.extend(_render_synthetic_properties(owner)) + body = "\n".join(lines) + return f"edge {type_name}: {from_type} -> {to_type} {{\n{body}\n}}" + + +#: A `//` line comment. Blanked out β€” replaced by spaces of the same length, +#: so every offset stays valid β€” before the schema is scanned: `schema show` +#: returns the source verbatim, comments included, and a `}` or a `node X {` +#: inside one is text, not structure. `//` is the only comment syntax the +#: engine accepts (`#` and `--` are parse errors, verified against the +#: binary). +_COMMENT_RE = re.compile(r"//[^\n]*") + +#: The head of a type declaration: `node NAME` or `edge NAME`, anywhere in +#: the source β€” not only at the start of a line. A hand-written schema may +#: indent its blocks or put two declarations on one line, and `schema show` +#: preserves both, so a merger that only recognised `node X {` at a line +#: start appended a second block and the engine refused the result +#: ("duplicate node name"). The KIND is captured, not just the name: a `.pg` +#: may legally hold `node Link` and `edge Link` side by side (the engine +#: accepts it), and matching on the name alone made an edge's fragment +#: overwrite the node's block. +_TYPE_HEAD_RE = re.compile(r"\b(node|edge)\s+(\w+)\b") + +#: The rest of an edge head, right after `edge NAME`: `: FROM -> TO`. +_EDGE_ENDPOINTS_RE = re.compile(r"\s*:\s*(\w+)\s*->\s*(\w+)") + +#: Optional whitespace and then the opening brace of a block body. +_BODY_OPEN_RE = re.compile(r"\s*\{") + + +def _blank_comments(pg: str) -> str: + """`pg` with every `//` comment replaced by spaces of the same length, so + offsets into it are offsets into the original.""" + return _COMMENT_RE.sub(lambda m: " " * len(m.group(0)), pg) + + +class _TypeBlock(NamedTuple): + kind: str + name: str + #: Span into the ORIGINAL source, comments and formatting included. + start: int + end: int + #: Edge endpoints; `None` for a node. + from_type: str | None + to_type: str | None + + +def _scan_type_blocks(existing_pg: str) -> list[_TypeBlock]: + """Every top-level type block in `existing_pg`, in source order. + + Scans a comment-blanked copy of the same length, so spans line up with + the original. A block starts at `node NAME` / `edge NAME` and ends at + the matching `}` of its braced body, or β€” for the brace-less `edge + NAME: FROM -> TO` form, a property-less edge the engine both accepts + and reproduces via `schema show` β€” right after the `TO` endpoint. This + connector's own builders always add `coco_key`, so they never emit the + brace-less form, but the schema being searched can legitimately contain + it, from another tool or a `managed_by=user` type. Scanning resumes + after each block's end, so a property named `node` or `edge` inside a + body is never mistaken for a declaration. + + Raises rather than guessing on a schema it can't edit safely β€” a node + head with no body, an edge head with no endpoints, or braces that never + balance: either way the splice point would land mid-declaration and + produce corrupt `.pg` that goes straight to `schema apply`. + """ + text = _blank_comments(existing_pg) + blocks: list[_TypeBlock] = [] + pos = 0 + while (head := _TYPE_HEAD_RE.search(text, pos)) is not None: + kind, name = head.group(1), head.group(2) + body_start = head.end() + from_type = to_type = None + if kind == "edge": + endpoints = _EDGE_ENDPOINTS_RE.match(text, body_start) + if endpoints is None: + raise ValueError( + f"Omnigraph schema declares edge {name!r} without " + f"`: FROM -> TO` endpoints; refusing to edit it" + ) + from_type, to_type = endpoints.group(1), endpoints.group(2) + body_start = endpoints.end() + body = _BODY_OPEN_RE.match(text, body_start) + if body is None: + if kind == "node": + raise ValueError( + f"Omnigraph schema declares node {name!r} without a " + f"`{{ ... }}` body; refusing to edit it" + ) + end = body_start + else: + depth = 0 + end = -1 + for i in range(body.end() - 1, len(text)): + if text[i] == "{": + depth += 1 + elif text[i] == "}": + depth -= 1 + if depth == 0: + end = i + 1 + break + if end == -1: + raise ValueError( + f"Omnigraph schema has an unterminated {kind} {name!r} block " + f"(unbalanced braces); refusing to edit it" + ) + blocks.append(_TypeBlock(kind, name, head.start(), end, from_type, to_type)) + pos = end + return blocks + + +class IncidentEdge(NamedTuple): + """An edge type that has a given node type at one of its ends.""" + + edge_type: str + role: str # "from" | "to" + + +def incident_edge_patterns(existing_pg: str, node_type_name: str) -> list[IncidentEdge]: + """Every (edge type, end) in `existing_pg` at which `node_type_name` + sits, in source order. A self-referencing edge type contributes both + ends.""" + patterns: list[IncidentEdge] = [] + for block in _scan_type_blocks(existing_pg): + if block.kind != "edge": + continue + if block.from_type == node_type_name: + patterns.append(IncidentEdge(block.name, "from")) + if block.to_type == node_type_name: + patterns.append(IncidentEdge(block.name, "to")) + return patterns + + +def edge_types_referencing(existing_pg: str, node_type_name: str) -> list[str]: + """Names of edge types in `existing_pg` that use `node_type_name` as an + endpoint. Used to refuse removing a node type out from under one.""" + return list( + dict.fromkeys( + p.edge_type for p in incident_edge_patterns(existing_pg, node_type_name) + ) + ) + + +def _find_type_block( + existing_pg: str, kind: str, type_name: str +) -> tuple[int, int] | None: + """Locate the `(start, end)` span of the `kind` type named `type_name` + in `existing_pg`, or `None` if it isn't present. + + `kind` is `"node"` or `"edge"` and is part of the match, not a hint: a + node and an edge may share a name in a valid `.pg`, and editing the + wrong one silently destroys a type. See `_scan_type_blocks` for what a + block's extent is. + + Raises on two blocks with the same kind and name: only the first would + ever be rewritten, leaving the second as a stale duplicate. + """ + if kind not in ("node", "edge"): + raise ValueError(f"type kind must be 'node' or 'edge', got {kind!r}") + matching = [ + block + for block in _scan_type_blocks(existing_pg) + if (block.kind, block.name) == (kind, type_name) + ] + if not matching: + return None + if len(matching) > 1: + raise ValueError( + f"Omnigraph schema declares {kind} {type_name!r} " + f"{len(matching)} times; refusing to edit an ambiguous schema" + ) + return matching[0].start, matching[0].end + + +#: An ownership declaration together with the whitespace that is its own: +#: the line break and indentation before it, or the run of spaces before it +#: in a one-line block β€” never more, or a release would eat the (blanked) +#: comment ending the previous line. Group 1 is the property's name. +_COCO_MANAGED_DECL_RE = re.compile( + rf"(?:\n[ \t]*|[ \t]+)?\b({COCO_MANAGED_PREFIX}\w+)\s*:\s*Bool\??" +) + + +def ownership_marker(existing_pg: str, kind: str, type_name: str) -> str | None: + """The ownership property the block for `kind` `type_name` declares + (`coco_managed_by_`, see `ownership_property`), or `None` for a + block no app of this connector currently owns β€” a user's, or another + tool's. + + Declarations only: the block is searched with its comments blanked, so + a user-owned block whose comment merely mentions the property is not + mistaken for the connector's (which once let an app drop delete it). + """ + span = _find_type_block(existing_pg, kind, type_name) + if span is None: + return None + start, end = span + m = _COCO_MANAGED_DECL_RE.search(_blank_comments(existing_pg)[start:end]) + return m.group(1) if m is not None else None + + +def release_ownership(existing_pg: str, kind: str, type_name: str, marker: str) -> str: + """Remove the declaration of the ownership property `marker` from the + block for `kind` `type_name`, leaving everything else in it β€” comments + included β€” and the rest of the schema untouched. A no-op if the block + is absent or does not declare `marker`; another app's marker on the + block is that app's to release. + + This is the one schema write a `managed_by=user` declaration causes: a + type the connector created and the app then handed to the user still + carried the ownership property, and a later drop of a node type it + referenced read that as current ownership and removed the user's edge + type with it. Dropping a nullable, never-written property is a + migration the engine applies without flags (verified). + """ + span = _find_type_block(existing_pg, kind, type_name) + if span is None: + return existing_pg + start, end = span + block = existing_pg[start:end] + blanked = _blank_comments(existing_pg)[start:end] + # Cut the declarations found on the comment-blanked view out of the + # original text; offsets line up because blanking preserves length. + for m in reversed(list(_COCO_MANAGED_DECL_RE.finditer(blanked))): + if m.group(1) == marker: + block = block[: m.start()] + block[m.end() :] + return existing_pg[:start] + block + existing_pg[end:] + + +def merge_type_into_schema( + existing_pg: str, kind: str, type_name: str, fragment: str +) -> str: + """Merge one type's rendered `.pg` fragment into the full schema source + read back from `schema show`. `kind` is `"node"` or `"edge"` β€” see + `_find_type_block` for why matching on the name alone is unsafe. + + Omnigraph's schema is applied whole-graph, not per type: `schema apply` + and `init` both treat their input as the *complete* desired schema and + silently drop any type omitted from it β€” verified against the engine: + applying a single new type's fragment alone dropped every other + existing type from the graph. Since each type reconciles independently + and only ever renders its own single-type fragment, applying that + directly would wipe the rest of the graph's schema on every type-level + sync. This does the merge textually instead: replace the existing + block for `type_name` if the current schema already has one, otherwise + append it. + + Pure and unit-tested directly β€” deliberately not exercised only via a + live store, since this is fiddly text handling. + + Read-merge-write is not atomic: two components applying schema + concurrently can race this, and a lost update would silently drop + whichever type the loser was adding or changing. Not addressed here. + Schema actions are rare (only on type creation or change), so the + window is small; closing it fully would need a much larger change β€” + one target owning the whole graph's schema β€” which is out of scope. + """ + span = _find_type_block(existing_pg, kind, type_name) + if span is not None: + start, end = span + return existing_pg[:start] + fragment + existing_pg[end:] + # Not present in the current schema -- append it. + sep = "\n\n" if existing_pg.strip() else "" + return existing_pg.rstrip("\n") + sep + fragment + "\n" + + +def remove_type_from_schema(existing_pg: str, kind: str, type_name: str) -> str: + """Remove the `kind` type named `type_name` from the schema source + entirely, leaving every other type untouched. A no-op if it isn't + present. + + This is the first half of the two-`schema apply` rebuild a `@key` (or + edge endpoint) change needs: the engine flatly rejects an in-place key + change in one step ("removing property constraints ... not supported + in schema migration v1", verified against the binary, regardless of + `--allow-data-loss` β€” this isn't a soft/hard-drop distinction, it's an + unsupported migration step outright) but accepts dropping the type and + then re-adding it with its new definition as two separate calls + (verified live, and a plain soft drop is enough β€” re-adding a type of + the same name right after succeeds cleanly, no `--allow-data-loss` + needed). + """ + span = _find_type_block(existing_pg, kind, type_name) + if span is None: + return existing_pg + start, end = span + before, after = existing_pg[:start], existing_pg[end:] + # Consume one adjoining blank-line separator so removing a middle block + # doesn't leave the two neighbors glued together by a doubled one. + if after.startswith("\n\n"): + after = after[2:] + elif before.endswith("\n\n"): + before = before[:-2] + return before + after + + +class PropertyValue(NamedTuple): + """One bound value: its property name, its GQ type, and the value itself. + + The type travels with the value because a parameterized query must declare + it in the signature (`query m($p_email: String)`). + """ + + name: str + pg_type: str + value: object + + +class Bind(NamedTuple): + """One parameter slot of a `Statement`: its label, `.pg` type and value. + + The label (`p_email`, `e_from`, `p_coco_key`, ...) is unique within its + statement; `render_query` turns it into a query-wide parameter name. + """ + + label: str + pg_type: str + value: object + + +class Statement(NamedTuple): + """One GQ statement with positional parameter slots. + + `body` holds a `$?` slot wherever a bound value goes, and `binds` lists + each slot's label, type and value in slot order. Parameter names are + allocated only when a query is rendered β€” `s{i}_{label}` for the i-th + statement β€” so any number of statements combine into one query with no + renaming, and a value never enters the statement text at all. + """ + + body: str + binds: tuple[Bind, ...] + + +class Query(NamedTuple): + """A rendered `query m(...) { ... }` source plus its bound params β€” the + form the transport sends, whether the block mutates or reads. Values + are never interpolated into `expr`, so a property value containing + quotes or braces cannot alter the query. + """ + + expr: str + params: dict[str, object] + + +#: Marks a parameter position in a `Statement.body`. `$` never appears in a +#: statement otherwise (identifiers are validated to `[A-Za-z0-9_]`), so the +#: slot can be found by plain splitting. +_SLOT = "$?" + + +def render_query(statements: Sequence[Statement]) -> Query: + """Render statements into ONE `query m(...) { ... }` β€” one CLI + invocation, one commit. N separate `mutate` calls would be N commits + and no atomicity. + + Each statement's binds become parameters `s{i}_{label}`, declared in the + signature in slot order and substituted for the statement's `$?` slots + in the same order. + """ + if not statements: + raise ValueError("render_query requires at least one statement") + signature: list[str] = [] + params: dict[str, object] = {} + bodies: list[str] = [] + for i, statement in enumerate(statements): + parts = statement.body.split(_SLOT) + if len(parts) != len(statement.binds) + 1: + raise ValueError( + f"statement has {len(parts) - 1} slots but " + f"{len(statement.binds)} bind{'s' if len(statement.binds) != 1 else ''}: " + f"{statement.body!r}" + ) + rendered = [parts[0]] + for bind, rest in zip(statement.binds, parts[1:], strict=True): + name = f"s{i}_{bind.label}" + if name in params: + raise ValueError( + f"Duplicate parameter label {bind.label!r} in {statement.body!r}" + ) + signature.append(f"${name}: {bind.pg_type}") + params[name] = bind.value + rendered.append(f"${name}{rest}") + bodies.append("".join(rendered)) + return Query(f"query m({', '.join(signature)}) {{ {' '.join(bodies)} }}", params) + + +def _bind(props: Sequence[PropertyValue], prefix: str) -> tuple[list[Bind], list[str]]: + """Validate `props` and turn them into binds plus `name: $?` assignments.""" + binds: list[Bind] = [] + assigns: list[str] = [] + seen: set[str] = set() + for prop in props: + validate_identifier(prop.name, "property name") + validate_pg_type(prop.pg_type) + if prop.name.startswith(COCO_PREFIX): + raise ValueError( + f"{prop.name!r} is reserved by the CocoIndex connector and cannot be " + f"supplied as a property value" + ) + if prop.name in seen: + raise ValueError(f"Duplicate property name: {prop.name!r}") + seen.add(prop.name) + binds.append(Bind(f"{prefix}_{prop.name}", prop.pg_type, prop.value)) + assigns.append(f"{prop.name}: {_SLOT}") + return binds, assigns + + +def _coco_key_bind(coco_key: str) -> Bind: + return Bind(f"p_{COCO_KEY}", "String", coco_key) + + +def build_node_upsert( + type_name: str, props: Sequence[PropertyValue], coco_key: str +) -> Statement: + """Keyed node `insert` is an upsert by the derived key tuple.""" + validate_identifier(type_name, "node type") + binds, assigns = _bind(props, "p") + assigns.append(f"{COCO_KEY}: {_SLOT}") + body = f"insert {type_name} {{ {', '.join(assigns)} }}" + return Statement(body, (*binds, _coco_key_bind(coco_key))) + + +def build_endpoint_stub( + type_name: str, key_props: Sequence[PropertyValue], coco_key: str +) -> Statement: + """Key-only upsert, so an edge can reference a node its owning component + has not written yet. The owner's own upsert later fills in the rest.""" + return build_node_upsert(type_name, key_props, coco_key) + + +def _delete_by_coco_key(type_name: str, coco_key: str) -> Statement: + body = f"delete {type_name} where {COCO_KEY} = {_SLOT}" + return Statement(body, (_coco_key_bind(coco_key),)) + + +def build_node_delete(type_name: str, coco_key: str) -> Statement: + validate_identifier(type_name, "node type") + return _delete_by_coco_key(type_name, coco_key) + + +def build_edge_delete(type_name: str, coco_key: str) -> Statement: + validate_identifier(type_name, "edge type") + return _delete_by_coco_key(type_name, coco_key) + + +def build_incident_edges_query( + node_type: str, coco_key: str, edge: IncidentEdge +) -> Statement: + """Read the `coco_key` of every `edge.edge_type` edge at the node of + `node_type` whose `coco_key` is `coco_key`, on the end `edge.role`. + + A read matches an edge type by its traversal spelling, which starts + with a lowercase letter (`worksAt` for `edge WorksAt`; the engine looks + it up case-insensitively, verified against the binary), and binds the + hop as `$e:worksAt` so the edge's own properties can be read back. The + node is found by `coco_key`, which every type this connector writes to + declares, rather than by its `@key`, so the query is the same for every + key type. + """ + validate_identifier(node_type, "node type") + validate_identifier(edge.edge_type, "edge type") + if edge.role not in ("from", "to"): + raise ValueError(f"edge role must be 'from' or 'to', got {edge.role!r}") + spelling = edge.edge_type[0].lower() + edge.edge_type[1:] + hop = f"$n $e:{spelling} $o" if edge.role == "from" else f"$o $e:{spelling} $n" + body = ( + f"match {{ $n: {node_type} {{ {COCO_KEY}: {_SLOT} }} {hop} }} " + f"return {{ $e.{COCO_KEY} as edge_key }}" + ) + return Statement(body, (_coco_key_bind(coco_key),)) + + +def build_edge_insert( + type_name: str, + from_ref: PropertyValue, + to_ref: PropertyValue, + props: Sequence[PropertyValue], + coco_key: str, +) -> Statement: + """Edge `insert` always creates a new edge β€” Omnigraph never deduplicates + and provides no settable id. Idempotence comes entirely from the connector + refusing to re-insert an edge whose `coco_key` it already tracks, and from + deleting by `coco_key` before re-inserting a changed one. + + `from_ref.name`/`to_ref.name` are ignored β€” the assignment is always + literally `from`/`to`; only `.pg_type` and `.value` address the endpoint. + """ + validate_identifier(type_name, "edge type") + validate_pg_type(from_ref.pg_type) + validate_pg_type(to_ref.pg_type) + for prop in props: + if prop.name in ("from", "to"): + raise ValueError( + f"{prop.name!r} is reserved for the edge endpoint and cannot be " + f"supplied as a property value" + ) + binds = [ + Bind("e_from", from_ref.pg_type, from_ref.value), + Bind("e_to", to_ref.pg_type, to_ref.value), + ] + assigns = [f"from: {_SLOT}", f"to: {_SLOT}"] + prop_binds, prop_assigns = _bind(props, "p") + binds += prop_binds + assigns += prop_assigns + assigns.append(f"{COCO_KEY}: {_SLOT}") + body = f"insert {type_name} {{ {', '.join(assigns)} }}" + return Statement(body, (*binds, _coco_key_bind(coco_key))) diff --git a/python/cocoindex/connectors/omnigraph/_target.py b/python/cocoindex/connectors/omnigraph/_target.py new file mode 100644 index 000000000..91093d7ed --- /dev/null +++ b/python/cocoindex/connectors/omnigraph/_target.py @@ -0,0 +1,2506 @@ +"""Omnigraph target connector: schemas, handlers, sink, and user-facing targets.""" + +from __future__ import annotations + +import contextlib +import dataclasses +import datetime +import functools +import hashlib +import re +import types +import typing +import uuid +from collections.abc import AsyncIterator, Callable, Sequence +from typing import Any, Generic, Literal, NamedTuple + +import cocoindex as coco +import msgspec +from cocoindex._internal.component_ctx import current_app_name +from cocoindex._internal.context_keys import ContextProvider +from cocoindex._internal.datatype import TypeChecker +from cocoindex._internal.memo_fingerprint import canonical_module_name +from cocoindex.connectorkits import statediff +from cocoindex.connectorkits.fingerprint import fingerprint_object +from cocoindex.connectorkits.target import ManagedBy +from cocoindex.connectors.omnigraph._client import ( + ConnectionFactory, + OmnigraphCliError, + _CliClient, +) +from cocoindex.connectors.omnigraph._gq import ( + PropertyValue, + Query, + Statement, + _find_type_block, + build_edge_delete, + build_edge_insert, + build_endpoint_stub, + build_incident_edges_query, + build_node_delete, + build_node_upsert, + edge_types_referencing, + incident_edge_patterns, + merge_type_into_schema, + ownership_marker, + ownership_property, + release_ownership, + remove_type_from_schema, + render_edge_type, + render_node_type, + render_query, + validate_identifier, + validate_pg_type, +) +from typing_extensions import TypeVar + + +def derive_coco_key(parts: object) -> str: + """Stable value for the synthetic `coco_key` property. + + Omnigraph gives edges no settable id, its auto-assigned ULID is not + filterable, and `where` accepts exactly one equality predicate β€” so the + connector declares `coco_key` on every type it manages and addresses + entities through it. This computes what goes in there. + + Fingerprints the key tuple rather than concatenating into a string: + `f"{a}_{b}"` collides whenever parts contain underscores, silently merging + two distinct entities onto one target state. + """ + return fingerprint_object(parts).hex() + + +ValueEncoder = Callable[[Any], Any] + + +class OmnigraphType(NamedTuple): + """Annotation overriding the default Python-to-`.pg` type mapping. + + Use as ``Annotated[int, OmnigraphType("I32")]`` when the default (``I64``) + is wider than the schema should declare. + """ + + pg_type: str + + +class PropertyDef(NamedTuple): + name: str + pg_type: str + #: Applied to non-`None` values before they are sent. `None` means the + #: built-in encoding for `pg_type` (ISO format for `Date`/`DateTime`, + #: mapped over a list of either), so a hand-built definition behaves + #: like one `NodeSchema.from_class` produced. + encoder: ValueEncoder | None = None + + +_SCALARS: dict[Any, str] = { + str: "String", + bool: "Bool", + int: "I64", + float: "F64", + datetime.date: "Date", + datetime.datetime: "DateTime", + # `bytes` deliberately has no mapping: Omnigraph's `Blob` is an external + # URI reference the engine *fetches* (a `file://` value), not inline + # bytes β€” a Python `bytes` value can never be a `Blob`, so it falls + # through to the generic "no mapping" TypeError below instead of being + # silently declared as one. +} + + +def _isoformat(value: Any) -> str: + return value.isoformat() # type: ignore[no-any-return] + + +#: `PropertyDef.encoder` for the scalar `.pg` types whose Python value isn't +#: `json.dumps`-safe as-is. Verified against the engine: `Date` accepts +#: `"2026-01-01"` and `DateTime` accepts ISO with or without a trailing `Z`, +#: so `.isoformat()` suffices for both. +_ENCODERS: dict[str, ValueEncoder] = { + "Date": _isoformat, + "DateTime": _isoformat, +} + +#: Encoders a key property of each `.pg` type may carry: the built-in one +#: for that type, and the `isoformat` method a person writing a schema by +#: hand would reach for, which is the same encoding. Keyed by type, not one +#: flat list: `datetime.date.isoformat` on a DateTime key drops the time +#: from every value, so two instants on one date were stored as a single +#: midnight node while tracked as two, and removing one declaration deleted +#: the node the other still declared. Anything else changes what the graph +#: keys on and is refused at `node_target`. +_KEY_ENCODERS: dict[str, tuple[ValueEncoder, ...]] = { + "Date": (_ENCODERS["Date"], datetime.date.isoformat), + "DateTime": (_ENCODERS["DateTime"], datetime.datetime.isoformat), +} + + +@functools.cache +def _list_encoder(element: ValueEncoder) -> ValueEncoder: + # Cached so `_encoder_for` hands back one object per element encoder: + # `_encoder_identity` tells the built-in encoding of a type from a custom + # one by identity. + return lambda values: [element(v) for v in values] + + +def _encoder_for(pg_type: str) -> ValueEncoder | None: + """Look up by the *resolved* pg_type string (stripped of `?`), not the + Python annotation β€” this way it applies the same whether the field came + from a plain `datetime.date` or an `Annotated[..., OmnigraphType(...)]` + override that happens to resolve to `Date`/`DateTime`. + + A list type `[T]` gets its element's encoder mapped over the list: the + schema accepts `[Date]`, and without this every write of such a value + died in `json.dumps` while the type itself had been declared fine. + """ + base = pg_type.rstrip("?") + if base.startswith("[") and base.endswith("]"): + element = _ENCODERS.get(base[1:-1]) + return _list_encoder(element) if element is not None else None + return _ENCODERS.get(base) + + +def _pg_type_for(annotation: Any) -> str: + """Map a Python annotation to a `.pg` type, honouring OmnigraphType. + + ``bool`` is checked before ``int`` because ``bool`` is a subclass of + ``int`` and would otherwise map to ``I64``. + + ``OmnigraphType`` lets the app author write an arbitrary `pg_type` + string, and that string is spliced directly into generated query text + downstream β€” so every value this function returns, on every path + (including the OmnigraphType override), goes through + `_gq.validate_pg_type` before it reaches the caller. + """ + origin = typing.get_origin(annotation) + if origin is typing.Annotated: + base, *meta = typing.get_args(annotation) + for m in meta: + if isinstance(m, OmnigraphType): + validate_pg_type(m.pg_type) + return m.pg_type + return _pg_type_for(base) + if origin in (typing.Union, types.UnionType): + args = [a for a in typing.get_args(annotation) if a is not type(None)] + if len(args) == 1: + inner = _pg_type_for(args[0]) + pg_type = inner if inner.endswith("?") else f"{inner}?" + validate_pg_type(pg_type) + return pg_type + raise TypeError(f"no Omnigraph type mapping for union {annotation!r}") + if origin in (list, Sequence): + (item,) = typing.get_args(annotation) + inner = _pg_type_for(item) + # Verified against the engine (2026-08-26): a list element must be a + # plain scalar. `[String?]` and `[I64?]` fail with "expected + # core_type"; `[[String]]` fails with "expected base_type". Emitting + # them would produce a `.pg` the engine rejects at init time, so fail + # here with a message that names the offending annotation. + if inner.endswith("?") or inner.startswith("["): + raise TypeError( + f"Omnigraph list elements must be non-nullable scalars; " + f"{annotation!r} maps to [{inner}], which the engine rejects" + ) + pg_type = f"[{inner}]" + validate_pg_type(pg_type) + return pg_type + # bool is a subclass of int, so it must be checked before int; dict + # lookup on the exact annotation naturally does this since `_SCALARS` + # keys on the type object itself, not an isinstance chain. + if annotation in _SCALARS: + pg_type = _SCALARS[annotation] + validate_pg_type(pg_type) + return pg_type + raise TypeError(f"no Omnigraph type mapping for {annotation!r}") + + +def _properties_from_class(target: type) -> dict[str, PropertyDef]: + """Introspect a dataclass into `PropertyDef`s, one per field, in order.""" + hints = typing.get_type_hints(target, include_extras=True) + properties: dict[str, PropertyDef] = {} + for field in dataclasses.fields(target): + validate_identifier(field.name, "property name") + # `id` is the one name Omnigraph reserves whether this schema ends + # up describing a node or an edge, and it's by far the likeliest + # collision in a real dataclass β€” so it's worth catching here, + # where the class name is in scope. The name-set that depends on + # which kind of type this becomes (an edge's `src`/`dst`/`from`/ + # `to`) is enforced by `_gq.render_edge_type` instead. + if field.name == "id": + raise ValueError( + f"{target.__name__}.id: 'id' is reserved by Omnigraph for " + f"the engine's own identity column and cannot be declared " + f"as a property. Rename the field (e.g. " + f"'{target.__name__.lower()}_id')." + ) + pg_type = _pg_type_for(hints[field.name]) + properties[field.name] = PropertyDef(field.name, pg_type, _encoder_for(pg_type)) + return properties + + +@dataclasses.dataclass(frozen=True) +class NodeSchema: + """A node type's properties and the one of them that is its `@key`.""" + + properties: dict[str, PropertyDef] + key: tuple[str, ...] + + @classmethod + async def from_class(cls, target: type, *, key: str | Sequence[str]) -> NodeSchema: + key_tuple = (key,) if isinstance(key, str) else tuple(key) + if not key_tuple: + raise ValueError( + f"{target.__name__} needs at least one key property: Omnigraph " + f"upserts by key on a keyed node `insert`, but an unkeyed " + f"`insert` is a strict insert β€” every re-run would duplicate " + f"every node" + ) + if len(key_tuple) > 1: + raise ValueError( + f"{target.__name__} declares a composite key {key_tuple!r}, but " + f"Omnigraph node types support exactly one @key property β€” the " + f"engine rejects a two-@key schema outright. Derive a single " + f"key field (e.g. a generated id) and key on that." + ) + properties = _properties_from_class(target) + for k in key_tuple: + if k not in properties: + raise ValueError( + f"key property {k!r} is not a field of {target.__name__}" + ) + if properties[k].pg_type.endswith("?"): + raise ValueError( + f"key property {k!r} of {target.__name__} must not be nullable" + ) + return cls(properties=properties, key=key_tuple) + + def render(self, type_name: str, *, owner: str) -> str: + return render_node_type( + type_name, + [(p.name, p.pg_type) for p in self.properties.values()], + key=self.key, + owner=owner, + ) + + +@dataclasses.dataclass(frozen=True) +class EdgeSchema: + """An edge type's own properties. + + An edge has no key of its own β€” its identity is always `(from_id, + to_id)` β€” so unlike `NodeSchema` there is none to declare. + """ + + properties: dict[str, PropertyDef] + + @classmethod + async def from_class(cls, target: type) -> EdgeSchema: + return cls(properties=_properties_from_class(target)) + + +# --------------------------------------------------------------------------- +# Container handlers: a node type or edge type, and how its schema evolves +# --------------------------------------------------------------------------- + + +class _TypeKey(NamedTuple): + """A managed type's tracking identity. + + `type_kind` is part of it because a `.pg` may legally declare `node Link` + and `edge Link` side by side β€” they are different types, and the sink has + to know which block it is editing. + + Deliberately does NOT carry the branch (or the store URI): both live on + the `ConnectionFactory`, which is resolved from `db_key` at action time, + never captured at declare time β€” a delete action runs in a process where + the declaring code never executed. Same split the sibling connectors + make: neo4j's `_TableKey` is `(db_key, table_name)` and keeps the + database name on the factory. Point two apps at different branches by + giving them different `ContextKey`s. + """ + + db_key: str + type_kind: str # "node" | "edge" + type_name: str + + +_TYPE_KEY_CHECKER: TypeChecker[tuple[str, str, str]] = TypeChecker(tuple[str, str, str]) + + +@dataclasses.dataclass(frozen=True) +class _TypeSpec: + schema: NodeSchema | EdgeSchema | None + key: tuple[str, ...] + from_type: str | None + to_type: str | None + managed_by: ManagedBy + #: The app declaring the type (`AppConfig.name`). Rendered into the + #: block as its ownership property, and persisted in the tracking + #: record so a removal β€” which runs from tracking alone β€” knows whose + #: blocks it may take along. + owner: str + # Populated only for edge specs: the *endpoint* node types' own key + # definitions. Needed by the sink to build endpoint stubs and the edge's + # `from`/`to` refs β€” `from_type`/`to_type` name the endpoint types, but + # nothing else here carries their key definitions, since `schema` above is + # the edge's own (possibly absent) property schema, not the endpoints'. + from_key_property: PropertyDef | None = None + to_key_property: PropertyDef | None = None + + +class _TypeMainRecord(msgspec.Struct, frozen=True, array_like=True): + """A managed type's identity β€” a change here forces a full rebuild. + + Deliberately does NOT carry neo4j's `has_schema` flag. + `_EdgeTypeHandler._render` emits a byte-identical fragment for + `schema=None` and for an empty schema, so tracking that distinction would + promote `None` ↔ `{}` β€” a difference the graph cannot observe β€” into a + destructive drop-and-recreate. + """ + + key: tuple[str, ...] + from_type: str | None + to_type: str | None + #: Constant for every record in one app's store β€” the store is the + #: app's β€” so it never differs between records and never forces a + #: rebuild; it is here because a removal has nothing else to read it + #: from. + owner: str + + +class _PropertyRecord(msgspec.Struct, frozen=True, array_like=True): + """One property as the type's tracking record remembers it. + + `pg_type` is the bare `.pg` type string: nullability is already spelled + in it as a trailing `?`, so neo4j's `(type, nullable)` split would only + restate the same bit. + + `encoder` is `_encoder_identity` of the property's encoder. The rendered + `.pg` never changes with an encoder, so without this a changed encoder + was invisible at the type level: reconciling each node re-encodes its + values and would notice, but a memoized declaring component is skipped + whole and never reconciles anything. A change here is a "lossy" + invalidation, which also invalidates every memo that declared into this + type, so its component runs again and its rows are re-upserted. + """ + + pg_type: str + encoder: str + + +def _code_digest(code: types.CodeType) -> str: + consts = tuple( + _code_digest(c) if isinstance(c, types.CodeType) else repr(c) + for c in code.co_consts + ) + payload = repr((code.co_code, consts, code.co_names)).encode() + return hashlib.sha256(payload).hexdigest()[:16] + + +def _encoder_identity(prop_def: PropertyDef) -> str: + """A stable name for what `prop_def.encoder` does, across runs. + + Empty for the built-in encoding of the property's type, whether the + definition left `encoder` unset or names the same function `from_class` + fills in β€” the two are one encoding. Otherwise the callable's qualified + name, with a digest of its code for a Python function, so editing a + lambda's body counts as a change the way swapping `str.lower` for + `str.upper` does. Closed-over values are not part of it, as with a + `@coco.fn` body's own helpers. + """ + encoder = prop_def.encoder + if encoder is None or encoder is _encoder_for(prop_def.pg_type): + return "" + module = ( + canonical_module_name(encoder) + if hasattr(encoder, "__module__") + else type(encoder).__module__ + ) + qualname = getattr(encoder, "__qualname__", None) or type(encoder).__qualname__ + code = getattr(encoder, "__code__", None) + if not isinstance(code, types.CodeType): + return f"{module}.{qualname}" + return f"{module}.{qualname}#{_code_digest(code)}" + + +_PROPERTY_SUBKEY_PREFIX = "prop:" + + +def _property_subkey(name: str) -> str: + return f"{_PROPERTY_SUBKEY_PREFIX}{name}" + + +def _property_name(subkey: str) -> str: + return subkey[len(_PROPERTY_SUBKEY_PREFIX) :] + + +#: Identity in `main`, one entry per property in `sub`. +_TypeTrackingRecord = statediff.MutualTrackingRecord[ + statediff.CompositeTrackingRecord[_TypeMainRecord, str, _PropertyRecord] +] + + +def _type_tracking_record_from_spec(spec: _TypeSpec) -> _TypeTrackingRecord: + """Build the tracking record for `spec`, already wrapped in + `MutualTrackingRecord`. + + Returns the *wrapped* record where neo4j's equivalent returns the bare + composite and wraps at the call site. Every caller here wants the wrapped + form, and `managed_by` riding along on the persisted record is the whole + point: it is what the removal path consults to decide it must not drop. + """ + sub = ( + { + _property_subkey(p.name): _PropertyRecord(p.pg_type, _encoder_identity(p)) + for p in spec.schema.properties.values() + } + if spec.schema is not None + else {} + ) + return statediff.MutualTrackingRecord( + tracking_record=statediff.CompositeTrackingRecord( + main=_TypeMainRecord( + key=spec.key, + from_type=spec.from_type, + to_type=spec.to_type, + owner=spec.owner, + ), + sub=sub, + ), + managed_by=spec.managed_by, + ) + + +class _TypeAction(NamedTuple): + key: _TypeKey + # The desired spec, carried through so `_apply_type_actions` can build the + # `_NodeHandler`/`_EdgeHandler` child for this type without needing a + # second lookup. `NON_EXISTENCE` when the type is being dropped β€” nothing + # to build a child from, and the sink returns no `ChildTargetDef` for a + # dropped type. + spec: _TypeSpec | coco.NonExistenceType + pg_fragment: str | None + main_action: statediff.DiffAction | None + property_actions: dict[str, statediff.DiffAction] + #: The app the action belongs to: whose ownership property the block + #: carries, and the only app whose edge types a node drop takes along. + owner: str + # A `managed_by=user` declaration of a type the connector created: the + # block's ownership property is stale and has to be removed, or a later + # drop of a node type it references takes the user's type along. + release_ownership: bool = False + + +_ChildInvalidation = Literal["destructive", "lossy"] | None + + +def _reconcile_removal( + key: _TypeKey, + prev_possible_records: typing.Collection[_TypeTrackingRecord], + prev_may_be_missing: bool, +) -> coco.TargetReconcileOutput[_TypeAction, _TypeTrackingRecord, Any] | None: + """Reconcile a type the app no longer declares. + + `resolve_system_transition` is what makes this safe: it returns `None` + when there is nothing tracked, and β€” crucially β€” when any previous record + is user-managed. The connector did not create a `managed_by=user` type and + does not get to delete it; dropping the block takes every row in it along, + which is the one irreversible thing this connector can do to data it was + explicitly told it does not manage. The old hand-rolled path dropped + unconditionally and structurally could not do better, because the desired + state is `NON_EXISTENCE` and the tracking record did not persist + `managed_by` at all. + """ + main_action, _ = statediff.diff_composite( + statediff.resolve_system_transition( + statediff.TrackingRecordTransition( + coco.NON_EXISTENCE, prev_possible_records, prev_may_be_missing + ) + ) + ) + if main_action is None: + # Safe to answer bare `None` *here*, and only here: the type is gone, + # so it owns no children that could be left without a provider. + # Answering `None` for a type that still exists is a shipped bug, not + # a style choice β€” this is a ROOT provider, the engine only refreshes + # a child's handler when the parent's own `reconcile()` output carries + # a fresh `ChildTargetDef`, and an unchanged type that returned `None` + # starved its own children the very next time it reconciled as + # unchanged: any node/edge declared under it then found no handler and + # the engine raised `RuntimeError: provider not ready for target state + # ...` β€” confirmed live, and confirmed against the house pattern + # (sqlite's `_TableHandler.reconcile` never returns bare `None` for + # exactly this reason). That is why `reconcile` below always returns an + # output for a type that exists, and why `_apply_type_actions` skips + # the schema write for an unchanged type but still builds and returns + # its child handler. + return None + # Every record in the store is this app's, so any of them names it. + owner = next(iter(prev_possible_records)).tracking_record.main.owner + return coco.TargetReconcileOutput( + action=_TypeAction( + key=key, + spec=coco.NON_EXISTENCE, + pg_fragment=None, + main_action=main_action, + property_actions={}, + owner=owner, + ), + sink=_type_sink, + tracking_record=coco.NON_EXISTENCE, + child_invalidation="destructive", + ) + + +class _TypeHandlerBase(coco.TargetHandler[_TypeSpec, _TypeTrackingRecord, Any]): + def _render(self, key: _TypeKey, spec: _TypeSpec) -> str: + raise NotImplementedError + + def reconcile( + self, + key: coco.StableKey, + desired_target_state: _TypeSpec | coco.NonExistenceType, + prev_possible_records: typing.Collection[_TypeTrackingRecord], + prev_may_be_missing: bool, + /, + ) -> coco.TargetReconcileOutput[_TypeAction, _TypeTrackingRecord, Any] | None: + key = _TypeKey(*_TYPE_KEY_CHECKER.check(key)) + if coco.is_non_existence(desired_target_state): + return _reconcile_removal(key, prev_possible_records, prev_may_be_missing) + + spec = desired_target_state + desired = _type_tracking_record_from_spec(spec) + + # --- The write path: what actually has to be applied. --- + # A SYSTEM declaration owns the desired schema, including when the + # previous declaration was USER-managed. Keep those previous schemas + # in the diff: filtering them out makes a simultaneous ownership and + # schema change look like an already-converged empty transition, then + # persists the new tracking record without ever applying its DDL. + # USER declarations still emit no schema action at all. + write_transition = ( + statediff.TrackingRecordTransition( + desired.tracking_record, + [p.tracking_record for p in prev_possible_records], + prev_may_be_missing, + ) + if spec.managed_by == ManagedBy.SYSTEM + else None + ) + main_action, property_transitions = statediff.diff_composite(write_transition) + property_actions: dict[str, statediff.DiffAction] = {} + if main_action is None: + # `_apply_type_schema` re-renders and re-issues this type's whole + # fragment on every write, so once a main action is scheduled the + # per-property actions are already subsumed by it. This is neo4j's + # `if main_action is None`, deliberately not sqlite's + # `in (None, "upsert")` β€” sqlite needs the wider gate because it + # emits per-column DDL, and we never do. + for sub_key, transition in property_transitions.items(): + action = statediff.diff(transition) + if action is not None: + property_actions[sub_key] = action + + # The reverse handoff. A SYSTEM declaration after a USER-managed + # record has to re-render the block even when the schema is otherwise + # unchanged: the release dropped `coco_managed_by_`, and until it is back + # a drop of a node type this type references treats it as the user's + # and refuses. + if ( + main_action is None + and not property_actions + and spec.managed_by == ManagedBy.SYSTEM + and any(p.managed_by == ManagedBy.USER for p in prev_possible_records) + ): + main_action = "upsert" + + # --- The tracked path: what CocoIndex has actually seen before. --- + # + # The same records, diffed as if `prev` were known complete and + # without `resolve_system_transition`: the non-nullable guard below + # and the child invalidation both need to see *through* ownership + # rather than past it. + # + # Reading the guard off the write path instead would silently disable + # it. Under `--reprocess` (or any other `prev_may_be_missing`) the + # main action becomes "upsert", which empties `property_actions` and + # leaves the guard nothing to fire on β€” reintroducing the raw + # `OG-MF-103` engine error it exists to replace. + tracked_main, tracked_transitions = statediff.diff_composite( + statediff.TrackingRecordTransition( + desired.tracking_record, + [p.tracking_record for p in prev_possible_records], + False, + ) + ) + tracked_actions: dict[str, statediff.DiffAction] = {} + if tracked_main is None: + for sub_key, transition in tracked_transitions.items(): + action = statediff.diff(transition) + if action is not None: + tracked_actions[sub_key] = action + + # A USER-managed type is never validated, let alone altered, by this + # connector: the schema is the user's, migrated with `omnigraph schema + # apply` on their own schedule, and the declaration is simply tracked + # from then on. An earlier guard compared the declaration against the + # *tracked* one and refused on any difference β€” and since reconcile + # never reads the live schema, the migration that guard advised could + # not unblock it: every later run failed the same way. The guard below + # is about a schema write this connector would itself issue, so it + # applies to SYSTEM-managed types only. + # `==`, not `is`: ManagedBy is a StrEnum, so a plain `"user"` string + # compares equal but is not identical β€” and an identity check would + # silently fall through to full SYSTEM management, letting the + # connector rewrite a schema the user said they own. The sibling + # connectors compare by value (see `connectorkits.statediff`). + if spec.managed_by != ManagedBy.USER and tracked_main is None: + # Verified in Task 1: `schema apply` rejects adding a NON-nullable + # property to an existing type (OG-MF-103, "requires a backfill and + # is not supported in schema migration v1"). Only initial creation + # or a full rebuild may declare one, so this fires on the in-place + # path only β€” `tracked_main is None`, never on a "replace". + # Failing here names the offending property in Python; letting it + # through surfaces as an opaque error from the Rust binary with + # nothing pointing at the user's dataclass. We do NOT silently + # escalate to a destructive rebuild β€” dropping a populated type + # because someone added a field is not a call the connector gets to + # make on the user's behalf. + # + # "upsert" counts as an addition alongside "insert": it means at + # least one candidate previous state lacked the property, and we + # don't get to pick which candidate is the engine's real state. + desired_sub = desired.tracking_record.sub + non_nullable_adds = sorted( + _property_name(sub_key) + for sub_key, action in tracked_actions.items() + if action in ("insert", "upsert") + and not desired_sub[sub_key].pg_type.endswith("?") + ) + if non_nullable_adds: + raise ValueError( + f"Cannot add non-nullable propert" + f"{'y' if len(non_nullable_adds) == 1 else 'ies'} " + f"{non_nullable_adds!r} to existing Omnigraph type: the engine " + f"only accepts non-nullable properties at initial creation. " + f"Make them optional (e.g. `str | None`), or drop and recreate " + f"the type deliberately." + ) + + child_invalidation: _ChildInvalidation = None + if tracked_main == "replace": + child_invalidation = "destructive" + elif tracked_main is None and any( + a != "insert" for a in tracked_actions.values() + ): + # A dropped or retyped property β€” and, after an interrupted update, + # one that only some candidate previous states carry β€” re-issues + # the whole fragment, so every child must re-upsert defensively. + child_invalidation = "lossy" + + # Handing a type the connector created to the user is the one thing + # a USER declaration writes: the block keeps its ownership property + # otherwise, and a later drop of a node type it references reads + # that as current ownership and removes the user's type with it. + release_ownership = spec.managed_by == ManagedBy.USER and any( + p.managed_by == ManagedBy.SYSTEM for p in prev_possible_records + ) + + return coco.TargetReconcileOutput( + action=_TypeAction( + key=key, + spec=spec, + pg_fragment=self._render(key, spec), + main_action=main_action, + property_actions=property_actions, + owner=spec.owner, + release_ownership=release_ownership, + ), + sink=_type_sink, + tracking_record=desired, + child_invalidation=child_invalidation, + ) + + +class _NodeTypeHandler(_TypeHandlerBase): + def _render(self, key: _TypeKey, spec: _TypeSpec) -> str: + assert isinstance(spec.schema, NodeSchema) + return spec.schema.render(key.type_name, owner=spec.owner) + + +class _EdgeTypeHandler(_TypeHandlerBase): + def _render(self, key: _TypeKey, spec: _TypeSpec) -> str: + assert spec.from_type is not None and spec.to_type is not None + props = ( + [(p.name, p.pg_type) for p in spec.schema.properties.values()] + if spec.schema is not None + else [] + ) + return render_edge_type( + key.type_name, spec.from_type, spec.to_type, props, owner=spec.owner + ) + + +# --------------------------------------------------------------------------- +# Child handlers: individual nodes and edges within a type +# --------------------------------------------------------------------------- + + +class _EntityTrackingRecord(msgspec.Struct, frozen=True, array_like=True): + """Fingerprint of the *encoded* property values, not the values themselves. + + Encoded, because that is what the graph holds: fingerprinting the raw + Python values let a changed `PropertyDef.encoder` rewrite every stored + value without change detection ever noticing. + + Internal-state size grows with target-state count, so storing full property + blobs for every entity in a large graph would bloat LMDB for no + reconciliation benefit. + + `array_like` for the same reason: this is the highest-cardinality record + the connector persists β€” one per node and per edge β€” so encoding it as an + object would pay for the literal key `"fingerprint"` on every single one, + which is exactly the overhead the paragraph above exists to avoid. The + siblings go further and track a bare `bytes` (neo4j's `_RowFingerprint`); + the struct is kept here only for the nominal typing. + """ + + fingerprint: bytes + + +class _NodeValue(NamedTuple): + properties: dict[str, Any] + + +class _EdgeValue(NamedTuple): + from_id: Any + to_id: Any + properties: dict[str, Any] + + +class _NodeAction(NamedTuple): + op: str # "upsert" | "delete" + key: _TypeKey + type_name: str + properties: Sequence[PropertyValue] + coco_key: str + #: On a delete: the key property in its transport encoding, so the sink + #: can turn the delete into a key-only upsert when an edge outside the + #: batch still references the node (see `_keep_referenced_nodes`). + key_properties: tuple[PropertyValue, ...] = () + + +class _EdgeAction(NamedTuple): + op: str # "insert" | "replace" | "delete" + key: _TypeKey + type_name: str + coco_key: str + from_id: Any + to_id: Any + properties: Sequence[PropertyValue] + from_type: str + to_type: str + from_key_property: PropertyDef + to_key_property: PropertyDef + + +def _entity_record(encoded: Sequence[PropertyValue]) -> _EntityTrackingRecord: + return _EntityTrackingRecord( + fingerprint=fingerprint_object({p.name: p.value for p in encoded}) + ) + + +def _needs_write( + desired: _EntityTrackingRecord, + prev_possible_records: typing.Collection[_EntityTrackingRecord], + prev_may_be_missing: bool, +) -> bool: + if prev_may_be_missing or not prev_possible_records: + return True + return not all(p.fingerprint == desired.fingerprint for p in prev_possible_records) + + +def _encode_properties( + properties: dict[str, Any], property_defs: dict[str, PropertyDef] +) -> tuple[PropertyValue, ...]: + """Resolve a raw `{name: value}` dict into typed, encoded `PropertyValue`s. + + Builders need each property's `pg_type` (to declare it in the query + signature), and `json.dumps` fails outright on `datetime.date`/ + `datetime.datetime` β€” both are `PropertyDef` concerns, so this is done + here, where the schema is in scope, rather than in `plan_commits`, which + is pure and has no schema to consult. + """ + out = [] + for name, value in properties.items(): + prop_def = property_defs[name] + out.append(_encode_property(prop_def, value)) + return tuple(out) + + +_EPOCH = datetime.datetime(1970, 1, 1, tzinfo=datetime.UTC) + + +def _tracking_key_value(prop_def: PropertyDef, value: Any) -> Any: + """The value a key property is tracked by: the graph's own identity. + + Omnigraph identifies a `DateTime`-keyed node by the instant, as epoch + milliseconds β€” `12:00+00:00` and `14:00+02:00` are one node, precision + below a millisecond is dropped, and a naive value is read as UTC (all + verified against the engine). Tracking by the ISO spelling gave one + graph node two tracking keys; removing one declaration then deleted the + node the other still declared. Every other key is tracked by the + encoded value the graph is sent, which is what it keys on. + + The instant can only be derived from a `datetime`, so a `DateTime` key + accepts nothing else: a hand-built `PropertyDef("at", "DateTime")` lets + a dict row carry an ISO string, and a string tracked by its spelling + reopens the two-keys-one-node hole. A `Date` key likewise takes a + `datetime.date` β€” and not a `datetime`, whose ISO form carries a time. + """ + base = prop_def.pg_type.rstrip("?") + if base == "DateTime": + if not isinstance(value, datetime.datetime): + raise TypeError( + f"Omnigraph key property {prop_def.name!r} is a DateTime: pass a " + f"`datetime.datetime`, not {type(value).__name__}. The graph " + f"identifies the node by the instant, which is derived from the " + f"datetime; a string would be tracked by its spelling, and two " + f"spellings of one instant would be two tracking keys for one node." + ) + if value.tzinfo is None: + value = value.replace(tzinfo=datetime.UTC) + return (value - _EPOCH) // datetime.timedelta(milliseconds=1) + if base == "Date" and ( + not isinstance(value, datetime.date) or isinstance(value, datetime.datetime) + ): + raise TypeError( + f"Omnigraph key property {prop_def.name!r} is a Date: pass a " + f"`datetime.date`, not {type(value).__name__}. The key is tracked by " + f"the date's ISO form, which only a `date` yields exactly." + ) + return _encode_property(prop_def, value).value + + +def _transport_key_value(prop_def: PropertyDef, tracked: Any) -> Any: + """The transport form of a key value the engine handed back as a + StableKey. Every tracking key is already the value the graph is sent β€” + except a `DateTime` key, tracked by its instant in epoch milliseconds, + which goes back out as ISO.""" + if prop_def.pg_type.rstrip("?") == "DateTime": + return (_EPOCH + datetime.timedelta(milliseconds=tracked)).isoformat() + return tracked + + +def _encode_property(prop_def: PropertyDef, value: Any) -> PropertyValue: + """Encode `value` for the transport: the definition's own encoder, or β€” + when it sets none β€” the built-in encoding for its `.pg` type. + + The built-in encoding follows the type, not the definition: a hand-built + `PropertyDef("at", "DateTime")` has no encoder set, and keying on that + field alone sent a raw `datetime` to `json.dumps` and a raw `date` into + the engine's StableKey, while the same schema inferred from a dataclass + worked because `from_class` filled the encoder in. + """ + encoder = prop_def.encoder + if encoder is None: + encoder = _encoder_for(prop_def.pg_type) + if value is not None and encoder is not None: + value = encoder(value) + return PropertyValue(prop_def.name, prop_def.pg_type, value) + + +class _NodeHandler(coco.TargetHandler[_NodeValue, _EntityTrackingRecord, Any]): + def __init__( + self, + type_name: str, + key_fields: tuple[str, ...], + key: _TypeKey, + property_defs: dict[str, PropertyDef], + ) -> None: + self._type_name = type_name + self._key_fields = key_fields + self._key = key + self._property_defs = property_defs + + def reconcile( + self, + key: coco.StableKey, + desired_target_state: _NodeValue | coco.NonExistenceType, + prev_possible_records: typing.Collection[_EntityTrackingRecord], + prev_may_be_missing: bool, + /, + ) -> coco.TargetReconcileOutput[_NodeAction, _EntityTrackingRecord, Any] | None: + # A composite key arrives as a tuple; a single-field key arrives as + # the bare scalar, so it must be wrapped before zipping with + # `key_fields`. + key_tuple = key if isinstance(key, tuple) else (key,) + if len(key_tuple) != len(self._key_fields): + raise ValueError( + f"Node type {self._type_name!r} declares key fields " + f"{self._key_fields!r} but the target state's key is " + f"{key_tuple!r}." + ) + coco_key = derive_coco_key(key_tuple) + + if coco.is_non_existence(desired_target_state): + # `prev_may_be_missing` means we may simply have lost track of a + # row that is still in the graph β€” after a `--reprocess`, after + # internal state was lost while the store persisted, or when a + # prior delete's tracking record was dropped before the delete + # landed. Only "nothing tracked AND the engine is certain nothing + # is missing" proves there is nothing to remove; anything else + # must still issue the delete, or the row survives untracked + # forever (nothing else ever removes it). Deleting a nonexistent + # entity by `coco_key` is a documented no-op, so the defensive + # delete costs nothing when it turns out to be unnecessary. Same + # guard every sibling connector uses, and the same reading the + # insert path below already applies to this exact signal. + if not prev_possible_records and not prev_may_be_missing: + return None + key_properties = tuple( + PropertyValue( + field, + self._property_defs[field].pg_type, + _transport_key_value(self._property_defs[field], value), + ) + for field, value in zip(self._key_fields, key_tuple, strict=True) + ) + return coco.TargetReconcileOutput( + action=_NodeAction( + "delete", + self._key, + self._type_name, + (), + coco_key, + key_properties=key_properties, + ), + sink=_entity_sink, + tracking_record=coco.NON_EXISTENCE, + ) + + # A keyed node `insert` merges by key tuple, so a new node and a + # changed node are both an "upsert" β€” unlike edges, there is no + # separate insert/replace distinction to make here. Encode first: + # the fingerprint covers what the engine is sent. + encoded = _encode_properties( + desired_target_state.properties, self._property_defs + ) + desired = _entity_record(encoded) + if not _needs_write(desired, prev_possible_records, prev_may_be_missing): + return None + return coco.TargetReconcileOutput( + action=_NodeAction("upsert", self._key, self._type_name, encoded, coco_key), + sink=_entity_sink, + tracking_record=desired, + ) + + +class _EdgeHandler(coco.TargetHandler[_EdgeValue, _EntityTrackingRecord, Any]): + def __init__( + self, + type_name: str, + key: _TypeKey, + from_type: str, + to_type: str, + from_key_property: PropertyDef, + to_key_property: PropertyDef, + property_defs: dict[str, PropertyDef], + ) -> None: + self._type_name = type_name + self._key = key + self._from_type = from_type + self._to_type = to_type + self._from_key_property = from_key_property + self._to_key_property = to_key_property + self._property_defs = property_defs + + def reconcile( + self, + key: coco.StableKey, + desired_target_state: _EdgeValue | coco.NonExistenceType, + prev_possible_records: typing.Collection[_EntityTrackingRecord], + prev_may_be_missing: bool, + /, + ) -> coco.TargetReconcileOutput[_EdgeAction, _EntityTrackingRecord, Any] | None: + # The edge's StableKey is its (from_id, to_id) tuple. + coco_key = derive_coco_key(key) + + if coco.is_non_existence(desired_target_state): + # `prev_may_be_missing` means we may simply have lost track of a + # row that is still in the graph β€” after a `--reprocess`, after + # internal state was lost while the store persisted, or when a + # prior delete's tracking record was dropped before the delete + # landed. Only "nothing tracked AND the engine is certain nothing + # is missing" proves there is nothing to remove; anything else + # must still issue the delete, or the row survives untracked + # forever (nothing else ever removes it). Deleting a nonexistent + # entity by `coco_key` is a documented no-op, so the defensive + # delete costs nothing when it turns out to be unnecessary. Same + # guard every sibling connector uses, and the same reading the + # insert path below already applies to this exact signal. + if not prev_possible_records and not prev_may_be_missing: + return None + return coco.TargetReconcileOutput( + action=_EdgeAction( + "delete", + self._key, + self._type_name, + coco_key, + None, + None, + (), + self._from_type, + self._to_type, + self._from_key_property, + self._to_key_property, + ), + sink=_entity_sink, + tracking_record=coco.NON_EXISTENCE, + ) + + # An edge `insert` is strict β€” it never merges, so a changed edge + # can't be re-inserted without duplicating it, and there is no edge + # `update`. A brand-new edge (nothing tracked yet, and the engine is + # certain nothing is missing) is the one case a bare insert is safe; + # anything already tracked that changed, OR anything whose prior + # state is merely uncertain (empty `prev_possible_records` with + # `prev_may_be_missing` β€” e.g. internal state lost while the target + # persists, or a crash between emitting a delete and it landing), + # must be deleted and re-inserted rather than risk duplicating an + # edge that's still there. Encoded here as a single "replace" op β€” + # `plan_commits` (Task 9) is what actually splits it into two + # mutations, since a delete and an insert can't share one mutation. + # Deleting a nonexistent edge by `coco_key` is a no-op, so "replace" + # is always safe even when the edge turns out not to have existed. + # + # Do NOT port this fork to `statediff.diff`: for edges it is inverted. + # `diff` answers "insert" for `(prev=[], prev_may_be_missing=True)` β€” + # the reading that suits a keyed upsert, where re-inserting converges. + # An Omnigraph edge has no key to upsert on, so that exact case is the + # one that needs "replace", and taking `diff`'s answer would duplicate + # a live edge every time internal state was lost. The container + # handler above uses `statediff` precisely because a type CAN be + # re-declared idempotently; an edge cannot. + encoded = _encode_properties( + desired_target_state.properties, self._property_defs + ) + desired = _entity_record(encoded) + if not prev_possible_records and not prev_may_be_missing: + op = "insert" + elif not _needs_write(desired, prev_possible_records, prev_may_be_missing): + return None + else: + op = "replace" + + return coco.TargetReconcileOutput( + action=_EdgeAction( + op, + self._key, + self._type_name, + coco_key, + desired_target_state.from_id, + desired_target_state.to_id, + encoded, + self._from_type, + self._to_type, + self._from_key_property, + self._to_key_property, + ), + sink=_entity_sink, + tracking_record=desired, + ) + + +# --------------------------------------------------------------------------- +# Commit planning: reconcile actions -> ordered Mutations, one per commit +# --------------------------------------------------------------------------- + +#: Chunk cap for render_query: a single commit combining more than this +#: many entities of one type risks an unreasonably large query. Deliberately +#: a module constant, not a public kwarg β€” tests monkeypatch it directly. +_MAX_ENTITIES_PER_TYPE = 8192 + +#: Cap on a whole commit, across every type in it. The per-type cap above +#: deliberately gives a large node type and a large edge type their own +#: budgets rather than making them share one β€” but with nothing bounding the +#: total, a sync touching N types produced a single commit of N x 8192 +#: entities (10 types = ~9.5 MB of GQ), and `_mutate_with_endpoint_retry` +#: re-sends that whole commit on every endpoint-stub retry. Set to 4x the +#: per-type cap: enough headroom for the common node-plus-its-edges shape to +#: stay in one commit, while still bounding the pathological case. +_MAX_ENTITIES_PER_COMMIT = 4 * _MAX_ENTITIES_PER_TYPE + + +#: `.pg` types a node key may have for that node type to be usable as an +#: edge endpoint. An edge's `from`/`to` is the engine's own node `id`, which +#: is always a String rendering of the key value β€” so the connector has to +#: reproduce that rendering exactly, and these are the types where it can: +#: `str` is itself, and Rust's i64/u32/... Display is `str(int)` digit for +#: digit. Verified against the engine, which renders the others in forms no +#: caller would guess: a `Date` key of 2026-01-05 has id `"20458"` (days +#: since the epoch) and a `DateTime` id is epoch milliseconds. +_ENDPOINT_KEY_TYPES = frozenset({"String", "I32", "I64", "U32", "U64"}) + + +def _endpoint_ref(value: Any) -> PropertyValue: + """Build the PropertyValue addressing an edge endpoint. + + Always `String`: an edge's `from`/`to` holds the endpoint node's `id`, + not its key property, and that id is a String whatever the key's own + declared type is β€” passing an `I64` for an `I64`-keyed endpoint is + rejected outright ("cannot assign/compare I64 with String for property + `to`", verified against the engine). `_validate_edge_endpoint` is what + keeps `str(value)` a faithful rendering of the id, by refusing endpoint + types whose key this can't reproduce. + + `.name` is ignored by `build_edge_insert` β€” only `.pg_type` and + `.value` address the endpoint. + """ + return PropertyValue("ref", "String", str(value)) + + +def _chunk_by_type( + items: Sequence[tuple[str, Statement]], cap: int, total_cap: int +) -> list[list[Statement]]: + """Split into chunks bounded two ways: no chunk holds more than `cap` + mutations for any single `type_name`, and none holds more than + `total_cap` overall. + + The per-type cap is checked per type, not across the whole phase, so a + large node type and a large edge type each get their own budget rather + than sharing one. `total_cap` is what keeps that from being unbounded: + without it a sync touching N types emitted one commit of N x `cap` + entities, which is re-sent whole on every endpoint-stub retry. + """ + chunks: list[list[Statement]] = [] + current: list[Statement] = [] + counts: dict[str, int] = {} + for type_name, mutation in items: + if counts.get(type_name, 0) >= cap or len(current) >= total_cap: + chunks.append(current) + current = [] + counts = {} + current.append(mutation) + counts[type_name] = counts.get(type_name, 0) + 1 + if current: + chunks.append(current) + return chunks + + +def plan_commits(actions: Sequence[Any]) -> list[Query]: + """Bucket reconcile actions into three phases and combine each phase's + mutations into one commit per phase (chunked past `_MAX_ENTITIES_PER_TYPE`). + + Phase A (pre-deletes): deletes of *replaced* edges only. Must precede + their re-insert in phase B β€” delete selects on `coco_key`, and a + replaced edge keeps the same one, so insert-then-delete would remove the + new edge too (verified against the engine: leaves zero edges). + + Phase B (upserts): node upserts, THEN edge inserts (both brand-new and + replaced). The order within the phase matters even though the whole + phase is one commit: statements in a combined query execute in the + order written, so an edge insert placed ahead of its own endpoint's + upsert fails with "not found" and forces the stub-and-retry path on + every run β€” a costly recovery meant for genuine cross-component + ordering violations, not for two writes the same component is issuing + together. Grouping nodes first makes that self-inflicted case + impossible; a failure can then only mean the endpoint really is owned + by another component that hasn't run yet. + + No endpoint stubs here β€” a keyed `insert` in Omnigraph is a + full-record replace, not a partial merge (verified against the engine), + so unconditionally stubbing an edge's endpoints would silently null out + any other nullable properties a node already has. Instead the sink + applies an edge insert optimistically and only builds a stub, for + *just* the specific endpoint the engine reports missing, after that + insert fails with a "not found" error β€” see `_apply_entity_actions`. + + Phase C (post-deletes): edge deletes, then node deletes β€” of removals + only (a replaced edge's delete is phase A, not this). Edges must be gone + before their nodes. + + A and C cannot merge into one commit even though both are deletes: A + must precede B and C must follow it. + """ + phase_a: list[tuple[str, Statement]] = [] + phase_b_nodes: list[tuple[str, Statement]] = [] + phase_b_edges: list[tuple[str, Statement]] = [] + phase_c_edges: list[tuple[str, Statement]] = [] + phase_c_nodes: list[tuple[str, Statement]] = [] + + for action in actions: + if isinstance(action, _NodeAction): + if action.op == "upsert": + phase_b_nodes.append( + ( + action.type_name, + build_node_upsert( + action.type_name, action.properties, action.coco_key + ), + ) + ) + elif action.op == "delete": + phase_c_nodes.append( + ( + action.type_name, + build_node_delete(action.type_name, action.coco_key), + ) + ) + else: + raise ValueError(f"unknown node action op: {action.op!r}") + elif isinstance(action, _EdgeAction): + if action.op == "delete": + phase_c_edges.append( + ( + action.type_name, + build_edge_delete(action.type_name, action.coco_key), + ) + ) + elif action.op in ("insert", "replace"): + if action.op == "replace": + phase_a.append( + ( + action.type_name, + build_edge_delete(action.type_name, action.coco_key), + ) + ) + from_ref = _endpoint_ref(action.from_id) + to_ref = _endpoint_ref(action.to_id) + phase_b_edges.append( + ( + action.type_name, + build_edge_insert( + action.type_name, + from_ref, + to_ref, + action.properties, + action.coco_key, + ), + ) + ) + else: + raise ValueError(f"unknown edge action op: {action.op!r}") + else: + raise TypeError(f"unknown action type: {type(action)!r}") + + phase_b = phase_b_nodes + phase_b_edges + phase_c = phase_c_edges + phase_c_nodes + commits: list[Query] = [] + for phase in (phase_a, phase_b, phase_c): + if not phase: + continue + for chunk in _chunk_by_type( + phase, _MAX_ENTITIES_PER_TYPE, _MAX_ENTITIES_PER_COMMIT + ): + commits.append(render_query(chunk)) + return commits + + +# --------------------------------------------------------------------------- +# Sinks +# --------------------------------------------------------------------------- + + +#: Emitted by the engine when an edge insert references a node that doesn't +#: exist yet β€” verified against the binary: `src '' not found in +#: ` for the `from` endpoint, `dst '' not found in ` for +#: `to`. The CLI colors its stderr with ANSI escapes even when not a TTY +#: (verified: piped subprocess output still carries them), and a leading +#: `\b` would break on that β€” the color-reset sequence right before "src"/ +#: "dst" ends in a letter (`\x1b[91m`), so there's no word boundary there. +#: `(src|dst) '...' not found in` is a tight enough anchor on its own; an +#: unrelated error would need this exact multi-token phrase to trigger a +#: spurious stub. +# Match the closing quote by its full error-message suffix rather than treating +# every apostrophe as the end of the key. String keys are allowed to contain +# apostrophes (for example, ``O'Brien``). +_ENDPOINT_NOT_FOUND_RE = re.compile(r"(src|dst) '([\s\S]*?)' not found in (\w+)") + + +def _parse_missing_endpoint(error: OmnigraphCliError) -> tuple[str, str, str] | None: + """Extract `(role, key, type_name)` from a "not found" failure, or + `None` if `error` doesn't match. Doesn't build the stub itself β€” the + caller tracks which `(type_name, key)` pairs it has already stubbed, so + it needs the parsed identity before deciding whether to act on it.""" + m = _ENDPOINT_NOT_FOUND_RE.search(str(error)) + if m is None: + return None + return m.group(1), m.group(2), m.group(3) + + +def _build_endpoint_stub( + role: str, key_str: str, type_name: str, edge_actions: Sequence[_EdgeAction] +) -> Statement | None: + """Build the key-only stub for the endpoint the engine reported missing. + Returns `None` if no action in this batch accounts for the (type, key) + β€” nothing safe to build a stub from. + + The engine's error names the endpoint's type and the string form of its + key value, but not the key's *field* name or its `pg_type` β€” those come + from matching (type, key) against the `_EdgeAction`s in this batch, + which is also how the actual (correctly typed) key value is recovered + rather than re-parsing it out of the error text. + """ + for action in edge_actions: + if action.op == "delete": + # A delete carries no endpoints (`from_id`/`to_id` are `None`) and + # needs none β€” deleting by `coco_key` never touches them. Without + # this, `str(None)` matches any node whose key value is literally + # "None" and can build a meaningless stub from the delete action. + continue + if ( + role == "src" + and action.from_type == type_name + and str(action.from_id) == key_str + ): + value, key_property = action.from_id, action.from_key_property + elif ( + role == "dst" + and action.to_type == type_name + and str(action.to_id) == key_str + ): + value, key_property = action.to_id, action.to_key_property + else: + continue + ref = _encode_property(key_property, value) + return build_endpoint_stub( + type_name, + [ref], + derive_coco_key((value,)), + ) + return None + + +async def _mutate_with_endpoint_retry( + client: _CliClient, + commit: Query, + *, + branch: str, + edge_actions: Sequence[_EdgeAction], +) -> None: + """Apply `commit` optimistically β€” without any endpoint stub β€” since the + common case is that both endpoints either already exist or aren't + needed at all. A keyed `insert` in Omnigraph is a full-record replace, + not a partial merge (verified against the engine), so unconditionally + stubbing every edge's endpoints was proven to silently null out a + node's other nullable properties whenever the owning component had + already written real values. Building a stub only reactively, from a + "not found" failure, is provably safe: the engine has just confirmed + the node doesn't exist, so there's nothing to wipe. + + The engine reports only the *first* missing endpoint per attempt + (verified: with both `from` and `to` absent, only `src` is reported), + and with up to 1024 components running concurrently, an edge component + preceding the components that own its endpoints is ordinary, not + exotic. A commit carries every edge of a component, so the number of + stub rounds it may need is the number of distinct endpoints those + edges reference β€” not two, which is the budget for a single edge and + failed any commit with three or more absent endpoints. + + The loop is bounded by exactly that set: a stub is only ever built for + an endpoint some edge in `edge_actions` references, and each `(type, + key)` is stubbed at most once. A "not found" that names an endpoint + already stubbed, or one no edge here references, propagates instead + of spinning. + """ + stubbed: set[tuple[str, str]] = set() + while True: + try: + await client.mutate(commit, branch=branch) + return + except OmnigraphCliError as e: + parsed = _parse_missing_endpoint(e) + if parsed is None: + raise + role, key_str, type_name = parsed + if (type_name, key_str) in stubbed: + raise + stub = _build_endpoint_stub(role, key_str, type_name, edge_actions) + if stub is None: + raise + stubbed.add((type_name, key_str)) + await client.mutate(render_query([stub]), branch=branch) + + +#: Prefix of every branch this connector creates. Nothing else creates +#: branches under it, so a branch carrying it that is visible while the +#: store lock is held was abandoned by a process that died mid-sync. +_SCRATCH_BRANCH_PREFIX = "coco_scratch_" + + +@contextlib.asynccontextmanager +async def _scratch_branch(client: _CliClient, *, frm: str) -> AsyncIterator[str]: + """Own a scratch branch for the duration of the block. + + Holds the store lock from before the branch is created until after it + is deleted: a non-main branch makes Omnigraph reject every schema apply + on the store, and schema reconciliation takes the same lock, so a + concurrent type update never observes the transient branch and fails + spuriously. The branch's own liveness lock is held for the same span, + which is what lets `_reap_abandoned_scratch_branches` tell an abandoned + branch from a live one even when the store lock did not cover both + parties. + + The branch is deleted whether the body succeeded or failed β€” `branch + merge` does not delete its source (we don't pass `--delete-branch`, + which would only cover the success path anyway), and a leftover + non-main branch blocks every subsequent `schema apply` on the store + (verified live). `branch_delete` itself raises when the branch doesn't + exist, so if `branch_create` never landed there is nothing to delete; + once it did, a cleanup failure must be reported β€” treating the sync as + successful would persist a hidden operational failure. + """ + async with client.store_lock(): + name = f"{_SCRATCH_BRANCH_PREFIX}{uuid.uuid4().hex}" + async with client.hold_scratch_branch(name): + await client.branch_create(name, frm=frm) + try: + yield name + finally: + await client.branch_delete(name) + + +async def _reap_abandoned_scratch_branches(client: _CliClient) -> list[str]: + """Delete every abandoned `coco_scratch_*` branch and return their names. + + Runs under the store lock, and deletes only a branch whose liveness + lock it can take: a live scratch branch holds that lock for its whole + lifetime inside `_scratch_branch`, so one whose lock is free was left + behind by a process that died between creating and deleting it. + Abandonment is established before deletion, not inferred from the + branch merely being visible β€” a reaper that took a different store + lock than the branch's owner (verified live with two spellings of one + store URI) deleted a branch that was still in use. User branches are + never touched. + """ + reaped: list[str] = [] + for name in await client.branch_list(): + if not name.startswith(_SCRATCH_BRANCH_PREFIX): + continue + async with client.claim_scratch_branch(name) as abandoned: + if abandoned: + await client.branch_delete(name) + reaped.append(name) + return reaped + + +async def _apply_schema_recovering_abandoned_branches( + client: _CliClient, schema_pg: str +) -> None: + """`schema apply`, recovering from a scratch branch an interrupted sync + left behind. + + Omnigraph refuses every schema change while any non-main branch exists + ("schema apply requires a graph with only main; found non-main + branches: ..."). An abandoned `coco_scratch_*` branch therefore blocked + the store's schema for good, since nothing ever listed or deleted it. + On that refusal, reap the connector's own abandoned branches and retry + once; if the refusal was about a user's branch, nothing is reaped and + the original error propagates. + """ + try: + await client.apply_schema(schema_pg) + except OmnigraphCliError as e: + if "non-main branches" not in str(e): + raise + if not await _reap_abandoned_scratch_branches(client): + raise + await client.apply_schema(schema_pg) + + +async def _keep_referenced_nodes( + client: _CliClient, actions: Sequence[_NodeAction | _EdgeAction], *, branch: str +) -> list[_NodeAction | _EdgeAction]: + """Turn the delete of a node that an edge outside `actions` still + references into a key-only upsert of that node. + + Deleting a node cascades to its edges in the graph (`affected_edges` in + the mutation's result, verified against the engine) β€” silently, from + the point of view of whichever component declared those edges: their + tracking still says they exist, so once the node came back the edges' + reconcile found nothing to do and they were gone for good. The edge is + the declared state that still holds, and the connector already has a + shape for "an edge whose endpoint nobody declares": the key-only stub + it writes when an edge arrives before its node. A referenced node is + reduced to that same stub, so the edge survives, and the node's next + declaration fills it back in. Edges this batch deletes itself do not + count, so a node undeclared together with all of its edges is deleted + outright, as before. + + The edge types to look at come from the live schema, one read per + batch; nodes of a type no edge type references cost nothing more. Each + remaining delete costs one read per (edge type, end) at which its type + sits. An edge inserted by another component between that read and this + batch's commit is still cascaded; the window is the same one every + unlocked entity write has. + """ + deletes = [a for a in actions if isinstance(a, _NodeAction) and a.op == "delete"] + if not deletes: + return list(actions) + existing = await client.read_schema() + if existing is None: + return list(actions) + patterns = { + type_name: incident_edge_patterns(existing, type_name) + for type_name in {a.type_name for a in deletes} + } + if not any(patterns.values()): + return list(actions) + deleted_edges = { + (a.type_name, a.coco_key) + for a in actions + if isinstance(a, _EdgeAction) and a.op == "delete" + } + out: list[_NodeAction | _EdgeAction] = [] + for action in actions: + if not (isinstance(action, _NodeAction) and action.op == "delete"): + out.append(action) + continue + referenced = False + for pattern in patterns[action.type_name]: + rows = await client.query( + render_query( + [ + build_incident_edges_query( + action.type_name, action.coco_key, pattern + ) + ] + ), + branch=branch, + ) + if any( + (pattern.edge_type, row["edge_key"]) not in deleted_edges + for row in rows + ): + referenced = True + break + if referenced: + action = _NodeAction( + "upsert", + action.key, + action.type_name, + action.key_properties, + action.coco_key, + ) + out.append(action) + return out + + +async def _apply_entity_actions( + context_provider: ContextProvider, actions: Sequence[_NodeAction | _EdgeAction] +) -> None: + if not actions: + return + by_db: dict[str, list[_NodeAction | _EdgeAction]] = {} + for action in actions: + by_db.setdefault(action.key.db_key, []).append(action) + + for db_key, db_actions in by_db.items(): + conn: ConnectionFactory = context_provider.get(db_key) + client = _CliClient(conn) + db_actions = await _keep_referenced_nodes( + client, db_actions, branch=conn.branch + ) + commits = plan_commits(db_actions) + if not commits: + continue + edge_actions = [a for a in db_actions if isinstance(a, _EdgeAction)] + if len(commits) == 1: + # One commit is atomic on its own, so send it straight to the + # live branch β€” no scratch branch, no merge. Deliberately NOT + # `_mutate_with_endpoint_retry`: that recovery is two commits + # (the stub, then the retried commit), which on the live branch + # would leave an orphan stub node behind β€” untracked, and + # visible to readers β€” if the retry then failed for an + # unrelated reason. A failed mutation applies nothing (verified + # against the engine; see TestMutationAtomicityLive), so the + # optimistic attempt costs nothing but the round-trip, and the + # moment recovery is actually needed this falls through to the + # scratch-branch path below and redoes the whole thing there. + try: + await client.mutate(commits[0], branch=conn.branch) + continue + except OmnigraphCliError as e: + if _parse_missing_endpoint(e) is None: + raise + # More than one commit needs atomicity across all of them together, + # and one CLI invocation is one commit β€” so apply the whole sequence + # on a scratch branch and merge it in as a single step. `branch merge` + # has no compare-and-swap precondition (as of 0.10.0 only `mutate` + # takes `--if-commit`); a conflicted merge simply surfaces as a + # non-zero exit from `branch_merge` and propagates like any other + # failure below. + async with _scratch_branch(client, frm=conn.branch) as scratch: + for commit in commits: + await _mutate_with_endpoint_retry( + client, commit, branch=scratch, edge_actions=edge_actions + ) + await client.branch_merge(scratch, into=conn.branch) + + +_entity_sink = coco.TargetActionSink.from_async_fn(_apply_entity_actions) + + +def _check_no_kind_clash(existing_pg: str, pending: Sequence[_TypeAction]) -> None: + """Refuse to add a type whose name is already taken by the OTHER kind. + + A `.pg` holding both `node Link` and `edge Link` is accepted by `schema + apply` β€” but at the mutation level the name resolves to the node type + alone, so `insert Link { from: ..., to: ... }` fails with "type `Link` + has no property `from`" (verified against the binary). The edge type + exists in the schema and can never be written to, read back, or + deleted. Fail here, naming both, rather than let a component declare + edges into a type nothing can address. + + Checked against the live schema rather than tracked at declare time: + the two kinds register under separate providers, so a node and an edge + of the same name can be declared by components that never see each + other, and only the graph itself has the whole picture. + """ + other = {"node": "edge", "edge": "node"} + for action in pending: + if coco.is_non_existence(action.spec): + continue + clashing = other[action.key.type_kind] + if _find_type_block(existing_pg, clashing, action.key.type_name) is None: + continue + raise ValueError( + f"Omnigraph type name {action.key.type_name!r} is already used by " + f"a {clashing} type in this graph, and a mutation can only ever " + f"address one of them β€” the name resolves to the {clashing}, so " + f"the {action.key.type_kind} type would be unwritable. Rename one " + f"of them." + ) + + +def _check_no_kind_clash_within_batch(pending: Sequence[_TypeAction]) -> None: + """Refuse a batch that writes a node and an edge under the same name. + + Same rule and same consequence as `_check_no_kind_clash` below, but + applied to the batch rather than the live schema β€” which is the only + place it can be caught when both types are new. The live-schema check + can only ever see types that are ALREADY in the graph, so two fresh + clashing types slip past it on an initialized store, and on a fresh + store it isn't reached at all (the `init` path returns first). + """ + kind_by_name: dict[str, str] = {} + for action in pending: + if coco.is_non_existence(action.spec): + continue + name = action.key.type_name + first_kind = kind_by_name.setdefault(name, action.key.type_kind) + if first_kind != action.key.type_kind: + raise ValueError( + f"Omnigraph type name {name!r} is declared as both a node and " + f"an edge in the same sync, and a mutation can only ever " + f"address one of them β€” the name resolves to the node, so the " + f"edge type would be unwritable. Rename one of them." + ) + + +class _BlockedRemoval(NamedTuple): + """A node type this batch removes or rebuilds, and the edge types not in + this batch that still point at it.""" + + node_name: str + is_drop: bool # False for a "replace" (a `@key` change) + edge_names: list[str] + + +def _blocked_node_removals( + existing_pg: str, pending: Sequence[_TypeAction] +) -> list[_BlockedRemoval]: + """Node removals in `pending` that would leave an edge type dangling. + + The engine rejects such a schema outright β€” `catalog error: edge 'X' has + an unresolved endpoint`, verified against the binary β€” and the connector + could emit one two different ways: dropping a node type leaves the + dangling edge in the FINAL schema, and a "replace" (a `@key` change) + leaves it in the INTERMEDIATE schema that the drop half applies first. + + Every mounted type is its own processing component, so a referencing + edge type's removal reaches the sink in its own batch, not alongside + the node's β€” unless the batcher happened to group them. That is why + this has to be checked here: the engine's own error names the edge but + offers no way forward, and arrives after the connector has already + decided to write. + """ + going_away = { + a.key.type_name: coco.is_non_existence(a.spec) + for a in pending + if a.key.type_kind == "node" + and (coco.is_non_existence(a.spec) or a.main_action == "replace") + } + if not going_away: + return [] + also_dropped = { + a.key.type_name + for a in pending + if a.key.type_kind == "edge" and coco.is_non_existence(a.spec) + } + blocked = [] + for node_name, is_drop in sorted(going_away.items()): + edge_names = sorted( + set(edge_types_referencing(existing_pg, node_name)) - also_dropped + ) + if edge_names: + blocked.append(_BlockedRemoval(node_name, is_drop, edge_names)) + return blocked + + +def _orphaned_endpoint_error( + blocked: _BlockedRemoval, *, foreign: dict[str, str | None], marker: str +) -> ValueError: + """`foreign` maps each referencing edge type this app does not own to + the ownership property its block does declare, if any β€” another app's + β€” so the message can say whose it is.""" + plural = len(blocked.edge_names) > 1 + if foreign: + many = len(foreign) > 1 + owned_by_others = "; ".join( + f"{name!r} is marked `{other}`" for name, other in foreign.items() if other + ) + advice = ( + f"{list(foreign)!r} {'are' if many else 'is'} not managed by this app " + f"({'their blocks declare' if many else 'its block declares'} no " + f"`{marker}` property{'; ' + owned_by_others if owned_by_others else ''}), " + f"so {'they' if many else 'it'} will not be removed for you: have the " + f"app that owns {'them' if many else 'it'} stop declaring " + f"{'them' if many else 'it'} first, remove {'them' if many else 'it'} " + f"from the schema yourself, or keep the node type." + ) + else: + advice = ( + f"Stop declaring {'those edge types' if plural else 'that edge type'} " + f"as well, or keep the node type. A `@key` change counts here too: it " + f"drops and recreates the type, so its endpoints go with it." + ) + return ValueError( + f"Cannot remove Omnigraph node type {blocked.node_name!r}: edge " + f"type{'s' if plural else ''} {blocked.edge_names!r} still " + f"{'reference' if plural else 'references'} it as an endpoint, " + f"and the engine rejects a schema with an unresolved endpoint. {advice}" + ) + + +async def _apply_type_schema( + client: _CliClient, actions: Sequence[_TypeAction] +) -> None: + """Fold every schema-changing action for one graph into one read and at + most two writes. + + Omnigraph's schema is applied whole-graph, not per type (verified + against the engine: applying a single new type's fragment alone + silently dropped every other existing type). Task 6 reconciles each + type independently and only ever renders that one type's fragment, so + every action funnels through here: read the current whole-graph + source, fold this batch's fragments into it in memory, and write the + result back. The whole batch arrives in one call β€” the sink's batcher + never has more than one group in flight β€” so doing a full + read-merge-write round-trip per action would be N times the CLI + invocations for the same end state. + + `kind` doesn't select a code path. A brand-new type's action is always + `"create"`, even when it's the *second* type ever declared on an + already-`init`'d graph, so `client.read_schema()` returning `None` is + the only signal for "not yet initialized" β€” it subsumes the old + create-or-alter stderr-sniffing fallback entirely. + + Two exceptions to the single write: + + `"drop"` removes the type's block. Nothing else deletes a dropped + type's nodes or edges: when a container target state disappears the + engine emits no per-child deletes, so if this didn't drop the block, + the type and every row in it would persist forever, untracked. + + `"replace"` (a `@key` or edge-endpoint change) cannot go through the + ordinary merge: the engine flatly rejects an in-place key change + ("removing property constraints ... not supported in schema migration + v1", verified against the binary β€” not a soft/hard-drop distinction, + `--allow-data-loss` doesn't help). It needs a genuine rebuild, so a + batch containing any `"replace"` costs one extra `schema apply` that + lands the replaced types' removal first, before the final apply + re-adds them under their new definitions. `child_invalidation= + "destructive"` on a `"replace"` already tells every child entity it + must be re-declared from scratch, which is exactly what the drop half + does at the schema level. + """ + pending = [ + a + for a in actions + if a.main_action is not None or a.property_actions or a.release_ownership + ] + if not pending: + return + + # Omnigraph applies schema source as a whole. Keep the read, in-memory + # merge, and every resulting apply under one store-scoped lock so two app + # updates cannot both read the same schema and overwrite each other's + # changes with competing complete-schema writes. + async with client.store_lock(): + await _apply_type_schema_locked(client, pending) + + +async def _apply_type_schema_locked( + client: _CliClient, pending: Sequence[_TypeAction] +) -> None: + # Before the `init` path below, not after it: on a fresh graph that path + # returns first, so a clash inside the very first batch was never checked. + _check_no_kind_clash_within_batch(pending) + + existing = await client.read_schema() + if existing is None: + # Nothing to read, nothing to drop from, and `init` takes the + # complete schema β€” so build the whole thing in memory and create + # the graph in a single call. + merged = "" + for action in pending: + if coco.is_non_existence(action.spec) or action.release_ownership: + # Nothing to drop from, and nothing of the user's to release. + continue + assert action.pg_fragment is not None + merged = merge_type_into_schema( + merged, action.key.type_kind, action.key.type_name, action.pg_fragment + ) + if merged: + await client.init_graph(merged) + return + + _check_no_kind_clash(existing, pending) + + # A batch is one app's, and that app's ownership property is what marks + # the blocks it may take along or release. + owners = {a.owner for a in pending} + assert len(owners) == 1, owners + marker = ownership_property(owners.pop()) + + # A node type's drop may find an edge type of this app still pointing + # at it. Every mounted type is its own component and all of them share + # one sink batcher, so when an app is dropped the edge type's own + # removal is queued behind this very batch β€” waiting for it here + # deadlocks until a timeout (verified live). Its removal is coming + # regardless, so take this app's own edge types along now; the edge's + # batch then finds nothing left to remove. An edge type marked as + # another app's is that app's to remove β€” its removal is not coming, + # and a marker that only said "some app of this connector" let one + # app's drop delete another's edges β€” and one without any marker is a + # user's or another tool's; neither is removed on their behalf, and a + # `@key` change keeps the edge declared, so all three still fail here. + taken_along: list[str] = [] + for removal in _blocked_node_removals(existing, pending): + foreign = { + name: found + for name in removal.edge_names + if (found := ownership_marker(existing, "edge", name)) != marker + } + if not removal.is_drop or foreign: + raise _orphaned_endpoint_error(removal, foreign=foreign, marker=marker) + taken_along.extend(removal.edge_names) + + replaced = [ + (a.key.type_kind, a.key.type_name) + for a in pending + if a.main_action == "replace" + ] + if replaced: + intermediate = existing + for kind, type_name in replaced: + intermediate = remove_type_from_schema(intermediate, kind, type_name) + await _apply_schema_recovering_abandoned_branches(client, intermediate) + existing = intermediate + + desired = existing + for action in pending: + if coco.is_non_existence(action.spec): + desired = remove_type_from_schema( + desired, action.key.type_kind, action.key.type_name + ) + elif action.release_ownership: + desired = release_ownership( + desired, action.key.type_kind, action.key.type_name, marker + ) + else: + assert action.pg_fragment is not None + desired = merge_type_into_schema( + desired, action.key.type_kind, action.key.type_name, action.pg_fragment + ) + for edge_name in taken_along: + desired = remove_type_from_schema(desired, "edge", edge_name) + # An unchanged schema needs no write β€” the edge type's own removal + # arriving after a node's batch took it along is the common case. + if desired != existing: + await _apply_schema_recovering_abandoned_branches(client, desired) + + +async def _apply_type_actions( + context_provider: ContextProvider, actions: Sequence[_TypeAction] +) -> list[coco.ChildTargetDef[_NodeHandler | _EdgeHandler] | None]: + actions_list = list(actions) + outputs: list[coco.ChildTargetDef[_NodeHandler | _EdgeHandler] | None] = [ + None + ] * len(actions_list) + + by_db: dict[str, list[int]] = {} + for i, action in enumerate(actions_list): + by_db.setdefault(action.key.db_key, []).append(i) + + for db_key, idxs in by_db.items(): + conn: ConnectionFactory = context_provider.get(db_key) + client = _CliClient(conn) + await _apply_type_schema(client, [actions_list[i] for i in idxs]) + + for i in idxs: + action = actions_list[i] + spec = action.spec + if coco.is_non_existence(spec): + # No child handler for a type that no longer exists. + continue + # An unchanged type wrote no schema, but its child handler must + # still be rebuilt here, or the engine leaves this type's own + # children without a provider on the next run. + if spec.from_type is None: + # Node type: `key` is the node's own key field names. + assert spec.schema is not None + outputs[i] = coco.ChildTargetDef( + handler=_NodeHandler( + action.key.type_name, + spec.key, + action.key, + spec.schema.properties, + ) + ) + else: + # Edge type: `from_type`/`to_type` name the endpoints, and + # `from_key_property`/`to_key_property` are the endpoints' own + # key definitions β€” needed to build correctly typed stubs. + assert spec.to_type is not None + assert spec.from_key_property is not None + assert spec.to_key_property is not None + outputs[i] = coco.ChildTargetDef( + handler=_EdgeHandler( + action.key.type_name, + action.key, + spec.from_type, + spec.to_type, + spec.from_key_property, + spec.to_key_property, + spec.schema.properties if spec.schema is not None else {}, + ) + ) + return outputs + + +_type_sink: coco.TargetActionSink[_TypeAction, _NodeHandler | _EdgeHandler] = ( + coco.TargetActionSink.from_async_fn(_apply_type_actions) +) + + +# --------------------------------------------------------------------------- +# Root provider registration +# --------------------------------------------------------------------------- + +_node_type_provider: coco.TargetStateProvider[_TypeSpec, _NodeHandler] = ( + coco.register_root_target_states_provider( + "cocoindex/omnigraph/node_type", _NodeTypeHandler() + ) +) +_edge_type_provider: coco.TargetStateProvider[_TypeSpec, _EdgeHandler] = ( + coco.register_root_target_states_provider( + "cocoindex/omnigraph/edge_type", _EdgeTypeHandler() + ) +) + + +# --------------------------------------------------------------------------- +# NodeTarget +# --------------------------------------------------------------------------- + +RowT = TypeVar("RowT", default=dict[str, Any]) + + +class NodeTarget( + coco.ResolvesTo["NodeTarget[RowT]"], Generic[RowT, coco.MaybePendingS] +): + """A target for writing nodes to an Omnigraph node type.""" + + _provider: coco.TargetStateProvider[_NodeValue, None, coco.MaybePendingS] + _schema: NodeSchema + _type_name: str + + def __init__( + self, + provider: coco.TargetStateProvider[_NodeValue, None, coco.MaybePendingS], + schema: NodeSchema, + type_name: str, + ) -> None: + self._provider = provider + self._schema = schema + self._type_name = type_name + + @property + def type_name(self) -> str: + return self._type_name + + @property + def schema(self) -> NodeSchema: + return self._schema + + def declare_node(self: NodeTarget[RowT], *, node: RowT) -> None: + """Declare a node (record) to be upserted to this node type. + + The target-state key is the graph's identity for the key value β€” + see `_tracking_key_value`: epoch milliseconds for a `DateTime`, the + ISO form for a `Date`, the value itself for a string or integer. A + raw `datetime` is not a type the engine's StableKey accepts, and + tracking a spelling the graph does not key on let two declarations + of one node get two tracking keys. + """ + properties = _record_to_dict(node, self._schema) + key: tuple[Any, ...] = tuple( + _tracking_key_value(self._schema.properties[k], properties[k]) + for k in self._schema.key + ) + coco.declare_target_state( + self._provider.target_state(key, _NodeValue(properties)) + ) + + def __coco_memo_key__(self) -> str: + return self._provider.memo_key + + +# --------------------------------------------------------------------------- +# EdgeTarget +# --------------------------------------------------------------------------- + + +class EdgeTarget( + coco.ResolvesTo["EdgeTarget[RowT]"], Generic[RowT, coco.MaybePendingS] +): + """A target for writing edges to an Omnigraph edge type.""" + + _provider: coco.TargetStateProvider[_EdgeValue, None, coco.MaybePendingS] + _schema: EdgeSchema | None + _type_name: str + _from_target: NodeTarget[Any] + _to_target: NodeTarget[Any] + + def __init__( + self, + provider: coco.TargetStateProvider[_EdgeValue, None, coco.MaybePendingS], + schema: EdgeSchema | None, + type_name: str, + from_target: NodeTarget[Any], + to_target: NodeTarget[Any], + ) -> None: + self._provider = provider + self._schema = schema + self._type_name = type_name + self._from_target = from_target + self._to_target = to_target + + @property + def type_name(self) -> str: + return self._type_name + + def declare_edge( + self: EdgeTarget[RowT], *, from_id: Any, to_id: Any, record: RowT | None = None + ) -> None: + """Declare an edge between the two nodes this edge type connects.""" + if record is not None and self._schema is None: + # Omnigraph is schema-first: an edge type declares its properties + # in `.pg`, and an insert naming one that isn't declared is + # rejected outright ("type `X` has no property `y`", verified + # against the engine). So a record cannot be carried by a + # schema-less edge type, and inferring the property types at + # write time wouldn't help β€” the type genuinely has no column to + # put them in. This is the one real difference from the Neo4j + # original a ported app hits: Neo4j's MERGE creates relationship + # properties on the fly, so the same code works there. + raise TypeError( + f"Edge type {self._type_name!r} was mounted without a schema, " + f"so it cannot carry a record. Pass the edge's own property " + f"schema (e.g. `EdgeSchema.from_class({type(record).__name__})`) " + f"when mounting it, or drop the `record=` argument." + ) + if record is None and self._schema is not None: + # The mirror of the guard above. Omitting the record leaves every + # declared property out of the insert, and Omnigraph rejects an + # insert missing a non-nullable one ("must provide non-nullable + # property") with an error naming neither this call nor the + # property. A schema of only nullable properties is fine -- that + # is what nullable means. + required = sorted( + name + for name, prop in self._schema.properties.items() + if not prop.pg_type.endswith("?") + ) + if required: + raise TypeError( + f"Edge type {self._type_name!r} declares non-nullable " + f"propert{'y' if len(required) == 1 else 'ies'} " + f"{required!r}, so `record=` is required. Pass the record, " + f"or declare those properties as optional." + ) + _check_endpoint_id(from_id, self._type_name, "from_id", self._from_target) + _check_endpoint_id(to_id, self._type_name, "to_id", self._to_target) + properties = _record_to_dict(record, self._schema) if record is not None else {} + key = (from_id, to_id) + coco.declare_target_state( + self._provider.target_state(key, _EdgeValue(from_id, to_id, properties)) + ) + + def __coco_memo_key__(self) -> str: + return self._provider.memo_key + + +def _endpoint_key_family(pg_type: str) -> str: + """`String` keys and integer keys are the two shapes `str(value)` can + reproduce as a node id; every integer width renders the same way.""" + return "String" if pg_type == "String" else "integer" + + +def _check_endpoint_id( + value: Any, type_name: str, what: str, endpoint: NodeTarget[Any] +) -> None: + """An edge endpoint is addressed by its node's single key VALUE, spliced + into the insert as `from: $e_from`. Checked here, where the offending + `declare_edge` call is on the stack, rather than at sink time β€” where + `_endpoint_ref` would raise the same underlying `TypeError` from inside + `plan_commits`, with nothing pointing at the declaration that caused it. + + The value is also checked against the endpoint type's own key: the id + is `str(value)`, so an integer handed to a String-keyed endpoint would + silently address whichever node's key happens to be those digits, and a + string handed to an integer-keyed endpoint can never match a node. + """ + try: + pg_type = _pg_type_for(type(value)) + except TypeError: + pg_type = None + if pg_type not in _ENDPOINT_KEY_TYPES: + raise TypeError( + f"Edge type {type_name!r}: {what}={value!r} is not usable as an " + f"endpoint reference. It must be the endpoint node's key value β€” a " + f"single string or integer, matching what `mount_edge_target` " + f"already required of the endpoint type's key." + ) + (key_field,) = endpoint.schema.key + key_type = endpoint.schema.properties[key_field].pg_type + if _endpoint_key_family(pg_type) != _endpoint_key_family(key_type): + raise TypeError( + f"Edge type {type_name!r}: {what}={value!r} is a {pg_type} value, but " + f"the endpoint node type {endpoint.type_name!r} is keyed by " + f"{key_field!r}: {key_type}. Pass that node's key value." + ) + + +def _record_to_dict( + record: Any, schema: NodeSchema | EdgeSchema | None +) -> dict[str, Any]: + """Extract `{property_name: raw_value}` from a dataclass or dict record. + + Values are NOT encoded here β€” `_encode_properties` (used by + `_NodeHandler`/`_EdgeHandler.reconcile`, above) applies `PropertyDef.encoder` + at commit time, once the property's `pg_type` is back in scope. + """ + property_names = schema.properties.keys() if schema is not None else None + if isinstance(record, dict): + if schema is None: + return dict(record) + # A dict record's keys are unchecked user input. A misspelled name used + # to be dropped without a word, leaving EVERY declared property `None` + # β€” so the key's `coco_key` was derived from `None` and the engine + # rejected a null into a non-nullable `@key` column, naming neither the + # record nor the typo. (The dataclass branch below needs no equivalent: + # its schema is normally derived from that same class by `from_class`, + # and a missing attribute already raises.) + unknown = sorted(set(record) - set(schema.properties)) + if unknown: + raise ValueError( + f"Record declares propert{'y' if len(unknown) == 1 else 'ies'} " + f"{unknown!r} not in the schema; declared properties are " + f"{sorted(schema.properties)!r}." + ) + missing = sorted( + name + for name, prop in schema.properties.items() + if name not in record and not prop.pg_type.endswith("?") + ) + if missing: + raise ValueError( + f"Record is missing non-nullable propert" + f"{'y' if len(missing) == 1 else 'ies'} {missing!r}. Only a " + f"nullable property (e.g. `str | None`) may be omitted." + ) + return {name: record.get(name) for name in schema.properties} + if property_names is None: + return {f.name: getattr(record, f.name) for f in dataclasses.fields(record)} + return {name: getattr(record, name) for name in property_names} + + +def _validate_edge_endpoint(target: NodeTarget[Any], role: str) -> None: + """Refuse, at mount time, a node type that cannot serve as an edge + endpoint β€” naming the type and what's wrong with it, rather than + failing mid-sync from whatever component happens to declare the edge. + + Two independent requirements: + + Its key must be one the connector can render the way the engine does. + An edge's `from`/`to` holds the endpoint's node `id`, a String rendering + of the key value, and only `String` and integer keys have a rendering + `str(value)` reproduces (see `_ENDPOINT_KEY_TYPES`). + + And it must be stubbable. An edge may be written before the component + owning one of its endpoints has run, so the connector recovers by + inserting a key-only stub for the endpoint the engine reports missing + (see `_build_endpoint_stub`). That stub can never satisfy a non-nullable + property outside the key β€” Omnigraph rejects the insert. + """ + schema = target.schema + if len(schema.key) != 1: + # `NodeSchema.from_class` guarantees this, but `NodeSchema` is public + # and can be built by hand, so don't let it fall through to an + # unpacking error with nothing naming the type. + raise ValueError( + f"Node type {target.type_name!r} cannot be referenced as an edge " + f"endpoint ({role}): its key is {schema.key!r}, and an endpoint is " + f"addressed by a single key value." + ) + (key_field,) = schema.key + key_type = schema.properties[key_field].pg_type + if key_type not in _ENDPOINT_KEY_TYPES: + raise ValueError( + f"Node type {target.type_name!r} cannot be referenced as an edge " + f"endpoint ({role}): its key {key_field!r} is {key_type}, and an " + f"edge addresses an endpoint by the node's id β€” a String rendering " + f"of the key that CocoIndex can only reproduce for " + f"{sorted(_ENDPOINT_KEY_TYPES)}. Key this type on a string or an " + f"integer instead." + ) + unstubbable = [ + p.name + for p in schema.properties.values() + if p.name not in schema.key and not p.pg_type.endswith("?") + ] + if unstubbable: + raise ValueError( + f"Node type {target.type_name!r} cannot be referenced as an edge " + f"endpoint ({role}): an edge may be written before the component " + f"that owns the node has run, so the connector inserts a key-only " + f"stub β€” which Omnigraph rejects while {unstubbable!r} " + f"{'is' if len(unstubbable) == 1 else 'are'} non-nullable. Declare " + f"{'it' if len(unstubbable) == 1 else 'them'} optional (e.g. " + f"`str | None`)." + ) + + +def _build_edge_spec( + schema: EdgeSchema | None, + from_target: NodeTarget[Any], + to_target: NodeTarget[Any], + managed_by: ManagedBy, + owner: str, +) -> _TypeSpec: + """Pure seam between `edge_target()` and `_TypeSpec` construction β€” kept + separate so the endpoint key-field wiring (Task 10 step 4) is directly + testable without reaching into `coco.TargetState`'s private value.""" + if schema is not None and not isinstance(schema, EdgeSchema): + raise TypeError( + f"An edge type's schema must be an EdgeSchema, not " + f"{type(schema).__name__}: an edge has no key of its own, its " + f"identity is always (from_id, to_id). Build it with " + f"`EdgeSchema.from_class(...)`." + ) + _validate_edge_endpoint(from_target, "from") + _validate_edge_endpoint(to_target, "to") + (from_key_field,) = from_target.schema.key + (to_key_field,) = to_target.schema.key + return _TypeSpec( + schema=schema, + key=(), + from_type=from_target.type_name, + to_type=to_target.type_name, + managed_by=managed_by, + owner=owner, + from_key_property=from_target.schema.properties[from_key_field], + to_key_property=to_target.schema.properties[to_key_field], + ) + + +# --------------------------------------------------------------------------- +# Module-level entry points +# --------------------------------------------------------------------------- + + +def node_target( + db: coco.ContextKey[ConnectionFactory], + type_name: str, + schema: NodeSchema, + *, + managed_by: ManagedBy = ManagedBy.SYSTEM, +) -> coco.TargetState[_NodeHandler]: + """Create a `TargetState` for an Omnigraph node type. + + Must run inside a component: the type's block is marked as owned by + the app declaring it, which is only known there. + """ + validate_identifier(type_name, "node type") + if not schema.key: + # `NodeSchema.from_class` guarantees a key, but `NodeSchema` is public + # and can be built by hand β€” and an empty key is silently catastrophic + # rather than merely wrong: the type renders with no `@key`, every row + # derives the same `coco_key` from the same empty tuple and so shares + # one StableKey (last write wins in tracking), while an unkeyed + # Omnigraph insert is a strict insert that duplicates the row on every + # re-run. `_validate_edge_endpoint` already refuses this, but only for + # types used as endpoints. + raise ValueError( + f"Node type {type_name!r} must declare a key: its schema's key is " + f"empty, so every row would share one identity. Build the schema " + f"with `NodeSchema.from_class(..., key=...)`, or pass a non-empty " + f"`key` to `NodeSchema`." + ) + for key_field in schema.key: + prop = schema.properties[key_field] + allowed = _KEY_ENCODERS.get(prop.pg_type.rstrip("?"), ()) + if prop.encoder is not None and prop.encoder not in allowed: + # The key is the node's identity, and an edge addresses that + # identity by the raw key value its author passed to + # `declare_edge`. A custom encoder would make the graph key on + # something else: changing one from `str.lower` to `str.upper` + # upserted a second node under the new spelling and left the old + # one behind, and an edge declared with the raw value would never + # find either. Only the fixed Date/DateTime encoders are allowed + # on a key, each on the type it encodes β€” pure functions of the + # value, never changing, and such types cannot be edge endpoints + # at all. + raise ValueError( + f"Node type {type_name!r}: key property {key_field!r} has a custom " + f"encoder. The key is the node's identity, and edges address it by " + f"the raw key value, so an encoder there would make the graph's key " + f"disagree with what edges reference. Normalize the key value before " + f"declaring the node instead; only the built-in ISO encoding may sit " + f"on a Date or DateTime key." + ) + type_key = _TypeKey(db_key=db.key, type_kind="node", type_name=type_name) + spec = _TypeSpec( + schema=schema, + key=schema.key, + from_type=None, + to_type=None, + managed_by=managed_by, + owner=current_app_name(), + ) + return _node_type_provider.target_state(type_key, spec) + + +def declare_node_target( + db: coco.ContextKey[ConnectionFactory], + type_name: str, + schema: NodeSchema, + *, + managed_by: ManagedBy = ManagedBy.SYSTEM, +) -> NodeTarget[Any, coco.PendingS]: + """Declare a node type target. + + Use this for node types that exist only as edge endpoints β€” no nodes + flow into this declaration's own handler. + """ + provider = coco.declare_target_state_with_child( + node_target(db, type_name, schema, managed_by=managed_by) + ) + return NodeTarget(provider, schema, type_name) + + +async def mount_node_target( + db: coco.ContextKey[ConnectionFactory], + type_name: str, + schema: NodeSchema, + *, + managed_by: ManagedBy = ManagedBy.SYSTEM, +) -> NodeTarget[Any]: + """Mount a node type target ready to receive `declare_node` calls.""" + provider = await coco.mount_target( + node_target(db, type_name, schema, managed_by=managed_by) + ) + return NodeTarget(provider, schema, type_name) + + +def edge_target( + db: coco.ContextKey[ConnectionFactory], + type_name: str, + from_target: NodeTarget[Any], + to_target: NodeTarget[Any], + schema: EdgeSchema | None = None, + *, + managed_by: ManagedBy = ManagedBy.SYSTEM, +) -> coco.TargetState[_EdgeHandler]: + """Create a `TargetState` for an Omnigraph edge type. + + Must run inside a component, like `node_target`. + """ + validate_identifier(type_name, "edge type") + spec = _build_edge_spec( + schema, from_target, to_target, managed_by, current_app_name() + ) + key = _TypeKey(db_key=db.key, type_kind="edge", type_name=type_name) + return _edge_type_provider.target_state(key, spec) + + +def declare_edge_target( + db: coco.ContextKey[ConnectionFactory], + type_name: str, + from_target: NodeTarget[Any], + to_target: NodeTarget[Any], + schema: EdgeSchema | None = None, + *, + managed_by: ManagedBy = ManagedBy.SYSTEM, +) -> EdgeTarget[Any, coco.PendingS]: + """Declare an edge type target.""" + provider = coco.declare_target_state_with_child( + edge_target( + db, type_name, from_target, to_target, schema, managed_by=managed_by + ) + ) + return EdgeTarget(provider, schema, type_name, from_target, to_target) + + +async def mount_edge_target( + db: coco.ContextKey[ConnectionFactory], + type_name: str, + from_target: NodeTarget[Any], + to_target: NodeTarget[Any], + schema: EdgeSchema | None = None, + *, + managed_by: ManagedBy = ManagedBy.SYSTEM, +) -> EdgeTarget[Any]: + """Mount an edge type target ready to receive `declare_edge` calls.""" + provider = await coco.mount_target( + edge_target( + db, type_name, from_target, to_target, schema, managed_by=managed_by + ) + ) + return EdgeTarget(provider, schema, type_name, from_target, to_target) + + +# --------------------------------------------------------------------------- +# Public exports +# --------------------------------------------------------------------------- + +__all__ = [ + "ConnectionFactory", + "EdgeSchema", + "EdgeTarget", + "NodeSchema", + "NodeTarget", + "OmnigraphType", + "PropertyDef", + "ValueEncoder", + "declare_edge_target", + "declare_node_target", + "edge_target", + "mount_edge_target", + "mount_node_target", + "node_target", +] diff --git a/python/tests/connectors/test_omnigraph_target.py b/python/tests/connectors/test_omnigraph_target.py new file mode 100644 index 000000000..145835a1d --- /dev/null +++ b/python/tests/connectors/test_omnigraph_target.py @@ -0,0 +1,7036 @@ +"""Tests for the Omnigraph target connector. + +Run with: + uv run pytest python/tests/connectors/test_omnigraph_target.py -v + +Builder unit tests run without a server. Live tests require the omnigraph +CLI at test/bin/omnigraph and are gated on OMNIGRAPH_TEST_STORE=1; CI installs +the pinned release there (see .github/workflows/_test.yml). +""" + +from __future__ import annotations + +import asyncio +import datetime +import json +import os +import shutil +import subprocess +import tempfile +import uuid +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Annotated, Any, cast + +import cocoindex as coco +import pytest +from cocoindex._internal.context_keys import ContextKey, ContextProvider +from cocoindex.connectorkits import statediff +from cocoindex.connectorkits.target import ManagedBy +from cocoindex.connectors import omnigraph +from cocoindex.connectors.omnigraph import _gq as ogq +from cocoindex.connectors.omnigraph import _target as ogt +from cocoindex.connectors.omnigraph._client import ( + ConnectionFactory, + OmnigraphCliError, + _CliClient, +) +from cocoindex.connectors.omnigraph._gq import ( + COCO_KEY, + Bind, + PropertyValue, + Query, + Statement, + build_edge_delete, + build_edge_insert, + build_endpoint_stub, + build_node_delete, + build_node_upsert, + merge_type_into_schema, + ownership_marker, + ownership_property, + release_ownership, + remove_type_from_schema, + render_edge_type, + render_node_type, + render_property, + render_query, + validate_identifier, + validate_pg_type, +) +from cocoindex.connectors.omnigraph._target import ( + EdgeSchema, + NodeSchema, + OmnigraphType, + PropertyDef, + _EdgeHandler, + _EdgeTypeHandler, + _EdgeValue, + _NodeHandler, + _NodeTypeHandler, + _NodeValue, + _type_tracking_record_from_spec, + _TypeKey, + _TypeSpec, + derive_coco_key, + plan_commits, +) + +from tests import common + +coco_env = common.create_test_env(__file__) + + +@dataclass +class _Doc: + slug: str + title: str + words: int + ratio: float + published: bool + on: datetime.date + at: datetime.datetime + note: str | None + + +class TestValidateIdentifier: + @pytest.mark.parametrize("name", ["Person", "_private", "T1", "a_b_c", "WorksAt"]) + def test_valid(self, name: str) -> None: + validate_identifier(name, "node type") + + @pytest.mark.parametrize( + "name", + [ + "my-type", + "123abc", + "", + "has space", + "semi;colon", + "a.b", + "back`tick", + "foo\n", + "foo\r", + "foo\t", + "foo\x00", + "foo\nbar", + ], + ) + def test_invalid(self, name: str) -> None: + with pytest.raises(ValueError, match="Invalid Omnigraph node type"): + validate_identifier(name, "node type") + + +class TestRenderProperty: + def test_plain(self) -> None: + assert render_property("title", "String", is_key=False) == "title: String" + + def test_key(self) -> None: + assert render_property("slug", "String", is_key=True) == "slug: String @key" + + def test_nullable_is_passed_through(self) -> None: + assert render_property("age", "I32?", is_key=False) == "age: I32?" + + def test_pg_type_cannot_inject_another_schema_block(self) -> None: + with pytest.raises(ValueError, match="Invalid Omnigraph property type"): + render_property( + "title", + "String\n}\nnode Surprise {\n slug: String @key", + is_key=False, + ) + + +class TestRenderNodeType: + def test_injects_coco_key_and_the_ownership_marker(self) -> None: + """Every block this connector renders carries `coco_key` and a + nullable, never-written `coco_managed_by_`. The latter is what + lets the schema sink tell this app's types from a user's and from + another app's β€” a user-managed type must declare `coco_key` too, so + that property alone cannot β€” and it is a property rather than a + comment because the engine only stores an applied source when + something structural changed, so a comment could never be released + on handoff.""" + assert render_node_type( + "Source", + [("slug", "String"), ("title", "String")], + key=("slug",), + owner="og", + ) == ( + "node Source {\n slug: String @key\n title: String\n" + " coco_key: String\n coco_managed_by_og: Bool?\n}" + ) + + def test_zero_properties_still_gets_coco_key(self) -> None: + assert render_node_type("Empty", [], key=(), owner="og") == ( + "node Empty {\n coco_key: String\n coco_managed_by_og: Bool?\n}" + ) + + def test_composite_key_raises(self) -> None: + """Omnigraph allows exactly one `@key` per node type: a two-`@key` + block is rejected outright at `init` with "node type Reading has + multiple @key constraints; only one is supported" (verified against + the binary). Rendering it anyway would emit a `.pg` the engine can + never accept, so it's refused here.""" + with pytest.raises(ValueError, match="exactly one @key property"): + render_node_type( + "Reading", + [("sensor", "String"), ("at", "DateTime"), ("v", "F64")], + key=("sensor", "at"), + owner="og", + ) + + def test_reserved_property_id_raises(self) -> None: + """`id` is the engine's own identity column, materialized from the + `@key` property. Declaring one fails at `init` with "physical schema + for 'node:X' must contain exactly one top-level `id` field; found 2", + and keying on it fails with the far more confusing "@key must + reference declared properties" β€” both verified against the binary.""" + with pytest.raises(ValueError, match=r"\['id'\] is reserved by Omnigraph"): + render_node_type( + "X", [("slug", "String"), ("id", "I64")], key=("slug",), owner="og" + ) + + def test_key_property_must_exist(self) -> None: + with pytest.raises(ValueError, match="key property 'missing' is not declared"): + render_node_type("X", [("a", "String")], key=("missing",), owner="og") + + def test_key_must_not_be_nullable(self) -> None: + with pytest.raises(ValueError, match="key property 'a' must not be nullable"): + render_node_type("X", [("a", "String?")], key=("a",), owner="og") + + def test_caller_may_not_shadow_coco_key(self) -> None: + with pytest.raises(ValueError, match="reserved"): + render_node_type( + "X", [("a", "String"), (COCO_KEY, "String")], key=("a",), owner="og" + ) + + +class TestRenderEdgeType: + def test_no_properties_still_gets_coco_key(self) -> None: + assert render_edge_type("Supports", "Source", "Claim", [], owner="og") == ( + "edge Supports: Source -> Claim {\n" + " coco_key: String\n coco_managed_by_og: Bool?\n}" + ) + + def test_with_properties(self) -> None: + assert render_edge_type( + "WorksAt", "Person", "Company", [("role", "String")], owner="og" + ) == ( + "edge WorksAt: Person -> Company {\n role: String\n" + " coco_key: String\n coco_managed_by_og: Bool?\n}" + ) + + +class TestOwnershipProperty: + def test_an_identifier_app_name_is_used_as_is(self) -> None: + assert ownership_property("meeting_notes") == "coco_managed_by_meeting_notes" + + def test_other_names_are_made_identifiers_without_colliding(self) -> None: + dashed, underscored = ownership_property("my-app"), ownership_property("my_app") + assert dashed != underscored + assert dashed.startswith("coco_managed_by_my_app_") + validate_identifier(dashed, "property name") + + +class TestMergeTypeIntoSchema: + """Omnigraph's schema is applied whole-graph, not per type β€” applying + one type's fragment alone silently drops every other type (verified + against the engine). This merge is what makes reconciling each type + independently safe; it is fiddly text handling and deserves direct + tests rather than only a live end-to-end check.""" + + def test_appends_a_new_type(self) -> None: + existing = ( + render_node_type("A", [("slug", "String")], key=("slug",), owner="og") + + "\n" + ) + b = render_node_type("B", [("slug", "String")], key=("slug",), owner="og") + merged = merge_type_into_schema(existing, "node", "B", b) + assert "node A {" in merged + assert "node B {" in merged + assert merged.index("node A {") < merged.index("node B {") + + def test_replaces_an_existing_type(self) -> None: + a1 = render_node_type("A", [("slug", "String")], key=("slug",), owner="og") + a2 = render_node_type( + "A", [("slug", "String"), ("title", "String")], key=("slug",), owner="og" + ) + merged = merge_type_into_schema(a1 + "\n", "node", "A", a2) + assert merged == a2 + "\n" + assert "title" in merged + + def test_leaves_unrelated_types_untouched(self) -> None: + a = render_node_type("A", [("slug", "String")], key=("slug",), owner="og") + b = render_node_type("B", [("slug", "String")], key=("slug",), owner="og") + existing = f"{a}\n\n{b}\n" + b2 = render_node_type( + "B", [("slug", "String"), ("title", "String")], key=("slug",), owner="og" + ) + merged = merge_type_into_schema(existing, "node", "B", b2) + assert a in merged + assert b2 in merged + assert b not in merged + + def test_replaces_the_brace_less_edge_form(self) -> None: + """A property-less edge is legal and renders brace-less (`edge E: A + -> B`, no `{ }` at all) β€” verified against the engine, which both + accepts this form and reproduces it via `schema show`. This + connector's own builders never emit it (they always add + `coco_key`), but the schema being merged into can legitimately + contain one from elsewhere, and the merge must still find its + exact extent β€” ending at the line, not spilling into whatever + follows.""" + existing = ( + "node A {\n slug: String @key\n coco_key: String\n}\n\n" + "node B {\n slug: String @key\n coco_key: String\n}\n\n" + "edge E: A -> B\n" + ) + e2 = render_edge_type("E", "A", "B", [("weight", "I64")], owner="og") + merged = merge_type_into_schema(existing, "edge", "E", e2) + assert "node A {\n slug: String @key\n coco_key: String\n}" in merged + assert "node B {\n slug: String @key\n coco_key: String\n}" in merged + assert e2 in merged + # The old brace-less form must be gone -- not just shadowed by the + # new braced one, which also starts with "edge E: A -> B". + assert "edge E: A -> B\n" not in merged + + def test_an_edge_fragment_never_touches_a_same_named_node_block(self) -> None: + """A `.pg` may legally hold `node Link` and `edge Link` side by side + β€” the engine accepts it (verified against the binary). Matching on + the name alone wrote the edge's fragment over the NODE's block, + silently destroying the node type and leaving TWO `edge Link` + blocks behind: the same whole-graph destruction class as applying a + single type's fragment alone, keyed on a name collision instead of + an omission.""" + existing = ( + "node Link {\n slug: String @key\n coco_key: String\n}\n\n" + "edge Link: A -> B {\n coco_key: String\n}\n" + ) + frag = render_edge_type("Link", "A", "B", [("weight", "I64?")], owner="og") + merged = merge_type_into_schema(existing, "edge", "Link", frag) + assert "node Link {\n slug: String @key\n coco_key: String\n}" in merged + assert merged.count("edge Link") == 1 + assert "weight: I64?" in merged + + def test_a_node_fragment_never_touches_a_same_named_edge_block(self) -> None: + existing = ( + "node Link {\n slug: String @key\n coco_key: String\n}\n\n" + "edge Link: A -> B {\n coco_key: String\n}\n" + ) + frag = render_node_type( + "Link", + [("slug", "String"), ("title", "String?")], + key=("slug",), + owner="og", + ) + merged = merge_type_into_schema(existing, "node", "Link", frag) + assert "edge Link: A -> B {\n coco_key: String\n}" in merged + assert merged.count("node Link") == 1 + assert "title: String?" in merged + + def test_unbalanced_braces_raise_instead_of_corrupting(self) -> None: + """With no closing brace the splice point used to stay at the + OPENING brace, so the fragment landed mid-declaration and the + corrupt result went straight to `schema apply`. Input is engine + output today, so this is defense in depth on the whole-graph write + path β€” but a bad edit there costs the entire schema.""" + broken = "node A {\n slug: String @key\n coco_key: String\n" + with pytest.raises(ValueError, match="unbalanced braces"): + merge_type_into_schema(broken, "node", "A", "node A {\n}") + + def test_duplicate_blocks_raise_instead_of_editing_only_the_first(self) -> None: + """Only the first block was ever found, so a merge rewrote one and + left the other as a stale duplicate, and a removal deleted one and + left the other behind β€” exactly the state the kind-blind match used + to create.""" + a = render_node_type("A", [("slug", "String")], key=("slug",), owner="og") + with pytest.raises(ValueError, match="declares node 'A' 2 times"): + merge_type_into_schema(a + "\n\n" + a + "\n", "node", "A", a) + + def test_rejects_an_unknown_kind(self) -> None: + with pytest.raises(ValueError, match="must be 'node' or 'edge'"): + merge_type_into_schema("", "relation", "A", "node A {\n}") + + # `schema show` returns the source exactly as it was written (verified + # against the engine), so a schema a person wrote or edited by hand keeps + # its formatting β€” and a merger that only recognised `node X {` at the + # start of a line, ended blocks at the first `}` it saw, and appended + # what it failed to find, produced schemas the engine then refused + # ("duplicate node name", "expected EOI or schema_decl"). + + def test_finds_an_indented_block(self) -> None: + existing = " node Person {\n slug: String @key\n coco_key: String\n }\n" + frag = render_node_type( + "Person", [("slug", "String"), ("age", "I64?")], key=("slug",), owner="og" + ) + merged = merge_type_into_schema(existing, "node", "Person", frag) + assert merged.count("node Person") == 1 + assert frag in merged + assert " slug: String @key" not in merged + + def test_finds_the_second_of_two_declarations_on_one_line(self) -> None: + person = "node Person { slug: String @key coco_key: String }" + company = "node Company { slug: String @key coco_key: String }" + existing = f"{person} {company}\n" + frag = render_node_type( + "Company", + [("slug", "String"), ("name", "String?")], + key=("slug",), + owner="og", + ) + merged = merge_type_into_schema(existing, "node", "Company", frag) + assert merged.count("node Company") == 1 + assert person in merged + assert frag in merged + + def test_a_brace_inside_a_comment_does_not_end_the_block(self) -> None: + existing = ( + "node Person {\n slug: String @key // closes } here\n coco_key: String\n}\n\n" + "node Company {\n slug: String @key\n coco_key: String\n}\n" + ) + frag = render_node_type( + "Person", [("slug", "String"), ("age", "I64?")], key=("slug",), owner="og" + ) + merged = merge_type_into_schema(existing, "node", "Person", frag) + assert merged.count("node Person") == 1 + assert "} here" not in merged + assert "node Company {\n slug: String @key\n coco_key: String\n}" in merged + + def test_a_declaration_inside_a_comment_is_not_a_block(self) -> None: + existing = ( + "// node Ghost { slug: String @key }\n" + "node A {\n slug: String @key\n coco_key: String\n}\n" + ) + assert remove_type_from_schema(existing, "node", "Ghost") == existing + frag = render_node_type( + "A", [("slug", "String"), ("t", "String?")], key=("slug",), owner="og" + ) + merged = merge_type_into_schema(existing, "node", "A", frag) + assert merged.startswith("// node Ghost { slug: String @key }\n") + assert merged.count("node A") == 1 + + def test_a_property_named_like_a_keyword_is_not_a_block(self) -> None: + existing = ( + "node A {\n slug: String @key\n edge: String?\n node: String?\n" + " coco_key: String\n}\n" + ) + assert remove_type_from_schema(existing, "edge", "String") == existing + frag = render_node_type("A", [("slug", "String")], key=("slug",), owner="og") + assert merge_type_into_schema(existing, "node", "A", frag) == frag + "\n" + + def test_edge_endpoints_are_read_from_an_indented_declaration(self) -> None: + from cocoindex.connectors.omnigraph._gq import edge_types_referencing + + existing = ( + "node A {\n slug: String @key\n coco_key: String\n}\n" + " edge E: A -> A {\n coco_key: String\n }\n" + "// edge Ghost: A -> A\n" + ) + assert edge_types_referencing(existing, "A") == ["E"] + + +class TestRemoveTypeFromSchema: + """The first half of the two-step `@key`-change rebuild: the engine + rejects an in-place key change but accepts drop-then-re-add as two + separate `schema apply` calls (verified against the engine).""" + + def test_removes_the_middle_block_without_doubling_the_separator(self) -> None: + a = render_node_type("A", [("slug", "String")], key=("slug",), owner="og") + b = render_node_type("B", [("slug", "String")], key=("slug",), owner="og") + c = render_node_type("C", [("slug", "String")], key=("slug",), owner="og") + existing = f"{a}\n\n{b}\n\n{c}\n" + removed = remove_type_from_schema(existing, "node", "B") + assert a in removed + assert c in removed + assert "node B {" not in removed + assert "\n\n\n" not in removed + + def test_removes_the_only_block(self) -> None: + a = render_node_type("A", [("slug", "String")], key=("slug",), owner="og") + removed = remove_type_from_schema(a + "\n", "node", "A") + assert "node A" not in removed + + def test_absent_type_is_a_no_op(self) -> None: + a = render_node_type("A", [("slug", "String")], key=("slug",), owner="og") + existing = a + "\n" + assert remove_type_from_schema(existing, "node", "Z") == existing + + def test_round_trips_with_merge_for_a_key_change(self) -> None: + """The exact rebuild sequence _apply_type_schema drives: remove the + old block, then merge the new definition back in β€” proving the two + functions compose into a schema with the type's key actually + changed and every other type still present.""" + a1 = render_node_type("A", [("slug", "String")], key=("slug",), owner="og") + b = render_node_type("B", [("slug", "String")], key=("slug",), owner="og") + existing = f"{a1}\n\n{b}\n" + without_a = remove_type_from_schema(existing, "node", "A") + a2 = render_node_type( + "A", [("slug", "String"), ("title", "String")], key=("title",), owner="og" + ) + rebuilt = merge_type_into_schema(without_a, "node", "A", a2) + assert a2 in rebuilt + assert b in rebuilt + assert a1 not in rebuilt # A's old (slug-keyed) definition is gone + + +class TestIncidentEdges: + SCHEMA = ( + "node Person {\n slug: String @key\n coco_key: String\n}\n" + "node Company {\n cid: I64 @key\n coco_key: String\n}\n" + "edge WorksAt: Person -> Company {\n coco_key: String\n}\n" + "edge Knows: Person -> Person {\n coco_key: String\n}\n" + ) + + def test_patterns_cover_both_ends_and_self_loops(self) -> None: + assert ogq.incident_edge_patterns(self.SCHEMA, "Person") == [ + ogq.IncidentEdge("WorksAt", "from"), + ogq.IncidentEdge("Knows", "from"), + ogq.IncidentEdge("Knows", "to"), + ] + assert ogq.incident_edge_patterns(self.SCHEMA, "Company") == [ + ogq.IncidentEdge("WorksAt", "to") + ] + assert ogq.incident_edge_patterns(self.SCHEMA, "Nope") == [] + + def test_query_matches_by_coco_key_in_traversal_spelling(self) -> None: + """The read side names an edge type by its traversal spelling, which + starts lowercase (`worksAt` for `edge WorksAt`, verified against + the engine, which looks it up case-insensitively). The node is + found by `coco_key`, bound as a parameter like every value.""" + out = render_query( + [ + ogq.build_incident_edges_query( + "Person", "k1", ogq.IncidentEdge("WorksAt", "from") + ) + ] + ) + assert out.expr == ( + "query m($s0_p_coco_key: String) { match { " + "$n: Person { coco_key: $s0_p_coco_key } $n $e:worksAt $o } " + "return { $e.coco_key as edge_key } }" + ) + assert out.params == {"s0_p_coco_key": "k1"} + incoming = render_query( + [ + ogq.build_incident_edges_query( + "Company", "k7", ogq.IncidentEdge("WorksAt", "to") + ) + ] + ) + assert "$o $e:worksAt $n }" in incoming.expr + + +class TestOwnershipDetection: + """`coco_managed_by_` is a declared property. Ownership detection + searched the block's raw text, comments included, so a user-owned block + whose comment merely mentioned the property counted as the connector's + β€” and an app drop deleted it. Detection and release must look at + declarations only.""" + + USER_EDGE = ( + "edge E: A -> B {\n" + " // coco_managed_by_og: Bool? was deliberately omitted\n" + " coco_key: String\n" + "}\n" + ) + OWNED_EDGE = ( + "edge E: A -> B {\n" + " coco_key: String // no coco_managed_by_og here, only below\n" + " coco_managed_by_og: Bool?\n" + "}\n" + ) + ANOTHER_APPS_EDGE = OWNED_EDGE.replace( + "coco_managed_by_og: Bool?", "coco_managed_by_other: Bool?" + ) + + def test_a_comment_does_not_confer_ownership(self) -> None: + assert ownership_marker(self.USER_EDGE, "edge", "E") is None + assert ( + release_ownership(self.USER_EDGE, "edge", "E", "coco_managed_by_og") + == self.USER_EDGE + ) + + def test_a_declaration_confers_ownership_whatever_the_comments_say(self) -> None: + assert ownership_marker(self.OWNED_EDGE, "edge", "E") == "coco_managed_by_og" + assert release_ownership( + self.OWNED_EDGE, "edge", "E", "coco_managed_by_og" + ) == ( + "edge E: A -> B {\n" + " coco_key: String // no coco_managed_by_og here, only below\n" + "}\n" + ) + + def test_the_marker_names_its_app(self) -> None: + """A release only ever drops this app's own marker: another app's + stays, so its drop still finds its type marked as its own.""" + assert ( + ownership_marker(self.ANOTHER_APPS_EDGE, "edge", "E") + == "coco_managed_by_other" + ) + assert ( + release_ownership(self.ANOTHER_APPS_EDGE, "edge", "E", "coco_managed_by_og") + == self.ANOTHER_APPS_EDGE + ) + + +class TestNodeMutations: + def test_upsert_is_a_statement_with_positional_slots(self) -> None: + """A builder returns a statement, not query text: its body holds a + `$?` slot per bound value, and the binds carry each slot's label, + type and value. Names are allocated only when a query is rendered, + so the value never touches the text and no renaming is ever needed + to combine statements.""" + m = build_node_upsert( + "Person", + [ + PropertyValue("email", "String", "ada@x.com"), + PropertyValue("display_name", "String", "Ada"), + ], + coco_key="k1", + ) + assert m == Statement( + "insert Person { email: $?, display_name: $?, coco_key: $? }", + ( + Bind("p_email", "String", "ada@x.com"), + Bind("p_display_name", "String", "Ada"), + Bind("p_coco_key", "String", "k1"), + ), + ) + rendered = render_query([m]) + assert rendered.expr == ( + "query m($s0_p_email: String, $s0_p_display_name: String, " + "$s0_p_coco_key: String) { insert Person { email: $s0_p_email, " + "display_name: $s0_p_display_name, coco_key: $s0_p_coco_key } }" + ) + assert rendered.params == { + "s0_p_email": "ada@x.com", + "s0_p_display_name": "Ada", + "s0_p_coco_key": "k1", + } + + def test_delete_uses_coco_key_single_predicate(self) -> None: + m = render_query([build_node_delete("Person", coco_key="k1")]) + assert m.expr == ( + "query m($s0_p_coco_key: String) " + "{ delete Person where coco_key = $s0_p_coco_key }" + ) + assert m.params == {"s0_p_coco_key": "k1"} + + def test_endpoint_stub_is_key_plus_coco_key(self) -> None: + m = render_query( + [ + build_endpoint_stub( + "Person", [PropertyValue("email", "String", "ada@x.com")], "k1" + ) + ] + ) + assert m.expr == ( + "query m($s0_p_email: String, $s0_p_coco_key: String) " + "{ insert Person { email: $s0_p_email, coco_key: $s0_p_coco_key } }" + ) + + def test_value_is_never_interpolated(self) -> None: + # A value containing GQ-syntax characters must not change the query's + # shape at all β€” asserting the full string, not just `not in`, is + # what actually proves the value never reaches the query text. + nasty = '", role: "pwned' + m = render_query( + [ + build_node_upsert( + "Person", [PropertyValue("email", "String", nasty)], coco_key="k1" + ) + ] + ) + assert m.expr == ( + "query m($s0_p_email: String, $s0_p_coco_key: String) " + "{ insert Person { email: $s0_p_email, coco_key: $s0_p_coco_key } }" + ) + assert nasty not in m.expr + assert m.params == {"s0_p_email": nasty, "s0_p_coco_key": "k1"} + + def test_type_is_never_interpolated(self) -> None: + # Unlike `value`, `pg_type` lands directly in the signature text β€” it + # must be validated, not merely passed through unexamined. + evil = ( + "String) { delete Person where coco_key = $p_coco_key } query z($x: String" + ) + with pytest.raises(ValueError, match="Invalid Omnigraph property type"): + build_node_upsert( + "Person", [PropertyValue("email", evil, "v")], coco_key="k1" + ) + + def test_empty_pg_type_rejected(self) -> None: + with pytest.raises(ValueError, match="Invalid Omnigraph property type"): + build_node_upsert( + "Person", [PropertyValue("email", "", "v")], coco_key="k1" + ) + + def test_bad_property_name_rejected(self) -> None: + with pytest.raises(ValueError, match="Invalid Omnigraph property name"): + build_node_upsert( + "Person", [PropertyValue("a b", "String", 1)], coco_key="k1" + ) + + def test_coco_key_named_prop_rejected(self) -> None: + with pytest.raises(ValueError, match="reserved"): + build_node_upsert( + "Person", [PropertyValue(COCO_KEY, "String", "x")], coco_key="k1" + ) + + def test_duplicate_property_names_rejected(self) -> None: + with pytest.raises(ValueError, match="Duplicate"): + build_node_upsert( + "Person", + [ + PropertyValue("email", "String", "a"), + PropertyValue("email", "String", "b"), + ], + coco_key="k1", + ) + + +class TestEdgeMutations: + def test_insert_carries_coco_key_not_id(self) -> None: + m = render_query( + [ + build_edge_insert( + "WorksAt", + PropertyValue("from", "String", "ada@x.com"), + PropertyValue("to", "String", "acme"), + [PropertyValue("role", "String", "Eng")], + coco_key="e1", + ) + ] + ) + assert m.expr == ( + "query m($s0_e_from: String, $s0_e_to: String, $s0_p_role: String, " + "$s0_p_coco_key: String) { insert WorksAt { from: $s0_e_from, " + "to: $s0_e_to, role: $s0_p_role, coco_key: $s0_p_coco_key } }" + ) + assert "id:" not in m.expr + assert m.params == { + "s0_e_from": "ada@x.com", + "s0_e_to": "acme", + "s0_p_role": "Eng", + "s0_p_coco_key": "e1", + } + + def test_insert_without_properties(self) -> None: + m = render_query( + [ + build_edge_insert( + "Supports", + PropertyValue("from", "String", "load-test"), + PropertyValue("to", "String", "lower-latency"), + [], + coco_key="e2", + ) + ] + ) + assert m.expr == ( + "query m($s0_e_from: String, $s0_e_to: String, $s0_p_coco_key: String) " + "{ insert Supports { from: $s0_e_from, to: $s0_e_to, " + "coco_key: $s0_p_coco_key } }" + ) + + def test_delete_uses_coco_key(self) -> None: + m = render_query([build_edge_delete("WorksAt", coco_key="e1")]) + assert m.expr == ( + "query m($s0_p_coco_key: String) " + "{ delete WorksAt where coco_key = $s0_p_coco_key }" + ) + + def test_no_edge_update_builder_exists(self) -> None: + import cocoindex.connectors.omnigraph._gq as gq + + assert not hasattr(gq, "build_edge_update") + + def test_endpoint_type_is_never_interpolated(self) -> None: + evil = ( + "String) { delete WorksAt where coco_key = $p_coco_key } query z($x: String" + ) + with pytest.raises(ValueError, match="Invalid Omnigraph property type"): + build_edge_insert( + "WorksAt", + PropertyValue("from", evil, "ada@x.com"), + PropertyValue("to", "String", "acme"), + [], + coco_key="e1", + ) + + def test_coco_key_named_prop_rejected(self) -> None: + with pytest.raises(ValueError, match="reserved"): + build_edge_insert( + "WorksAt", + PropertyValue("from", "String", "a"), + PropertyValue("to", "String", "b"), + [PropertyValue(COCO_KEY, "String", "x")], + coco_key="e1", + ) + + def test_from_named_prop_rejected(self) -> None: + with pytest.raises(ValueError, match="reserved"): + build_edge_insert( + "WorksAt", + PropertyValue("from", "String", "a"), + PropertyValue("to", "String", "b"), + [PropertyValue("from", "String", "x")], + coco_key="e1", + ) + + def test_to_named_prop_rejected(self) -> None: + with pytest.raises(ValueError, match="reserved"): + build_edge_insert( + "WorksAt", + PropertyValue("from", "String", "a"), + PropertyValue("to", "String", "b"), + [PropertyValue("to", "String", "x")], + coco_key="e1", + ) + + def test_duplicate_property_names_rejected(self) -> None: + with pytest.raises(ValueError, match="Duplicate"): + build_edge_insert( + "WorksAt", + PropertyValue("from", "String", "a"), + PropertyValue("to", "String", "b"), + [ + PropertyValue("role", "String", "x"), + PropertyValue("role", "String", "y"), + ], + coco_key="e1", + ) + + +class TestValidatePgType: + @pytest.mark.parametrize( + "pg_type", + [ + "String", + "Bool", + "I32", + "I64", + "U32", + "U64", + "F32", + "F64", + "Date", + "DateTime", + "Blob", + "String?", + "I32?", + "DateTime?", + "Vector(768)", + "Vector(3)?", + "[String]", + "[I32]?", + "enum(a, b, c)", + "enum(a,b,c)?", + "enum(only_one)", + ], + ) + def test_valid(self, pg_type: str) -> None: + validate_pg_type(pg_type) # must not raise + + @pytest.mark.parametrize( + "pg_type", + [ + "", + "String) { delete Person where coco_key = $p_coco_key } query z($x: String", + "Strin g", + "Vector()", + "Vector(abc)", + "Vector(0)", + "[Bogus]", + "[String", + "enum()", + "enum(1abc)", + "Foo", + "string", + "String??", + ], + ) + def test_invalid(self, pg_type: str) -> None: + with pytest.raises(ValueError, match="Invalid Omnigraph property type"): + validate_pg_type(pg_type) + + +class TestDeriveCocoKey: + def test_deterministic(self) -> None: + from cocoindex.connectors.omnigraph._target import derive_coco_key + + assert derive_coco_key(("a", "b")) == derive_coco_key(("a", "b")) + + def test_order_matters(self) -> None: + from cocoindex.connectors.omnigraph._target import derive_coco_key + + assert derive_coco_key(("a", "b")) != derive_coco_key(("b", "a")) + + def test_no_underscore_collision(self) -> None: + """The neo4j connector's f"{a}_{b}_{c}_{d}" form collides here; a + fingerprint over the tuple must not. See spec decision 4.""" + from cocoindex.connectors.omnigraph._target import derive_coco_key + + assert derive_coco_key(("x", "y_z")) != derive_coco_key(("x_y", "z")) + + def test_type_is_significant(self) -> None: + from cocoindex.connectors.omnigraph._target import derive_coco_key + + assert derive_coco_key((1, 2)) != derive_coco_key(("1", "2")) + + def test_single_part_key(self) -> None: + from cocoindex.connectors.omnigraph._target import derive_coco_key + + assert derive_coco_key(("slug-1",)) != derive_coco_key(("slug-2",)) + + def test_is_hex(self) -> None: + from cocoindex.connectors.omnigraph._target import derive_coco_key + + k = derive_coco_key(("a", "b")) + int(k, 16) + assert k.islower() + + +class TestTypeMapping: + @pytest.mark.asyncio + async def test_scalar_mapping(self) -> None: + schema = await NodeSchema.from_class(_Doc, key="slug") + assert {n: p.pg_type for n, p in schema.properties.items()} == { + "slug": "String", + "title": "String", + "words": "I64", + "ratio": "F64", + "published": "Bool", + "on": "Date", + "at": "DateTime", + "note": "String?", + } + + @pytest.mark.asyncio + async def test_date_and_datetime_get_isoformat_encoders(self) -> None: + """`PropertyDef.encoder` is what lets a `date`/`datetime` field reach + the transport layer without crashing `json.dumps` β€” verified against + the engine that `Date` accepts `"2026-01-01"` and `DateTime` accepts + ISO with or without a trailing `Z`, so `.isoformat()` suffices.""" + schema = await NodeSchema.from_class(_Doc, key="slug") + on = datetime.date(2026, 1, 1) + at = datetime.datetime(2026, 1, 1, 12, 30) # noqa: DTZ001 + assert schema.properties["on"].encoder is not None + assert schema.properties["on"].encoder(on) == on.isoformat() + assert schema.properties["at"].encoder is not None + assert schema.properties["at"].encoder(at) == at.isoformat() + assert schema.properties["slug"].encoder is None + assert schema.properties["note"].encoder is None + + def test_explicit_date_and_datetime_definitions_encode_by_type(self) -> None: + """A hand-built `PropertyDef("at", "DateTime")` has no encoder set, + and encoding was keyed on that field: a `datetime` then reached + `json.dumps` raw, and a `date` key reached the engine's StableKey + raw. The built-in ISO encoding follows the `.pg` type, whether the + definition came from a dataclass or was written out by hand.""" + at = datetime.datetime(2026, 1, 1, 12, 30, tzinfo=datetime.UTC) + assert ogt._encode_property(PropertyDef("at", "DateTime"), at).value == ( + at.isoformat() + ) + assert ( + ogt._encode_property( + PropertyDef("on", "Date?"), datetime.date(2026, 1, 5) + ).value + == "2026-01-05" + ) + assert ogt._encode_property( + PropertyDef("days", "[Date]"), [datetime.date(2026, 1, 5)] + ).value == ["2026-01-05"] + assert ( + ogt._tracking_key_value( + PropertyDef("on", "Date"), datetime.date(2026, 1, 5) + ) + == "2026-01-05" + ) + + def test_isoformat_given_explicitly_counts_as_the_built_in_key_encoder( + self, + ) -> None: + """The obvious thing to write by hand is the same thing the connector + applies by default; it must not be refused as a custom key encoder.""" + for pg_type, encoder in ( + ("DateTime", datetime.datetime.isoformat), + ("Date", datetime.date.isoformat), + ): + schema = NodeSchema( + properties={"k": PropertyDef("k", pg_type, encoder)}, key=("k",) + ) + + def mount(schema: NodeSchema = schema) -> None: + omnigraph.node_target(_fresh_db("iso"), "T", schema) + + _in_component(mount) + + @pytest.mark.asyncio + async def test_list_of_dates_encodes_every_element(self) -> None: + """`list[datetime.date]` maps to `[Date]`, which the schema accepts β€” + but `json.dumps` cannot serialize the elements, so the encoder has + to reach into the list. A nullable list stays `None` when absent.""" + + @dataclass + class _L: + slug: str + days: list[datetime.date] + maybe: list[datetime.datetime] | None + + schema = await NodeSchema.from_class(_L, key="slug") + assert schema.properties["days"].pg_type == "[Date]" + assert schema.properties["maybe"].pg_type == "[DateTime]?" + encoded = ogt._encode_properties( + {"slug": "a", "days": [datetime.date(2026, 1, 1)], "maybe": None}, + schema.properties, + ) + by_name = {p.name: p.value for p in encoded} + assert by_name["days"] == ["2026-01-01"] + assert by_name["maybe"] is None + json.dumps(by_name) + + @pytest.mark.asyncio + async def test_key_normalised_to_tuple(self) -> None: + schema = await NodeSchema.from_class(_Doc, key="slug") + assert schema.key == ("slug",) + + @pytest.mark.asyncio + async def test_composite_key_raises(self) -> None: + """Same engine limit `_gq.render_node_type` enforces (one `@key` per + node type), caught here so the error names the dataclass rather than + surfacing only later out of `.render()`.""" + + @dataclass + class _R: + sensor: str + at: datetime.datetime + v: float + + with pytest.raises(ValueError, match="composite key"): + await NodeSchema.from_class(_R, key=["sensor", "at"]) + + @pytest.mark.asyncio + async def test_id_field_raises(self) -> None: + """`id` is reserved by Omnigraph whether the schema ends up + describing a node or an edge, and it's the likeliest collision in a + real dataclass β€” so it's rejected here, naming the class.""" + + @dataclass + class _WithId: + id: str + title: str + + with pytest.raises(ValueError, match=r"_WithId\.id: 'id' is reserved"): + await NodeSchema.from_class(_WithId, key="id") + + @pytest.mark.asyncio + async def test_empty_key_raises(self) -> None: + """An unkeyed node `insert` is a strict insert, not an upsert β€” every + re-run would duplicate every node. Verified against the live engine: + inserting the same unkeyed node twice yields two rows.""" + + with pytest.raises(ValueError, match="at least one key property"): + await NodeSchema.from_class(_Doc, key=[]) + + @pytest.mark.asyncio + async def test_nullable_key_raises_at_from_class(self) -> None: + """Same rule `_gq.render_node_type` enforces, but caught here so the + error names the dataclass and field the user actually got wrong, + rather than surfacing only later out of `.render()`.""" + + with pytest.raises( + ValueError, match=r"key property 'note' of _Doc must not be nullable" + ): + await NodeSchema.from_class(_Doc, key="note") + + @pytest.mark.asyncio + async def test_annotation_override(self) -> None: + @dataclass + class _S: + slug: str + small: Annotated[int, OmnigraphType("I32")] + + schema = await NodeSchema.from_class(_S, key="slug") + assert schema.properties["small"].pg_type == "I32" + + @pytest.mark.asyncio + async def test_omnigraph_type_override_is_validated(self) -> None: + """`OmnigraphType` is a security boundary: its `pg_type` is spliced + directly into generated query text downstream, so a crafted value + must be rejected here rather than silently accepted into a schema.""" + + @dataclass + class _Evil: + slug: str + small: Annotated[ + str, + OmnigraphType("String) { delete X where y = $z } query w($a: String"), + ] + + with pytest.raises(ValueError, match="Invalid Omnigraph property type"): + await NodeSchema.from_class(_Evil, key="slug") + + @pytest.mark.asyncio + async def test_nullable_list_element_raises(self) -> None: + """`[I64?]` is rejected by the engine with "expected core_type".""" + + @dataclass + class _N: + slug: str + vals: list[int | None] + + with pytest.raises(TypeError, match="non-nullable scalars"): + await NodeSchema.from_class(_N, key="slug") + + @pytest.mark.asyncio + async def test_nested_list_raises(self) -> None: + """`[[String]]` is rejected by the engine with "expected base_type".""" + + @dataclass + class _NN: + slug: str + vals: list[list[str]] + + with pytest.raises(TypeError, match="non-nullable scalars"): + await NodeSchema.from_class(_NN, key="slug") + + @pytest.mark.asyncio + async def test_unmapped_type_raises(self) -> None: + @dataclass + class _U: + slug: str + weird: complex + + with pytest.raises(TypeError, match="no Omnigraph type mapping for"): + await NodeSchema.from_class(_U, key="slug") + + @pytest.mark.asyncio + async def test_bytes_has_no_mapping(self) -> None: + """Omnigraph's `Blob` is an external URI reference (the engine fetches + a `file://` value), not inline bytes β€” a Python `bytes` value can + never be a `Blob`, so it must fall through to the generic + no-mapping error rather than silently being declared as one.""" + + @dataclass + class _B: + slug: str + payload: bytes + + with pytest.raises(TypeError, match="no Omnigraph type mapping for"): + await NodeSchema.from_class(_B, key="slug") + + @pytest.mark.asyncio + async def test_render(self) -> None: + @dataclass + class _Src: + slug: str + title: str + + schema = await NodeSchema.from_class(_Src, key="slug") + assert schema.render("Source", owner="og") == ( + "node Source {\n slug: String @key\n title: String\n" + " coco_key: String\n coco_managed_by_og: Bool?\n}" + ) + + +def _spec(schema: NodeSchema, managed_by: ManagedBy = ManagedBy.SYSTEM) -> _TypeSpec: + return _TypeSpec( + schema=schema, + key=schema.key, + from_type=None, + to_type=None, + managed_by=managed_by, + owner="og", + ) + + +class TestEdgeSchema: + @pytest.mark.asyncio + async def test_from_class_needs_no_key(self) -> None: + """An edge's identity is always `(from_id, to_id)`; its own properties + never form a key. Building edge properties through + `NodeSchema.from_class` forced a meaningless `key=` onto every edge + dataclass, which the example then had to explain away.""" + schema = await EdgeSchema.from_class(_AttendedRel) + assert schema.properties == { + "is_organizer": PropertyDef("is_organizer", "Bool") + } + assert not hasattr(schema, "key") + + @pytest.mark.asyncio + async def test_from_class_rejects_id(self) -> None: + @dataclass + class _WithId: + id: str + + with pytest.raises(ValueError, match="'id' is reserved"): + await EdgeSchema.from_class(_WithId) + + def test_a_node_schema_is_not_an_edge_schema(self) -> None: + """Keyed schemas describe node types only; handing one to an edge + type is refused at mount, where the mistake is.""" + node_schema = NodeSchema( + properties={"slug": PropertyDef("slug", "String")}, key=("slug",) + ) + + def mount_with_a_node_schema() -> None: + omnigraph.edge_target( + _fresh_db("node_schema"), + "E", + _bare_node_target(node_schema, "A"), + _bare_node_target(node_schema, "B"), + node_schema, # type: ignore[arg-type] + ) + + with pytest.raises(TypeError, match="EdgeSchema"): + _in_component(mount_with_a_node_schema) + + +class TestNodeTypeReconcile: + @pytest.mark.asyncio + async def test_absent_creates(self) -> None: + schema = await NodeSchema.from_class(_Doc, key="slug") + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), _spec(schema), [], True + ) + assert out is not None + assert out.action.main_action == "insert" + assert out.child_invalidation is None + + @pytest.mark.asyncio + async def test_unchanged_is_noop(self) -> None: + """A root provider that owns children must never return bare + `None` for "no change needed" β€” the engine only refreshes a + child's handler when reconcile()'s own output carries a fresh + ChildTargetDef, so bare None here would starve this type's own + node/edge children on the very next unchanged run.""" + schema = await NodeSchema.from_class(_Doc, key="slug") + prev = _type_tracking_record_from_spec(_spec(schema)) + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), _spec(schema), [prev], False + ) + assert out is not None + assert out.action.main_action is None and not out.action.property_actions + assert out.child_invalidation is None + + @pytest.mark.asyncio + async def test_property_added_is_additive(self) -> None: + @dataclass + class _V1: + slug: str + + @dataclass + class _V2: + slug: str + title: str | None + + prev = _type_tracking_record_from_spec( + _spec(await NodeSchema.from_class(_V1, key="slug")) + ) + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(await NodeSchema.from_class(_V2, key="slug")), + [prev], + False, + ) + assert out is not None + assert out.action.main_action is None + assert out.action.property_actions == {"prop:title": "insert"} + assert out.child_invalidation is None + + @pytest.mark.asyncio + async def test_non_nullable_property_added_raises(self) -> None: + """Engine rejects this via schema apply (OG-MF-103); fail in Python instead.""" + + @dataclass + class _V1: + slug: str + + @dataclass + class _V2: + slug: str + title: str # non-nullable addition + + prev = _type_tracking_record_from_spec( + _spec(await NodeSchema.from_class(_V1, key="slug")) + ) + with pytest.raises(ValueError, match="non-nullable"): + _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(await NodeSchema.from_class(_V2, key="slug")), + [prev], + False, + ) + + @pytest.mark.asyncio + async def test_nullable_property_added_is_additive(self) -> None: + @dataclass + class _V1: + slug: str + + @dataclass + class _V2: + slug: str + title: str | None + + prev = _type_tracking_record_from_spec( + _spec(await NodeSchema.from_class(_V1, key="slug")) + ) + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(await NodeSchema.from_class(_V2, key="slug")), + [prev], + False, + ) + assert out is not None + assert out.action.main_action is None + assert out.action.property_actions == {"prop:title": "insert"} + assert out.child_invalidation is None + + @pytest.mark.asyncio + async def test_property_dropped_is_lossy(self) -> None: + @dataclass + class _V1: + slug: str + title: str + + @dataclass + class _V2: + slug: str + + prev = _type_tracking_record_from_spec( + _spec(await NodeSchema.from_class(_V1, key="slug")) + ) + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(await NodeSchema.from_class(_V2, key="slug")), + [prev], + False, + ) + assert out is not None + assert out.child_invalidation == "lossy" + + def test_encoder_change_is_lossy(self) -> None: + """A changed encoder rewrites what the graph holds for every row, but + the rendered `.pg` is byte-identical, so the only way a memoized + declaring component ever learns of it is through the type's own + tracking record: the change has to invalidate the children like a + retyped property does. A hand-built definition without an encoder + and one carrying the built-in encoder for its type are the same + encoding and must not count as a change.""" + + def schema(encoder: ogt.ValueEncoder | None) -> NodeSchema: + return NodeSchema( + properties={ + "slug": PropertyDef("slug", "String"), + "name": PropertyDef("name", "String", encoder), + "at": PropertyDef("at", "DateTime", ogt._ENCODERS["DateTime"]), + }, + key=("slug",), + ) + + prev = _type_tracking_record_from_spec(_spec(schema(str.lower))) + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), _spec(schema(str.upper)), [prev], False + ) + assert out is not None + assert out.action.main_action is None + assert out.child_invalidation == "lossy" + + by_hand = NodeSchema( + properties={ + "slug": PropertyDef("slug", "String"), + "name": PropertyDef("name", "String", str.lower), + "at": PropertyDef("at", "DateTime"), + }, + key=("slug",), + ) + assert ( + _type_tracking_record_from_spec(_spec(by_hand)).tracking_record.sub + == prev.tracking_record.sub + ) + + @pytest.mark.asyncio + async def test_key_change_is_destructive(self) -> None: + @dataclass + class _V: + slug: str + title: str + + prev = _type_tracking_record_from_spec( + _spec(await NodeSchema.from_class(_V, key="slug")) + ) + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(await NodeSchema.from_class(_V, key="title")), + [prev], + False, + ) + assert out is not None + assert out.action.main_action == "replace" + assert out.child_invalidation == "destructive" + + @pytest.mark.asyncio + async def test_user_managed_schema_drift_is_tracked_not_rejected(self) -> None: + """User-managed means the schema is the user's: they migrate it with + `omnigraph schema apply`, then declare the wider dataclass. The + connector used to compare that declaration against what it had + tracked and refuse β€” advising exactly that migration, which it then + rejected again on every run, because the check read tracking + history and never the live schema. No validation, no schema write: + the new declaration is simply tracked from here on. That holds even + for a non-nullable addition β€” whether the user's migration was legal + is the engine's call, made when they applied it.""" + + @dataclass + class _V1: + slug: str + + @dataclass + class _V2: + slug: str + title: str # non-nullable addition + + prev = _type_tracking_record_from_spec( + _spec(await NodeSchema.from_class(_V1, key="slug"), ManagedBy.USER) + ) + spec_v2 = _spec(await NodeSchema.from_class(_V2, key="slug"), ManagedBy.USER) + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), spec_v2, [prev], False + ) + assert out is not None + assert out.action.main_action is None and not out.action.property_actions + assert out.tracking_record == _type_tracking_record_from_spec(spec_v2) + + @pytest.mark.asyncio + async def test_handing_a_type_to_the_user_releases_ownership( + self, + ) -> None: + """A block the connector rendered declares `coco_managed_by_`. Once the + app declares the type `managed_by=user` that is stale β€” and a later + drop of a node type read it as current ownership, taking the user's + edge type (and its edges) along. The handoff has to write once to + drop the property; after that the block is the user's.""" + schema = await NodeSchema.from_class(_Doc, key="slug") + prev = _type_tracking_record_from_spec(_spec(schema, ManagedBy.SYSTEM)) + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), _spec(schema, ManagedBy.USER), [prev], False + ) + assert out is not None + assert out.action.main_action is None and not out.action.property_actions + assert out.action.release_ownership is True + + @pytest.mark.asyncio + async def test_taking_a_type_back_from_the_user_rewrites_its_block(self) -> None: + """The reverse handoff. A system-managed declaration after a + user-managed record wrote nothing when the schema was otherwise + unchanged, so the ownership property stayed absent and the next drop of a + node type it referenced treated it as the user's. Reclaiming must + re-render the block even with nothing else to change.""" + schema = await NodeSchema.from_class(_Doc, key="slug") + prev = _type_tracking_record_from_spec(_spec(schema, ManagedBy.USER)) + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(schema, ManagedBy.SYSTEM), + [prev], + False, + ) + assert out is not None + assert out.action.main_action == "upsert" + assert out.action.release_ownership is False + + @pytest.mark.asyncio + async def test_a_type_already_user_managed_releases_nothing(self) -> None: + schema = await NodeSchema.from_class(_Doc, key="slug") + prev = _type_tracking_record_from_spec(_spec(schema, ManagedBy.USER)) + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), _spec(schema, ManagedBy.USER), [prev], False + ) + assert out is not None + assert out.action.release_ownership is False + + @pytest.mark.asyncio + async def test_user_managed_first_run_adopts(self) -> None: + """The first run of a managed_by=user type has NO tracking records, + so the tracked diff is empty and there is nothing to disagree with. + This used to raise, making the mode unusable on run one even against + a perfectly matching graph, and with an error text ("differs from the + tracked one") that misdescribed the cause. Nothing is tracked yet; + there is nothing to differ from.""" + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(await NodeSchema.from_class(_Doc, key="slug"), ManagedBy.USER), + [], + True, + ) + assert out is not None + # No schema write, but a child handler all the same -- node/edge + # upserts must still flow into a user-managed type. + assert out.action.main_action is None and not out.action.property_actions + assert out.child_invalidation is None + + @pytest.mark.asyncio + async def test_user_managed_first_run_adopts_with_a_plain_string(self) -> None: + """`ManagedBy` is a StrEnum, so `managed_by="user"` compares equal but + is NOT identical. Under an `is` check it fell through to full SYSTEM + management β€” the connector would rewrite a schema the user declared + they own. Same expectation as the enum case above.""" + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec( + await NodeSchema.from_class(_Doc, key="slug"), cast(ManagedBy, "user") + ), + [], + True, + ) + assert out is not None + assert out.action.main_action is None and not out.action.property_actions + + @pytest.mark.asyncio + async def test_user_managed_unchanged_is_noop(self) -> None: + """The second run: tracked, unchanged, still no schema write. This is + the second of the two bare-`None` sites `_noop_output` replaced, and + it had no coverage.""" + spec = _spec(await NodeSchema.from_class(_Doc, key="slug"), ManagedBy.USER) + prev = _type_tracking_record_from_spec(spec) + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), spec, [prev], False + ) + assert out is not None + assert out.action.main_action is None and not out.action.property_actions + + @pytest.mark.asyncio + async def test_undeclaring_a_user_managed_type_does_not_drop_it(self) -> None: + """Un-declaring a `managed_by=user` type must be a no-op, not a drop. + + The connector did not create the type and does not own it, so pulling + the mount out of the app is a statement about the app, not about the + graph. Dropping the block here deletes the user's type AND every row + in it -- the one irreversible thing this connector can do to data it + was explicitly told it does not manage. `omnigraph.mdx` promises + exactly this ("Removing a `managed_by="user"` type from your app is + also never destructive") and nothing enforced it: the drop path never + consulted ownership, and structurally could not, because the desired + state is NON_EXISTENCE and the tracking record did not persist + `managed_by` at all. + """ + spec = _spec(await NodeSchema.from_class(_Doc, key="slug"), ManagedBy.USER) + prev = _type_tracking_record_from_spec(spec) + assert ( + _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), coco.NON_EXISTENCE, [prev], False + ) + is None + ) + + @pytest.mark.asyncio + async def test_multiple_divergent_prev_records_forces_action(self) -> None: + """`prev_possible_records` can hold more than one candidate after an + interrupted update. Picking one arbitrarily (as opposed to requiring + every candidate to agree with `desired` before treating state as + converged) can silently skip reconciliation: here `desired` matches + the first candidate exactly but diverges from the second, so an + action must still be emitted.""" + + @dataclass + class _Narrow: + slug: str + + @dataclass + class _Wide: + slug: str + title: str + + narrow_schema = await NodeSchema.from_class(_Narrow, key="slug") + wide_schema = await NodeSchema.from_class(_Wide, key="slug") + matching_prev = _type_tracking_record_from_spec(_spec(narrow_schema)) + divergent_prev = _type_tracking_record_from_spec(_spec(wide_schema)) + + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(narrow_schema), + [matching_prev, divergent_prev], + False, + ) + assert out is not None + assert out.action.main_action is None + assert out.action.property_actions == {"prop:title": "delete"} + assert out.child_invalidation == "lossy" + + +def _edge_spec( + schema: NodeSchema, + from_type: str = "Source", + to_type: str = "Claim", + managed_by: ManagedBy = ManagedBy.SYSTEM, +) -> _TypeSpec: + return _TypeSpec( + schema=schema, + key=(), + from_type=from_type, + to_type=to_type, + managed_by=managed_by, + owner="og", + ) + + +class TestEdgeTypeReconcile: + @pytest.mark.asyncio + async def test_absent_creates(self) -> None: + @dataclass + class _Props: + role: str + + schema = await NodeSchema.from_class(_Props, key="role") + out = _EdgeTypeHandler().reconcile( + _TypeKey("og", "edge", "WorksAt"), _edge_spec(schema), [], True + ) + assert out is not None + assert out.action.main_action == "insert" + assert out.child_invalidation is None + + @pytest.mark.asyncio + async def test_unchanged_is_noop(self) -> None: + """See TestNodeTypeReconcile.test_unchanged_is_noop: bare None + here would starve this edge type's own child entities.""" + + @dataclass + class _Props: + role: str + + schema = await NodeSchema.from_class(_Props, key="role") + prev = _type_tracking_record_from_spec(_edge_spec(schema)) + out = _EdgeTypeHandler().reconcile( + _TypeKey("og", "edge", "WorksAt"), _edge_spec(schema), [prev], False + ) + assert out is not None + assert out.action.main_action is None and not out.action.property_actions + assert out.child_invalidation is None + + @pytest.mark.asyncio + async def test_property_added_is_additive(self) -> None: + @dataclass + class _V1: + role: str + + @dataclass + class _V2: + role: str + since: str | None + + prev = _type_tracking_record_from_spec( + _edge_spec(await NodeSchema.from_class(_V1, key="role")) + ) + out = _EdgeTypeHandler().reconcile( + _TypeKey("og", "edge", "WorksAt"), + _edge_spec(await NodeSchema.from_class(_V2, key="role")), + [prev], + False, + ) + assert out is not None + assert out.action.main_action is None + assert out.action.property_actions == {"prop:since": "insert"} + assert out.child_invalidation is None + + @pytest.mark.asyncio + async def test_non_nullable_property_added_raises(self) -> None: + @dataclass + class _V1: + role: str + + @dataclass + class _V2: + role: str + since: str # non-nullable addition + + prev = _type_tracking_record_from_spec( + _edge_spec(await NodeSchema.from_class(_V1, key="role")) + ) + with pytest.raises(ValueError, match="non-nullable"): + _EdgeTypeHandler().reconcile( + _TypeKey("og", "edge", "WorksAt"), + _edge_spec(await NodeSchema.from_class(_V2, key="role")), + [prev], + False, + ) + + @pytest.mark.asyncio + async def test_property_dropped_is_lossy(self) -> None: + @dataclass + class _V1: + role: str + since: str + + @dataclass + class _V2: + role: str + + prev = _type_tracking_record_from_spec( + _edge_spec(await NodeSchema.from_class(_V1, key="role")) + ) + out = _EdgeTypeHandler().reconcile( + _TypeKey("og", "edge", "WorksAt"), + _edge_spec(await NodeSchema.from_class(_V2, key="role")), + [prev], + False, + ) + assert out is not None + assert out.child_invalidation == "lossy" + + @pytest.mark.asyncio + async def test_from_type_change_is_destructive(self) -> None: + @dataclass + class _Props: + role: str + + schema = await NodeSchema.from_class(_Props, key="role") + prev = _type_tracking_record_from_spec( + _edge_spec(schema, from_type="Source", to_type="Claim") + ) + out = _EdgeTypeHandler().reconcile( + _TypeKey("og", "edge", "WorksAt"), + _edge_spec(schema, from_type="Person", to_type="Claim"), + [prev], + False, + ) + assert out is not None + assert out.action.main_action == "replace" + assert out.child_invalidation == "destructive" + + @pytest.mark.asyncio + async def test_to_type_change_is_destructive(self) -> None: + @dataclass + class _Props: + role: str + + schema = await NodeSchema.from_class(_Props, key="role") + prev = _type_tracking_record_from_spec( + _edge_spec(schema, from_type="Source", to_type="Claim") + ) + out = _EdgeTypeHandler().reconcile( + _TypeKey("og", "edge", "WorksAt"), + _edge_spec(schema, from_type="Source", to_type="Report"), + [prev], + False, + ) + assert out is not None + assert out.action.main_action == "replace" + assert out.child_invalidation == "destructive" + + +class TestTypeOwnershipMatrix: + """Every combination of tracked ownership and declared ownership. + + Twelve cases, small enough to be exhaustive: the previous state is + nothing / system-managed / user-managed / one of each (an interrupted + update leaves several candidates behind), crossed with a declaration that + is system-managed, user-managed, or absent. None of this was answerable + before `managed_by` started riding on the persisted tracking record. + """ + + @staticmethod + async def _prev( + ownerships: list[ManagedBy], + ) -> list[ogt._TypeTrackingRecord]: + """Previous records that differ ONLY in ownership, so each case below + isolates the ownership decision from any schema difference.""" + schema = await NodeSchema.from_class(_Doc, key="slug") + return [_type_tracking_record_from_spec(_spec(schema, m)) for m in ownerships] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("prev_ownerships", "writes"), + [ + ([], True), + ([ManagedBy.SYSTEM], False), + ([ManagedBy.USER], True), + ([ManagedBy.SYSTEM, ManagedBy.USER], True), + ], + ids=["no-prev", "system", "user", "mixed"], + ) + async def test_declared_system_managed( + self, prev_ownerships: list[ManagedBy], writes: bool + ) -> None: + """A system-managed declaration writes DDL when the graph might not + already match it: on a first run, where nothing is tracked and the + engine reports `prev_may_be_missing`, and whenever any tracked record + was user-managed β€” the block then lacks its ownership property, and only a + re-render puts it back. A system-managed record already carrying this + exact schema means there is nothing to apply.""" + schema = await NodeSchema.from_class(_Doc, key="slug") + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(schema, ManagedBy.SYSTEM), + await self._prev(prev_ownerships), + not prev_ownerships, + ) + assert out is not None + assert bool(out.action.main_action or out.action.property_actions) is writes + + @pytest.mark.asyncio + async def test_user_to_system_handoff_applies_a_schema_change(self) -> None: + @dataclass + class _V1: + slug: str + + @dataclass + class _V2: + slug: str + title: str | None + + old_schema = await NodeSchema.from_class(_V1, key="slug") + new_schema = await NodeSchema.from_class(_V2, key="slug") + prev = _type_tracking_record_from_spec(_spec(old_schema, ManagedBy.USER)) + + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(new_schema, ManagedBy.SYSTEM), + [prev], + False, + ) + + assert out is not None + assert out.action.main_action is None + assert out.action.property_actions == {"prop:title": "insert"} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "prev_ownerships", + [[], [ManagedBy.SYSTEM], [ManagedBy.USER], [ManagedBy.SYSTEM, ManagedBy.USER]], + ids=["no-prev", "system", "user", "mixed"], + ) + async def test_declared_user_managed( + self, prev_ownerships: list[ManagedBy] + ) -> None: + """A user-managed declaration never writes DDL, whatever is tracked -- + including a type this connector used to manage itself. Handing a type + over is allowed; what is not allowed is touching `.pg` afterwards.""" + schema = await NodeSchema.from_class(_Doc, key="slug") + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(schema, ManagedBy.USER), + await self._prev(prev_ownerships), + not prev_ownerships, + ) + assert out is not None + assert out.action.main_action is None and not out.action.property_actions + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("prev_ownerships", "drops"), + [ + ([], False), + ([ManagedBy.SYSTEM], True), + ([ManagedBy.USER], False), + ([ManagedBy.SYSTEM, ManagedBy.USER], False), + ], + ids=["no-prev", "system", "user", "mixed"], + ) + async def test_undeclared( + self, prev_ownerships: list[ManagedBy], drops: bool + ) -> None: + """Removal is the case that matters, because it is the irreversible + one. The connector drops only what it owns outright, and a single + user-managed candidate vetoes the drop: after an interrupted update we + do not get to pick which candidate is the engine's real state, and + guessing wrong deletes rows the connector never created.""" + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + coco.NON_EXISTENCE, + await self._prev(prev_ownerships), + False, + ) + if not drops: + assert out is None + return + assert out is not None + assert coco.is_non_existence(out.action.spec) + assert out.action.main_action == "delete" + + +class TestGuardsUnderPrevMayBeMissing: + """The non-nullable-addition guard reads a diff computed with + `prev_may_be_missing=False`, and these tests are what pin that. + + Collapsing the two diffs into one looks like a simplification and + silently disables the guard: with `prev_may_be_missing` set (a + `--reprocess`, or internal state lost while the graph persists) the main + action becomes "upsert", which empties `property_actions` and leaves it + nothing to fire on. The failure mode is not a wrong error message -- it is + the raw `OG-MF-103` engine error coming back. + """ + + @pytest.mark.asyncio + async def test_user_managed_drift_writes_nothing_even_when_prev_may_be_missing( + self, + ) -> None: + """`prev_may_be_missing` turns the write path's main action into an + "upsert" for a system-managed type. A user-managed one must still + emit no schema action at all, and must not be rejected either.""" + + @dataclass + class _V1: + slug: str + + @dataclass + class _V2: + slug: str + title: str | None + + prev = _type_tracking_record_from_spec( + _spec(await NodeSchema.from_class(_V1, key="slug"), ManagedBy.USER) + ) + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(await NodeSchema.from_class(_V2, key="slug"), ManagedBy.USER), + [prev], + True, + ) + assert out is not None + assert out.action.main_action is None and not out.action.property_actions + + @pytest.mark.asyncio + async def test_non_nullable_addition_still_raises(self) -> None: + @dataclass + class _V1: + slug: str + + @dataclass + class _V2: + slug: str + title: str # non-nullable addition + + prev = _type_tracking_record_from_spec( + _spec(await NodeSchema.from_class(_V1, key="slug")) + ) + with pytest.raises(ValueError, match="non-nullable"): + _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(await NodeSchema.from_class(_V2, key="slug")), + [prev], + True, + ) + + @pytest.mark.asyncio + async def test_property_on_only_some_prev_records_is_lossy(self) -> None: + """An accepted behaviour change from the statediff port. + + A nullable property that only SOME candidate previous states carry + used to be classified as a plain additive alter; it is lossy now. + More conservative, matches neo4j and falkordb, and reachable only + after an interrupted update left several candidates behind -- exactly + the moment when guessing "additive" is least safe. + """ + + @dataclass + class _V1: + slug: str + + @dataclass + class _V2: + slug: str + title: str | None + + narrow = await NodeSchema.from_class(_V1, key="slug") + wide = await NodeSchema.from_class(_V2, key="slug") + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), + _spec(wide), + [ + _type_tracking_record_from_spec(_spec(narrow)), + _type_tracking_record_from_spec(_spec(wide)), + ], + False, + ) + assert out is not None + assert out.action.property_actions == {"prop:title": "upsert"} + assert out.child_invalidation == "lossy" + + +class TestTypeTrackingRecordRoundTrip: + @pytest.mark.asyncio + async def test_managed_by_survives_the_engine_decoder(self) -> None: + """`managed_by` is only worth persisting if it comes back. + + The engine stores whatever `reconcile` returns and hands it back as + `prev_possible_records` on the next run, decoded against this + handler's own annotation. Refusing to drop a user-managed type is + decided entirely from that restored flag, so a record that + round-trips without it silently reinstates the drop. + """ + from cocoindex._internal import serde + + schema = await NodeSchema.from_class(_Doc, key="slug") + out = _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), _spec(schema, ManagedBy.USER), [], True + ) + assert out is not None + assert not coco.is_non_existence(out.tracking_record) + + hint = serde.unwrap_element_type( + serde.get_param_annotation(_NodeTypeHandler().reconcile, 2) + ) + restored = serde.make_deserialize_fn(hint)(serde.serialize(out.tracking_record)) + assert restored == out.tracking_record + assert restored.managed_by == ManagedBy.USER + + # And the restored record really is enough to veto the drop. + assert ( + _NodeTypeHandler().reconcile( + _TypeKey("og", "node", "Doc"), coco.NON_EXISTENCE, [restored], False + ) + is None + ) + + +_NK = _TypeKey("og", "node", "Source") +_EK = _TypeKey("og", "edge", "Supports") + + +_SOURCE_PROPS = { + "slug": PropertyDef("slug", "String"), + "title": PropertyDef("title", "String"), +} +_SUPPORTS_PROPS = {"w": PropertyDef("w", "I64")} + + +def _node_handler(title_encoder: ogt.ValueEncoder | None = None) -> _NodeHandler: + props = dict(_SOURCE_PROPS) + if title_encoder is not None: + props["title"] = PropertyDef("title", "String", title_encoder) + return _NodeHandler("Source", ("slug",), _NK, props) + + +def _edge_handler(w_encoder: ogt.ValueEncoder | None = None) -> _EdgeHandler: + props = dict(_SUPPORTS_PROPS) + if w_encoder is not None: + props["w"] = PropertyDef("w", "I64", w_encoder) + return _EdgeHandler( + "Supports", + _EK, + "Source", + "Claim", + _SOURCE_PROPS["slug"], + _SOURCE_PROPS["slug"], + props, + ) + + +class TestNodeReconcile: + def test_new_node_upserts(self) -> None: + out = _node_handler().reconcile( + "s1", _NodeValue({"slug": "s1", "title": "T"}), [], False + ) + assert out is not None and out.action.op == "upsert" + + def test_encoder_change_forces_a_write(self) -> None: + """Change detection has to see what the engine will be sent, not the + raw Python value: swapping a property's encoder changes every stored + value while leaving every raw value alone, so a fingerprint taken + before encoding scheduled no write at all.""" + v = _NodeValue({"slug": "s1", "title": "Title"}) + first = _node_handler(str.lower).reconcile("s1", v, [], False) + assert first is not None + assert not isinstance(first.tracking_record, coco.NonExistenceType) + out = _node_handler(str.upper).reconcile( + "s1", v, [first.tracking_record], False + ) + assert out is not None and out.action.op == "upsert" + assert {p.name: p.value for p in out.action.properties}["title"] == "TITLE" + + def test_unchanged_is_noop(self) -> None: + h, v = _node_handler(), _NodeValue({"slug": "s1", "title": "T"}) + first = h.reconcile("s1", v, [], False) + assert first is not None + assert not isinstance(first.tracking_record, coco.NonExistenceType) + assert h.reconcile("s1", v, [first.tracking_record], False) is None + + def test_changed_reupserts(self) -> None: + h = _node_handler() + first = h.reconcile("s1", _NodeValue({"slug": "s1", "title": "T"}), [], False) + assert first is not None + assert not isinstance(first.tracking_record, coco.NonExistenceType) + out = h.reconcile( + "s1", + _NodeValue({"slug": "s1", "title": "T2"}), + [first.tracking_record], + False, + ) + assert out is not None and out.action.op == "upsert" + + def test_prev_may_be_missing_forces_write(self) -> None: + h, v = _node_handler(), _NodeValue({"slug": "s1", "title": "T"}) + first = h.reconcile("s1", v, [], False) + assert first is not None + assert not isinstance(first.tracking_record, coco.NonExistenceType) + out = h.reconcile("s1", v, [first.tracking_record], True) + assert out is not None and out.action.op == "upsert" + + def test_undeclared_deletes(self) -> None: + h = _node_handler() + first = h.reconcile("s1", _NodeValue({"slug": "s1", "title": "T"}), [], False) + assert first is not None + assert not isinstance(first.tracking_record, coco.NonExistenceType) + out = h.reconcile("s1", coco.NON_EXISTENCE, [first.tracking_record], False) + assert out is not None and out.action.op == "delete" + # The sink turns this delete into a key-only stub if an edge still + # references the node, so the key travels with it. + assert out.action.key_properties == (PropertyValue("slug", "String", "s1"),) + + def test_undeclared_and_never_written_is_noop(self) -> None: + assert _node_handler().reconcile("s1", coco.NON_EXISTENCE, [], False) is None + + def test_undeclared_with_uncertain_prev_still_deletes(self) -> None: + """`prev_possible_records=[]` with `prev_may_be_missing` does NOT mean + the entity was never written β€” it means we lost track of it. + + The engine reports that state after a `--reprocess`, after internal + state is lost while the graph persists, and when a prior delete's + tracking record was dropped before the delete itself landed. Treating + it as "never existed" orphans the row in the graph forever: nothing + else ever deletes it, and nothing tracks it any more. + + Deleting by `coco_key` is a documented no-op when the entity isn't + there, so emitting the delete is safe even if it turns out to be + unnecessary. This is the guard every sibling connector uses, and it is + the same reasoning the insert path already applies one method below β€” + which is why the two must not disagree. + """ + out = _node_handler().reconcile("s1", coco.NON_EXISTENCE, [], True) + assert out is not None and out.action.op == "delete" + + def test_coco_key_is_derived_from_the_key_tuple(self) -> None: + out = _node_handler().reconcile("s1", _NodeValue({"slug": "s1"}), [], False) + assert out is not None + assert out.action.coco_key == derive_coco_key(("s1",)) + + def test_scalar_key_equals_singleton_tuple_key(self) -> None: + """A single-field key may arrive as a bare scalar or as a 1-tuple; + both must zip identically against `key_fields` and yield the same + `coco_key`.""" + h = _node_handler() + scalar_out = h.reconcile("s1", _NodeValue({"slug": "s1"}), [], False) + tuple_out = h.reconcile(("s1",), _NodeValue({"slug": "s1"}), [], False) + assert scalar_out is not None and tuple_out is not None + assert scalar_out.action.coco_key == tuple_out.action.coco_key + + +class TestEdgeReconcile: + def test_new_edge_inserts(self) -> None: + out = _edge_handler().reconcile(("a", "b"), _EdgeValue("a", "b", {}), [], False) + assert out is not None and out.action.op == "insert" + + def test_encoder_change_forces_a_replace(self) -> None: + """Edge counterpart of the node case: the fingerprint must cover the + encoded value, or an encoder change never reaches the graph.""" + v = _EdgeValue("a", "b", {"w": 2}) + first = _edge_handler(lambda w: w * 10).reconcile(("a", "b"), v, [], False) + assert first is not None + assert not isinstance(first.tracking_record, coco.NonExistenceType) + out = _edge_handler(lambda w: w * 100).reconcile( + ("a", "b"), v, [first.tracking_record], False + ) + assert out is not None and out.action.op == "replace" + assert {p.name: p.value for p in out.action.properties}["w"] == 200 + + def test_unchanged_is_noop(self) -> None: + h, v = _edge_handler(), _EdgeValue("a", "b", {"w": 1}) + first = h.reconcile(("a", "b"), v, [], False) + assert first is not None + assert not isinstance(first.tracking_record, coco.NonExistenceType) + assert h.reconcile(("a", "b"), v, [first.tracking_record], False) is None + + def test_changed_edge_is_replace_not_insert(self) -> None: + """A strict insert would duplicate; there is no edge update. Must be replace.""" + h = _edge_handler() + first = h.reconcile(("a", "b"), _EdgeValue("a", "b", {"w": 1}), [], False) + assert first is not None + assert not isinstance(first.tracking_record, coco.NonExistenceType) + out = h.reconcile( + ("a", "b"), _EdgeValue("a", "b", {"w": 2}), [first.tracking_record], False + ) + assert out is not None and out.action.op == "replace" + + def test_prev_may_be_missing_forces_replace(self) -> None: + h, v = _edge_handler(), _EdgeValue("a", "b", {}) + first = h.reconcile(("a", "b"), v, [], False) + assert first is not None + assert not isinstance(first.tracking_record, coco.NonExistenceType) + out = h.reconcile(("a", "b"), v, [first.tracking_record], True) + assert out is not None and out.action.op == "replace" + + def test_undeclared_deletes(self) -> None: + h = _edge_handler() + first = h.reconcile(("a", "b"), _EdgeValue("a", "b", {}), [], False) + assert first is not None + assert not isinstance(first.tracking_record, coco.NonExistenceType) + out = h.reconcile( + ("a", "b"), coco.NON_EXISTENCE, [first.tracking_record], False + ) + assert out is not None and out.action.op == "delete" + + def test_undeclared_with_uncertain_prev_still_deletes(self) -> None: + """Edge counterpart of the node case: an edge whose tracking was lost + must still be deleted, or it survives in the graph untracked forever. + + `_EdgeHandler.reconcile` already refuses to trust `(prev=[], + prev_may_be_missing=True)` on the insert side β€” it forces "replace" + there precisely because the edge might still be live. The delete side + must read that same signal the same way. + """ + out = _edge_handler().reconcile(("a", "b"), coco.NON_EXISTENCE, [], True) + assert out is not None and out.action.op == "delete" + + def test_repointing_changes_coco_key(self) -> None: + h = _edge_handler() + a = h.reconcile(("s1", "c1"), _EdgeValue("s1", "c1", {}), [], False) + b = h.reconcile(("s1", "c2"), _EdgeValue("s1", "c2", {}), [], False) + assert a is not None and b is not None + assert a.action.coco_key != b.action.coco_key + + def test_endpoint_metadata_is_carried_for_stubs(self) -> None: + """The sink needs each endpoint's full key definition for stubs.""" + out = _edge_handler().reconcile(("a", "b"), _EdgeValue("a", "b", {}), [], False) + assert out is not None + assert out.action.from_type == "Source" and out.action.to_type == "Claim" + assert out.action.from_key_property == PropertyDef("slug", "String") + + def test_empty_and_prev_may_be_missing_is_replace_not_insert(self) -> None: + """Empty `prev_possible_records` does not mean the edge is absent: the + engine reports (empty, prev_may_be_missing=True) both when a prior + delete's tracking record was dropped before the delete itself landed, + and when internal state was lost outright while the target persists + (e.g. db_path repointed). A bare insert in either case would silently + duplicate a still-live edge. Must fall through to replace β€” deleting + a nonexistent edge by coco_key is a no-op, so replace is safe even if + the edge turns out not to have existed.""" + out = _edge_handler().reconcile(("a", "b"), _EdgeValue("a", "b", {}), [], True) + assert out is not None and out.action.op == "replace" + + +def _client() -> _CliClient: + return _CliClient(ConnectionFactory(store="file:///tmp/g.omni")) + + +#: Stands in wherever `_apply_type_schema` needs a spec that merely exists. +_PLACEHOLDER_SPEC = _TypeSpec( + schema=None, + key=(), + from_type=None, + to_type=None, + managed_by=ManagedBy.SYSTEM, + owner="og", +) + + +def _type_action( + main_action: statediff.DiffAction | None, + type_name: str, + fragment: str | None, + *, + type_kind: str = "node", + property_actions: dict[str, statediff.DiffAction] | None = None, + release_ownership: bool = False, + owner: str = "og", +) -> ogt._TypeAction: + """A `_TypeAction` carrying only what `_apply_type_schema` reads β€” whether + anything has to be written at all (`main_action`/`property_actions`), + whether the type is being removed (`spec`), the type name, and the + fragment. + + `spec`'s *contents* are what the sink needs to build a child handler and + are never consulted on the schema path, so a placeholder stands in for + every action but a removal β€” which the path detects with + `coco.is_non_existence(spec)`, and which `_reconcile_removal` is the only + producer of, always paired with `main_action="delete"`. + """ + return ogt._TypeAction( + key=_TypeKey("og", type_kind, type_name), + spec=coco.NON_EXISTENCE if main_action == "delete" else _PLACEHOLDER_SPEC, + pg_fragment=fragment, + main_action=main_action, + property_actions=property_actions or {}, + release_ownership=release_ownership, + owner=owner, + ) + + +class TestCliArgv: + def test_mutate_argv_passes_payloads_by_file(self) -> None: + """Neither the GQ source nor the params may appear IN argv. + + A commit is a whole component's writes β€” at the 8,192-entity cap + that is ~1.3 MB of expression and ~1.0 MB of params, past darwin's + 1 MB ARG_MAX and far past Linux's 128 KiB per-argument limit (which + caps the expression alone at ~800 entities). `create_subprocess_exec` + then raises `OSError: [Errno 7] Argument list too long`, which is + not an `OmnigraphCliError` and is caught nowhere. + """ + argv = _client()._mutate_argv("/tmp/q.gq", "/tmp/p.json", branch="main") + assert argv[:2] == ["omnigraph", "mutate"] + assert argv[argv.index("--query") + 1] == "/tmp/q.gq" + assert argv[argv.index("--params-file") + 1] == "/tmp/p.json" + # The inline forms are what put a payload in argv -- none of them. + assert "--expr" not in argv + assert "-e" not in argv + assert "--query-string" not in argv + assert "--params" not in argv + assert argv[argv.index("--store") + 1] == "file:///tmp/g.omni" + assert "--branch" in argv and "main" in argv + assert "--json" in argv and "--quiet" in argv + + @pytest.mark.asyncio + async def test_mutate_writes_expr_and_params_to_the_files_it_names( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The argv test above pins the shape; this pins that the files + those paths point at actually hold the mutation, read back from + disk while the call is in flight. Reopening them here also verifies + that the original handles are closed, which Windows requires.""" + seen: dict[str, str] = {} + + async def fake_run(self: object, argv: list[str]) -> dict[str, object]: + query_path = Path(argv[argv.index("--query") + 1]) + params_path = Path(argv[argv.index("--params-file") + 1]) + seen["query_path"] = str(query_path) + seen["params_path"] = str(params_path) + seen["query"] = query_path.read_text() + seen["params"] = params_path.read_text() + return {} + + monkeypatch.setattr(_CliClient, "_run", fake_run) + m = Query("query m($p_a: String) { insert P { a: $p_a } }", {"p_a": "1"}) + await _client().mutate(m, branch="main") + assert seen["query"] == m.expr + assert json.loads(seen["params"]) == {"p_a": "1"} + assert not Path(seen["query_path"]).exists() + assert not Path(seen["params_path"]).exists() + + def test_merge_argv_has_no_cas_flag(self) -> None: + argv = _client()._merge_argv("scratch", into="main") + assert "--if-commit" not in argv + assert "branch" in argv and "merge" in argv + + def test_init_takes_uri_positionally(self) -> None: + argv = _client()._init_argv("/tmp/s.pg") + assert "--schema" in argv and "/tmp/s.pg" in argv + assert argv[-1] == "file:///tmp/g.omni" # positional, not --store + assert "--store" not in argv + + +def _node(op: str, slug: str, name: str = "Source") -> ogt._NodeAction: + return ogt._NodeAction( + op, + _TypeKey("og", "node", name), + name, + (PropertyValue("slug", "String", slug),), + derive_coco_key((slug,)), + ) + + +def _edge(op: str, a: str, b: str) -> ogt._EdgeAction: + return ogt._EdgeAction( + op, + _TypeKey("og", "edge", "Supports"), + "Supports", + derive_coco_key((a, b)), + a, + b, + (), + "Source", + "Claim", + PropertyDef("slug", "String"), + PropertyDef("slug", "String"), + ) + + +class TestStoreLocks: + def test_equivalent_file_uris_share_one_lock(self) -> None: + """The store lock is what makes a visible scratch branch provably + abandoned. It used to be keyed by the raw URI text, so + `file:///a/store` and `file:///a/./store` took different locks: a + schema apply through one reaped the scratch branch another process + was still using through the other, and that process then failed + its own cleanup.""" + paths = { + _CliClient(ConnectionFactory(store=f"file://{root}")).store_lock_path + for root in ( + "/tmp/locks/store", + "/tmp/locks/./store", + "/tmp/locks/../locks/store", + "/tmp//locks/store", + ) + } + assert len(paths) == 1 + other = _CliClient(ConnectionFactory(store="file:///tmp/locks/other")) + assert other.store_lock_path not in paths + + def test_non_file_uris_are_normalised_too(self) -> None: + a = _CliClient(ConnectionFactory(store="s3://Bucket/graphs/./g/")) + b = _CliClient(ConnectionFactory(store="s3://bucket/graphs/g")) + assert a.store_lock_path == b.store_lock_path + + @pytest.mark.asyncio + async def test_reaper_skips_a_scratch_branch_whose_owner_is_alive( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Abandonment is established before deletion: a live scratch branch + is held under a per-branch lock for its whole lifetime, and the + reaper deletes only branches whose lock it can take.""" + deleted: list[str] = [] + + async def fake_branch_list(self: object) -> list[str]: + return ["coco_scratch_live", "coco_scratch_dead", "main"] + + async def fake_branch_delete(self: object, name: str) -> None: + deleted.append(name) + + monkeypatch.setattr(_CliClient, "branch_list", fake_branch_list) + monkeypatch.setattr(_CliClient, "branch_delete", fake_branch_delete) + client = _client() + async with client.hold_scratch_branch("coco_scratch_live"): + reaped = await ogt._reap_abandoned_scratch_branches(client) + assert reaped == ["coco_scratch_dead"] + assert deleted == ["coco_scratch_dead"] + + +class TestCliCancellation: + @pytest.mark.asyncio + async def test_cancelling_a_call_kills_and_reaps_the_child( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Cancelling the task awaiting a CLI call must not leave the CLI + running: asyncio's `communicate()` does nothing to the child on + cancellation, so a cancelled `mutate` kept writing to the store + after the connector had given up on it.""" + started: list[asyncio.subprocess.Process] = [] + real_exec = asyncio.create_subprocess_exec + + async def spying_exec(*argv: str, **kw: Any) -> asyncio.subprocess.Process: + proc = await real_exec(*argv, **kw) + started.append(proc) + return proc + + monkeypatch.setattr(asyncio, "create_subprocess_exec", spying_exec) + client = _CliClient(ConnectionFactory(store="file:///tmp/g.omni", cli="sleep")) + task = asyncio.create_task(client._run(["sleep", "30"])) + while not started: + await asyncio.sleep(0.01) + (proc,) = started + try: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + # Killed AND reaped: a returncode means `wait()` collected it. + assert proc.returncode is not None + finally: + if proc.returncode is None: + proc.kill() + await proc.wait() + + +class TestRenderQuery: + def test_single_statement_gets_the_first_prefix(self) -> None: + m = render_query( + [Statement("insert P { a: $? }", (Bind("p_a", "String", "1"),))] + ) + assert m.params == {"s0_p_a": "1"} + assert m.expr == "query m($s0_p_a: String) { insert P { a: $s0_p_a } }" + + def test_statements_with_the_same_labels_never_collide(self) -> None: + """Every builder uses the same labels (`p_coco_key`, ...). Naming + happens per statement at render time, so two of them share one + query without any rewriting of the statement text.""" + one = Statement( + "insert P { coco_key: $? }", (Bind("p_coco_key", "String", "A"),) + ) + two = Statement( + "insert P { coco_key: $? }", (Bind("p_coco_key", "String", "B"),) + ) + m = render_query([one, two]) + assert m.params == {"s0_p_coco_key": "A", "s1_p_coco_key": "B"} + assert m.expr == ( + "query m($s0_p_coco_key: String, $s1_p_coco_key: String) " + "{ insert P { coco_key: $s0_p_coco_key } insert P { coco_key: $s1_p_coco_key } }" + ) + + def test_a_label_that_is_a_prefix_of_another_is_rendered_whole(self) -> None: + m = render_query( + [ + Statement( + "insert P { a: $?, ab: $? }", + (Bind("p_a", "String", "1"), Bind("p_ab", "String", "2")), + ) + ] + ) + assert m.expr == ( + "query m($s0_p_a: String, $s0_p_ab: String) " + "{ insert P { a: $s0_p_a, ab: $s0_p_ab } }" + ) + + def test_slot_and_bind_counts_must_agree(self) -> None: + with pytest.raises(ValueError, match="2 slots but 1 bind"): + render_query( + [Statement("insert P { a: $?, b: $? }", (Bind("p_a", "String", 1),))] + ) + + def test_duplicate_labels_in_one_statement_are_rejected(self) -> None: + with pytest.raises(ValueError, match="Duplicate parameter label"): + render_query( + [ + Statement( + "insert P { a: $?, b: $? }", + (Bind("p_x", "String", 1), Bind("p_x", "String", 2)), + ) + ] + ) + + def test_empty_raises(self) -> None: + with pytest.raises(ValueError): + render_query([]) + + +class TestPlanCommits: + def test_upsert_only_is_one_commit(self) -> None: + assert len(plan_commits([_node("upsert", f"s{i}") for i in range(5)])) == 1 + + def test_upserts_and_deletes_never_share_a_commit(self) -> None: + commits = plan_commits([_node("upsert", "a"), _node("delete", "gone")]) + assert len(commits) == 2 + for c in commits: + assert not ("insert" in c.expr and "delete" in c.expr) + + def test_replace_deletes_before_it_inserts(self) -> None: + """Delete selects on coco_key and the replacement reuses it β€” insert-first + would delete the new edge too. Verified against the engine.""" + commits = plan_commits([_edge("replace", "a", "b")]) + exprs = [c.expr for c in commits] + del_i = next(i for i, e in enumerate(exprs) if "delete Supports" in e) + ins_i = next(i for i, e in enumerate(exprs) if "insert Supports" in e) + assert del_i < ins_i + + def test_three_phases_when_replace_and_removal_coexist(self) -> None: + commits = plan_commits([_edge("replace", "a", "b"), _node("delete", "gone")]) + assert len(commits) == 3 + + def test_edge_insert_has_no_endpoint_stubs(self) -> None: + """Endpoint stubs are no longer emitted unconditionally by + plan_commits β€” a keyed insert is a full-record replace in Omnigraph + (verified against the engine), so stubbing every edge's endpoints + would silently null out a node's other nullable properties. The + sink builds a stub reactively, only after an insert fails with a + "not found" error for that specific endpoint β€” see + TestEndpointRetryLive.""" + (commit,) = plan_commits([_edge("insert", "a", "b")]) + assert "insert Source" not in commit.expr + assert "insert Claim" not in commit.expr + assert "insert Supports" in commit.expr + + def test_node_upserts_precede_edge_inserts_whatever_the_action_order( + self, + ) -> None: + """Actions reach `plan_commits` in reconcile order, which is not the + commit order they need: here the edge is listed first, and both its + endpoints' upserts are listed after it. Appending in arrival order + would emit the edge insert ahead of the nodes it references, inside + the same commit, and the engine would refuse it.""" + commits = plan_commits( + [ + _edge("insert", "a", "b"), + _node("upsert", "a"), + _node("upsert", "b", name="Claim"), + ] + ) + (phase_b,) = commits + assert phase_b.expr.index("insert Source") < phase_b.expr.index( + "insert Supports" + ) + assert phase_b.expr.index("insert Claim") < phase_b.expr.index( + "insert Supports" + ) + + def test_endpoint_refs_are_always_string(self) -> None: + """An edge's `from`/`to` holds the endpoint's node `id`, which is a + String whatever the key's declared type is: passing an `I64` for an + `I64`-keyed endpoint is refused outright ("cannot assign/compare I64 + with String for property `to`", verified against the engine). The + pg_type used to be inferred from the value's Python type, so any + non-string-keyed endpoint β€” an int-keyed Meeting, say β€” failed on + every edge insert.""" + action = ogt._EdgeAction( + "insert", + _TypeKey("og", "edge", "Attended"), + "Attended", + derive_coco_key(("Ada", 7)), + "Ada", + 7, + (), + "Person", + "Meeting", + PropertyDef("name", "String"), + PropertyDef("meeting_id", "I64"), + ) + (commit,) = plan_commits([action]) + assert "$s0_e_to: String" in commit.expr + assert "$s0_e_to: I64" not in commit.expr + assert commit.params["s0_e_to"] == "7" + assert commit.params["s0_e_from"] == "Ada" + # ...but the coco_key still derives from the ORIGINAL values, so it + # matches what `_EdgeHandler` tracked and what a later delete selects. + assert commit.params["s0_p_coco_key"] == derive_coco_key(("Ada", 7)) + + def test_edge_deletes_precede_node_deletes(self) -> None: + commits = plan_commits([_edge("delete", "a", "b"), _node("delete", "a")]) + joined = " || ".join(c.expr for c in commits) + assert joined.index("delete Supports") < joined.index("delete Source") + + def test_chunks_when_over_the_cap(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(ogt, "_MAX_ENTITIES_PER_TYPE", 4) + assert len(plan_commits([_node("upsert", f"s{i}") for i in range(10)])) == 3 + + def test_cap_is_per_type_not_global(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(ogt, "_MAX_ENTITIES_PER_TYPE", 4) + actions = [_node("upsert", f"s{i}") for i in range(3)] + actions += [_node("upsert", f"c{i}", name="Claim") for i in range(3)] + assert len(plan_commits(actions)) == 1 + + def test_a_total_cap_bounds_the_commit_across_types( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The per-type cap is deliberate, but on its own it bounded nothing: + N types each at the cap produced ONE commit of N x cap entities + (10 types measured at ~9.5 MB of GQ), and + `_mutate_with_endpoint_retry` re-sends the whole commit on every + endpoint-stub retry. + """ + monkeypatch.setattr(ogt, "_MAX_ENTITIES_PER_TYPE", 4) + monkeypatch.setattr(ogt, "_MAX_ENTITIES_PER_COMMIT", 5) + actions = [_node("upsert", f"s{i}") for i in range(3)] + actions += [_node("upsert", f"c{i}", name="Claim") for i in range(3)] + commits = plan_commits(actions) + assert len(commits) > 1 + assert max(c.expr.count("insert ") for c in commits) <= 5 + + def test_mixed_workload_respects_all_phase_invariants(self) -> None: + """One call exercising all four action kinds together β€” a node + upsert, an edge replace, an edge removal, and a node removal β€” the + exact combination the three-phase design exists to keep safe. The + commit sequence must never delete something it just wrote, and + never write something it is about to delete; verified directly via + each mutation's own bound coco_key, not just by argument.""" + replace_coco_key = derive_coco_key(("r1", "r2")) + removed_edge_coco_key = derive_coco_key(("d1", "d2")) + removed_node_coco_key = derive_coco_key(("gone",)) + + actions = [ + _node("upsert", "keep"), + _edge("replace", "r1", "r2"), + _edge("delete", "d1", "d2"), + _node("delete", "gone"), + ] + commits = plan_commits(actions) + assert len(commits) == 3 + phase_a, phase_b, phase_c = commits + + # Phase A: only the replace's own delete β€” nothing written yet. + assert "delete Supports" in phase_a.expr + assert "insert" not in phase_a.expr + assert replace_coco_key in phase_a.params.values() + + # Phase B: the kept node's upsert and the replace's re-insert, + # under the SAME coco_key phase A just deleted β€” never landing + # without that delete having already applied ahead of it. + assert "insert Source" in phase_b.expr + assert "insert Supports" in phase_b.expr + assert "delete" not in phase_b.expr + assert replace_coco_key in phase_b.params.values() + # ...and the node comes FIRST. Statements in a combined query + # execute in the order written, so an edge insert ahead of its own + # endpoint's upsert would fail "not found" and force the sink's + # stub-and-retry recovery on every single run. + assert phase_b.expr.index("insert Source") < phase_b.expr.index( + "insert Supports" + ) + + # Phase C: the two removals, edge before node, under coco_keys that + # never appear in phase B β€” nothing here was just written. + assert phase_c.expr.index("delete Supports") < phase_c.expr.index( + "delete Source" + ) + assert removed_edge_coco_key in phase_c.params.values() + assert removed_node_coco_key in phase_c.params.values() + assert replace_coco_key not in phase_c.params.values() + assert not set(phase_c.params.values()) & set(phase_b.params.values()) + + +class TestPlanCommitsEncoding: + @pytest.mark.asyncio + async def test_date_and_datetime_properties_survive_json_dumps(self) -> None: + """A dataclass with date/datetime/optional fields must round-trip + through plan_commits into a Query whose params survive + json.dumps β€” this is what _CliClient._mutate_argv does with them.""" + + @dataclass + class _Event: + slug: str + on: datetime.date + at: datetime.datetime + note: str | None + + schema = await NodeSchema.from_class(_Event, key="slug") + handler = _NodeHandler( + "Event", ("slug",), _TypeKey("og", "node", "Event"), schema.properties + ) + out = handler.reconcile( + "e1", + _NodeValue( + { + "slug": "e1", + "on": datetime.date(2026, 1, 1), + "at": datetime.datetime(2026, 1, 1, 12, 0), # noqa: DTZ001 + "note": None, + } + ), + [], + False, + ) + assert out is not None + (commit,) = plan_commits([out.action]) + json.dumps(commit.params) # must not raise + # plan_commits always renders a chunk with render_query, even a + # single one β€” so params carry the `s0_` commit prefix. + assert commit.params["s0_p_on"] == "2026-01-01" + assert commit.params["s0_p_at"] == "2026-01-01T12:00:00" + assert commit.params["s0_p_note"] is None + + +class TestParseMissingEndpoint: + @pytest.mark.parametrize( + ("role", "key", "type_name"), + [ + ("src", "O'Brien", "Person"), + ("dst", "a'b'c", "Claim"), + ("src", "", "EmptyKey"), + ], + ) + def test_preserves_apostrophes_in_key( + self, role: str, key: str, type_name: str + ) -> None: + error = OmnigraphCliError(f"{role} '{key}' not found in {type_name}") + + assert ogt._parse_missing_endpoint(error) == (role, key, type_name) + + +class TestBuildEndpointStub: + def test_uses_the_endpoint_schemas_declared_key_type(self) -> None: + action = ogt._EdgeAction( + "insert", + _TypeKey("og", "edge", "Supports"), + "Supports", + derive_coco_key((7, "c1")), + 7, + "c1", + (), + "Source", + "Claim", + PropertyDef("source_id", "I32"), + PropertyDef("slug", "String"), + ) + + stub = ogt._build_endpoint_stub("src", "7", "Source", [action]) + + assert stub is not None + rendered = render_query([stub]) + assert "$s0_p_source_id: I32" in rendered.expr + assert "$s0_p_source_id: I64" not in rendered.expr + assert rendered.params["s0_p_source_id"] == 7 + + def test_delete_actions_are_not_stub_candidates(self) -> None: + """A delete `_EdgeAction` carries `from_id=None`/`to_id=None`, so + `str(action.from_id)` is the literal `"None"`. + + A node whose key value is the string `"None"` therefore matched a + delete action and built a meaningless endpoint stub. A delete needs no + endpoints anyway: deleting by `coco_key` never touches them. + """ + h = _edge_handler() + first = h.reconcile(("a", "b"), _EdgeValue("a", "b", {}), [], False) + assert first is not None + assert not isinstance(first.tracking_record, coco.NonExistenceType) + deleted = h.reconcile( + ("a", "b"), coco.NON_EXISTENCE, [first.tracking_record], False + ) + assert deleted is not None and deleted.action.op == "delete" + + assert ( + ogt._build_endpoint_stub("src", "None", "Source", [deleted.action]) is None + ) + + +class TestApplyEntityActions: + @pytest.mark.asyncio + async def test_branch_create_failure_propagates_original_error( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Cleanup must never mask the original failure: if branch_create + itself fails, branch_delete(scratch) has nothing to delete and + raises too β€” that second error must be suppressed so the original + branch_create failure is what actually propagates.""" + + async def fake_branch_create(self: object, name: str, *, frm: str) -> None: + raise OmnigraphCliError("branch create failed: disk full") + + async def fake_branch_delete(self: object, name: str) -> None: + raise OmnigraphCliError(f"branch {name!r} does not exist") + + monkeypatch.setattr(_CliClient, "branch_create", fake_branch_create) + monkeypatch.setattr(_CliClient, "branch_delete", fake_branch_delete) + # The delete below makes the sink read the schema first, to learn + # whether some edge type could still reference the node. + monkeypatch.setattr(_CliClient, "read_schema", _noop_read_schema) + + db = ContextKey[ConnectionFactory](f"test_db_{uuid.uuid4().hex}") + cp = ContextProvider() + cp.provide(db, ConnectionFactory(store="file:///tmp/whatever.omni")) + + key = _TypeKey(db.key, "node", "Source") + actions: list[ogt._NodeAction | ogt._EdgeAction] = [ + ogt._NodeAction( + "upsert", + key, + "Source", + (PropertyValue("slug", "String", "a"),), + derive_coco_key(("a",)), + ), + ogt._NodeAction( + "delete", + key, + "Source", + (), + derive_coco_key(("gone",)), + ), + ] + with pytest.raises(OmnigraphCliError, match="branch create failed"): + await ogt._apply_entity_actions(cp, actions) + + @pytest.mark.asyncio + async def test_branch_cleanup_failure_is_reported( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def noop(self: object, *args: object, **kwargs: object) -> None: + return None + + async def fail_delete(self: object, name: str) -> None: + raise OmnigraphCliError(f"failed to delete {name}") + + monkeypatch.setattr(_CliClient, "mutate", noop) + monkeypatch.setattr(_CliClient, "branch_create", noop) + monkeypatch.setattr(_CliClient, "branch_merge", noop) + monkeypatch.setattr(_CliClient, "branch_delete", fail_delete) + monkeypatch.setattr(_CliClient, "read_schema", _noop_read_schema) + + db = ContextKey[ConnectionFactory](f"test_db_{uuid.uuid4().hex}") + cp = ContextProvider() + cp.provide(db, ConnectionFactory(store="file:///tmp/whatever.omni")) + key = _TypeKey(db.key, "node", "Source") + actions: list[ogt._NodeAction | ogt._EdgeAction] = [ + ogt._NodeAction( + "upsert", + key, + "Source", + (PropertyValue("slug", "String", "a"),), + derive_coco_key(("a",)), + ), + ogt._NodeAction("delete", key, "Source", (), derive_coco_key(("gone",))), + ] + + with pytest.raises(OmnigraphCliError, match="failed to delete"): + await ogt._apply_entity_actions(cp, actions) + + @pytest.mark.asyncio + async def test_single_commit_skips_the_scratch_branch( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A batch that plans down to exactly one commit must go straight to + `mutate` on the real branch β€” no branch_create/branch_merge at all.""" + calls: list[str] = [] + + async def fake_mutate(self: object, mutation: Query, *, branch: str) -> None: + calls.append(f"mutate:{branch}") + + async def fake_branch_create(self: object, name: str, *, frm: str) -> None: + calls.append("branch_create") + + monkeypatch.setattr(_CliClient, "mutate", fake_mutate) + monkeypatch.setattr(_CliClient, "branch_create", fake_branch_create) + + db = ContextKey[ConnectionFactory](f"test_db_{uuid.uuid4().hex}") + cp = ContextProvider() + cp.provide( + db, ConnectionFactory(store="file:///tmp/whatever.omni", branch="main") + ) + + key = _TypeKey(db.key, "node", "Source") + actions: list[ogt._NodeAction | ogt._EdgeAction] = [ + ogt._NodeAction( + "upsert", + key, + "Source", + (PropertyValue("slug", "String", "a"),), + derive_coco_key(("a",)), + ), + ] + await ogt._apply_entity_actions(cp, actions) + assert calls == ["mutate:main"] + + @pytest.mark.asyncio + async def test_a_node_an_edge_still_references_becomes_a_stub( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Deleting a node cascades to its edges in the graph, while the + edges' own tracking, in whatever component declared them, still + says they exist β€” so they were never re-inserted. A node that an + edge outside this batch still references is reduced to a key-only + stub instead; one whose only edges this batch deletes too is + deleted for real.""" + schema = TestOrphanedEdgeEndpoints.SCHEMA + queried: list[tuple[str, dict[str, object]]] = [] + commits: list[Query] = [] + edge_ab = derive_coco_key(("a", "co")) + edge_bb = derive_coco_key(("b", "co")) + + async def fake_read_schema(self: object) -> str | None: + return schema + + async def fake_query( + self: object, query: Query, *, branch: str + ) -> list[dict[str, object]]: + queried.append((query.expr, query.params)) + (coco_key,) = query.params.values() + return [ + { + "edge_key": edge_ab + if coco_key == derive_coco_key(("a",)) + else edge_bb + } + ] + + async def fake_mutate(self: object, mutation: Query, *, branch: str) -> None: + commits.append(mutation) + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "query", fake_query) + monkeypatch.setattr(_CliClient, "mutate", fake_mutate) + _patch_cli_branch_calls(monkeypatch) + + db = ContextKey[ConnectionFactory](f"test_db_{uuid.uuid4().hex}") + cp = ContextProvider() + cp.provide(db, ConnectionFactory(store="file:///tmp/whatever.omni")) + + def delete(slug: str) -> ogt._NodeAction: + return ogt._NodeAction( + "delete", + _TypeKey(db.key, "node", "Person"), + "Person", + (), + derive_coco_key((slug,)), + key_properties=(PropertyValue("slug", "String", slug),), + ) + + await ogt._apply_entity_actions( + cp, + [ + delete("a"), + delete("b"), + ogt._EdgeAction( + "delete", + _TypeKey(db.key, "edge", "WorksAt"), + "WorksAt", + edge_bb, + None, + None, + (), + "Person", + "Company", + PropertyDef("slug", "String"), + PropertyDef("cid", "I64"), + ), + ], + ) + + # One query per referencing edge type and side: Person is only ever + # a `from` of WorksAt here. + assert [params for _, params in queried] == [ + {"s0_p_coco_key": derive_coco_key(("a",))}, + {"s0_p_coco_key": derive_coco_key(("b",))}, + ] + assert all("$n $e:worksAt $o" in expr for expr, _ in queried) + # `a` is still referenced: a key-only upsert, in the upsert phase. + # `b`'s only edge goes in this batch: a real delete, after the edge. + assert [c.expr for c in commits] == [ + ( + "query m($s0_p_slug: String, $s0_p_coco_key: String) " + "{ insert Person { slug: $s0_p_slug, coco_key: $s0_p_coco_key } }" + ), + ( + "query m($s0_p_coco_key: String, $s1_p_coco_key: String) " + "{ delete WorksAt where coco_key = $s0_p_coco_key " + "delete Person where coco_key = $s1_p_coco_key }" + ), + ] + assert commits[0].params == { + "s0_p_slug": "a", + "s0_p_coco_key": derive_coco_key(("a",)), + } + + @pytest.mark.asyncio + async def test_missing_endpoint_escalates_a_single_commit_to_a_branch( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Recovering a missing endpoint takes TWO commits β€” the stub, then + the retried commit β€” so doing it on the live branch would leave an + orphan stub node behind, untracked and visible to readers, if the + retry then failed. The single-commit fast path therefore escalates + to the scratch branch the moment it sees a "not found", rather than + stubbing in place. Nothing is lost by the failed first attempt: a + failed mutation applies nothing.""" + calls: list[str] = [] + + async def fake_mutate(self: object, mutation: Query, *, branch: str) -> None: + calls.append(f"mutate:{branch}") + if branch == "main": + raise OmnigraphCliError("dst 'c1' not found in Claim") + + async def fake_branch_create(self: object, name: str, *, frm: str) -> None: + calls.append("branch_create") + + async def fake_branch_merge(self: object, name: str, *, into: str) -> None: + calls.append("branch_merge") + + async def fake_branch_delete(self: object, name: str) -> None: + calls.append("branch_delete") + + monkeypatch.setattr(_CliClient, "mutate", fake_mutate) + monkeypatch.setattr(_CliClient, "branch_create", fake_branch_create) + monkeypatch.setattr(_CliClient, "branch_merge", fake_branch_merge) + monkeypatch.setattr(_CliClient, "branch_delete", fake_branch_delete) + + db = ContextKey[ConnectionFactory](f"test_db_{uuid.uuid4().hex}") + cp = ContextProvider() + cp.provide( + db, ConnectionFactory(store="file:///tmp/whatever.omni", branch="main") + ) + + key = _TypeKey(db.key, "edge", "Supports") + actions: list[ogt._NodeAction | ogt._EdgeAction] = [ + ogt._EdgeAction( + "insert", + key, + "Supports", + derive_coco_key(("s1", "c1")), + "s1", + "c1", + (), + "Source", + "Claim", + PropertyDef("slug", "String"), + PropertyDef("slug", "String"), + ), + ] + await ogt._apply_entity_actions(cp, actions) + + # The failed live attempt, then the whole thing redone on scratch. + assert calls[0] == "mutate:main" + assert calls[1] == "branch_create" + assert "mutate:main" not in calls[1:] # no stub written to the live branch + assert "branch_merge" in calls and calls[-1] == "branch_delete" + + @staticmethod + def _two_edge_actions(db_key: str) -> list[ogt._NodeAction | ogt._EdgeAction]: + key = _TypeKey(db_key, "edge", "Supports") + return [ + ogt._EdgeAction( + "insert", + key, + "Supports", + derive_coco_key((s, c)), + s, + c, + (), + "Source", + "Claim", + PropertyDef("slug", "String"), + PropertyDef("slug", "String"), + ) + for s, c in (("s1", "c1"), ("s2", "c2")) + ] + + @pytest.mark.asyncio + async def test_stubs_every_distinct_missing_endpoint_in_the_batch( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Two edges whose four endpoints are all absent need four stub + rounds β€” the engine names only the first missing endpoint per + attempt β€” and every one of them must be taken. A fixed two-round + budget failed this batch on its third attempt.""" + present: set[tuple[str, str]] = set() + + async def fake_mutate(self: object, mutation: Query, *, branch: str) -> None: + if "insert Supports" in mutation.expr: + for i in range(2): + frm, to = ( + mutation.params[f"s{i}_e_from"], + mutation.params[f"s{i}_e_to"], + ) + if ("Source", frm) not in present: + raise OmnigraphCliError(f"src '{frm}' not found in Source") + if ("Claim", to) not in present: + raise OmnigraphCliError(f"dst '{to}' not found in Claim") + return + # Anything else is a key-only endpoint stub. + type_name = mutation.expr.split("insert ")[1].split(" ")[0] + present.add((type_name, str(mutation.params["s0_p_slug"]))) + + async def noop(self: object, *args: object, **kwargs: object) -> None: + return None + + monkeypatch.setattr(_CliClient, "mutate", fake_mutate) + monkeypatch.setattr(_CliClient, "branch_create", noop) + monkeypatch.setattr(_CliClient, "branch_merge", noop) + monkeypatch.setattr(_CliClient, "branch_delete", noop) + + db = ContextKey[ConnectionFactory](f"test_db_{uuid.uuid4().hex}") + cp = ContextProvider() + cp.provide(db, ConnectionFactory(store="file:///tmp/whatever.omni")) + await ogt._apply_entity_actions(cp, self._two_edge_actions(db.key)) + assert present == { + ("Source", "s1"), + ("Claim", "c1"), + ("Source", "s2"), + ("Claim", "c2"), + } + + @pytest.mark.asyncio + async def test_endpoint_still_missing_after_its_stub_fails_loudly( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A "not found" that names an endpoint already stubbed in this + batch means the stub did not take; that must propagate, not spin.""" + attempts = 0 + + async def fake_mutate(self: object, mutation: Query, *, branch: str) -> None: + nonlocal attempts + if "insert Supports" in mutation.expr: + attempts += 1 + raise OmnigraphCliError("src 's1' not found in Source") + + async def noop(self: object, *args: object, **kwargs: object) -> None: + return None + + monkeypatch.setattr(_CliClient, "mutate", fake_mutate) + monkeypatch.setattr(_CliClient, "branch_create", noop) + monkeypatch.setattr(_CliClient, "branch_merge", noop) + monkeypatch.setattr(_CliClient, "branch_delete", noop) + + db = ContextKey[ConnectionFactory](f"test_db_{uuid.uuid4().hex}") + cp = ContextProvider() + cp.provide(db, ConnectionFactory(store="file:///tmp/whatever.omni")) + with pytest.raises(OmnigraphCliError, match="not found"): + await ogt._apply_entity_actions(cp, self._two_edge_actions(db.key)) + # The live attempt, the first scratch attempt, and exactly one retry + # after the stub β€” never a second retry for the same endpoint. + assert attempts == 3 + + @pytest.mark.asyncio + async def test_scratch_branch_blocks_schema_changes_until_deleted( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + branch_created = asyncio.Event() + allow_mutations = asyncio.Event() + schema_started = asyncio.Event() + schema_read = asyncio.Event() + calls: list[str] = [] + + async def fake_mutate(self: object, mutation: Query, *, branch: str) -> None: + calls.append(f"mutate:{branch}") + await allow_mutations.wait() + + async def fake_branch_create(self: object, name: str, *, frm: str) -> None: + calls.append("branch_create") + branch_created.set() + + async def fake_branch_merge(self: object, name: str, *, into: str) -> None: + calls.append("branch_merge") + + async def fake_branch_delete(self: object, name: str) -> None: + calls.append("branch_delete") + + async def fake_read_schema(self: object) -> str | None: + calls.append("schema_read") + schema_read.set() + return "node Existing {\n slug: String @key\n coco_key: String\n}\n" + + async def fake_apply_schema(self: object, schema_pg: str) -> None: + calls.append("schema_apply") + + monkeypatch.setattr(_CliClient, "mutate", fake_mutate) + monkeypatch.setattr(_CliClient, "branch_create", fake_branch_create) + monkeypatch.setattr(_CliClient, "branch_merge", fake_branch_merge) + monkeypatch.setattr(_CliClient, "branch_delete", fake_branch_delete) + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "apply_schema", fake_apply_schema) + + store = f"file:///tmp/{uuid.uuid4().hex}.omni" + conn = ConnectionFactory(store=store) + db = ContextKey[ConnectionFactory](f"test_db_{uuid.uuid4().hex}") + cp = ContextProvider() + cp.provide(db, conn) + edge_action = ogt._EdgeAction( + "replace", + _TypeKey(db.key, "edge", "Supports"), + "Supports", + derive_coco_key(("s1", "c1")), + "s1", + "c1", + (), + "Source", + "Claim", + PropertyDef("slug", "String"), + PropertyDef("slug", "String"), + ) + schema_action = _type_action( + "insert", + "Added", + "node Added {\n slug: String @key\n coco_key: String\n}", + ) + + async def reconcile_schema() -> None: + schema_started.set() + await ogt._apply_type_schema(_CliClient(conn), [schema_action]) + + entity_task = asyncio.create_task(ogt._apply_entity_actions(cp, [edge_action])) + await branch_created.wait() + schema_task = asyncio.create_task(reconcile_schema()) + await schema_started.wait() + try: + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(schema_read.wait(), timeout=0.05) + finally: + allow_mutations.set() + await asyncio.gather(entity_task, schema_task) + + assert calls.index("branch_delete") < calls.index("schema_read") + + +class TestReadSchema: + """Unit coverage for read_schema's not-found detection, mocked so it + doesn't need the live binary β€” TestApplyTypeActionsLive below is what + proves it against the engine's actual wording.""" + + @pytest.mark.asyncio + async def test_returns_none_on_the_actual_dataset_not_found_wording( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Verified against the binary: the real message is `storage: + Dataset at path was not found: Not found: ...` β€” "Dataset" + and "not found" are not adjacent, so a naive `"dataset not found"` + substring check would never match this.""" + + async def fake_run(self: object, argv: list[str]) -> dict[str, object]: + raise OmnigraphCliError( + "omnigraph schema show exited 1: Error: \n 0: storage: " + "Dataset at path /tmp/g.omni/__manifest was not found: " + "Not found: /tmp/g.omni/__manifest/_versions" + ) + + monkeypatch.setattr(_CliClient, "_run", fake_run) + assert await _client().read_schema() is None + + @pytest.mark.asyncio + async def test_unrelated_error_propagates( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """An unrelated failure that happens to say "not found" (a bad type + reference, say) must propagate as-is, not be read as "uninitialized".""" + + async def fake_run(self: object, argv: list[str]) -> dict[str, object]: + raise OmnigraphCliError( + "omnigraph schema show exited 1: Error: \n 0: type " + "'Ghost' referenced in edge definition not found in schema" + ) + + monkeypatch.setattr(_CliClient, "_run", fake_run) + with pytest.raises(OmnigraphCliError, match="Ghost"): + await _client().read_schema() + + @pytest.mark.asyncio + async def test_a_store_path_containing_dataset_is_not_read_as_uninitialized( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The store URI is echoed into the CLI's error text (verified against + the binary), so a store living under `~/datasets/` hands the word + "dataset" to *every* failure it ever reports. + + Testing for "dataset" and "not found" as independent substrings + therefore misread any unrelated not-found failure on such a store as + "graph not initialized" β€” and `_apply_type_schema` would then take the + `init_graph` branch against a populated store. Match the engine's + actual phrasing instead. + """ + + async def fake_run(self: object, argv: list[str]) -> dict[str, object]: + raise OmnigraphCliError( + "omnigraph schema show exited 1: Error: \n 0: type 'Ghost' " + "referenced in edge definition not found in schema " + "(store file:///Users/me/datasets/kg.omni)" + ) + + monkeypatch.setattr(_CliClient, "_run", fake_run) + with pytest.raises(OmnigraphCliError, match="Ghost"): + await _client().read_schema() + + +class TestApplyTypeSchema: + """Unit coverage for _apply_type_schema's init-vs-merge branching, + mocked so it doesn't need the live binary.""" + + _REFUSAL = ( + "schema apply exited 1: schema apply requires a graph with only main; " + "found non-main branches: {names}" + ) + + @pytest.mark.asyncio + async def test_abandoned_scratch_branch_is_reaped_and_the_apply_retried( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A process killed between creating and deleting its scratch branch + leaves a `coco_scratch_*` branch behind, and Omnigraph refuses every + later `schema apply` while it exists. Nothing ever recovered it. Any + such branch seen while holding the store lock is abandoned β€” a live + one is held under that same lock β€” so the sink deletes it and + retries the apply once.""" + calls: list[str] = [] + branches = ["coco_scratch_deadbeef", "main"] + + async def fake_read_schema(self: object) -> str | None: + return "node A {\n slug: String @key\n coco_key: String\n}\n" + + async def fake_apply_schema(self: object, pg_fragment: str) -> None: + calls.append("apply") + stale = [b for b in branches if b != "main"] + if stale: + raise OmnigraphCliError( + TestApplyTypeSchema._REFUSAL.format(names=", ".join(stale)) + ) + + async def fake_branch_list(self: object) -> list[str]: + calls.append("branch_list") + return list(branches) + + async def fake_branch_delete(self: object, name: str) -> None: + calls.append(f"branch_delete:{name}") + branches.remove(name) + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "apply_schema", fake_apply_schema) + monkeypatch.setattr(_CliClient, "branch_list", fake_branch_list, raising=False) + monkeypatch.setattr(_CliClient, "branch_delete", fake_branch_delete) + + fragment = "node B {\n slug: String @key\n coco_key: String\n}" + await ogt._apply_type_schema(_client(), [_type_action("insert", "B", fragment)]) + assert calls == [ + "apply", + "branch_list", + "branch_delete:coco_scratch_deadbeef", + "apply", + ] + + @pytest.mark.asyncio + async def test_a_user_branch_is_never_reaped( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Only the connector's own `coco_scratch_*` prefix is fair game. A + user's branch blocks the apply exactly as documented, and the + refusal propagates untouched.""" + calls: list[str] = [] + + async def fake_read_schema(self: object) -> str | None: + return "node A {\n slug: String @key\n coco_key: String\n}\n" + + async def fake_apply_schema(self: object, pg_fragment: str) -> None: + calls.append("apply") + raise OmnigraphCliError( + TestApplyTypeSchema._REFUSAL.format(names="staging") + ) + + async def fake_branch_list(self: object) -> list[str]: + return ["main", "staging"] + + async def fake_branch_delete(self: object, name: str) -> None: + calls.append(f"branch_delete:{name}") + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "apply_schema", fake_apply_schema) + monkeypatch.setattr(_CliClient, "branch_list", fake_branch_list, raising=False) + monkeypatch.setattr(_CliClient, "branch_delete", fake_branch_delete) + + fragment = "node B {\n slug: String @key\n coco_key: String\n}" + with pytest.raises(OmnigraphCliError, match="non-main branches: staging"): + await ogt._apply_type_schema( + _client(), [_type_action("insert", "B", fragment)] + ) + assert calls == ["apply"] + + @pytest.mark.asyncio + async def test_releasing_ownership_drops_only_the_ownership_property( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + marked = " coco_key: String\n coco_managed_by_og: Bool?\n" + existing = ( + f"node Person {{\n slug: String @key\n{marked}}}\n\n" + f"node Company {{\n slug: String @key\n{marked}}}\n" + ) + applied: list[str] = [] + + async def fake_read_schema(self: object) -> str | None: + return existing + + async def record_apply(self: object, pg_fragment: str) -> None: + applied.append(pg_fragment) + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "apply_schema", record_apply) + + await ogt._apply_type_schema( + _client(), + [_type_action(None, "Person", "unused", release_ownership=True)], + ) + assert applied == [ + ( + "node Person {\n slug: String @key\n coco_key: String\n}\n\n" + f"node Company {{\n slug: String @key\n{marked}}}\n" + ) + ] + + @pytest.mark.asyncio + async def test_inits_when_graph_absent( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + calls: list[str] = [] + + async def fake_read_schema(self: object) -> str | None: + return None + + async def fake_init_graph(self: object, pg_fragment: str) -> None: + calls.append(f"init:{pg_fragment}") + + async def fake_apply_schema(self: object, pg_fragment: str) -> None: + calls.append(f"apply:{pg_fragment}") + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "init_graph", fake_init_graph) + monkeypatch.setattr(_CliClient, "apply_schema", fake_apply_schema) + + fragment = "node A {\n slug: String @key\n coco_key: String\n}" + await ogt._apply_type_schema(_client(), [_type_action("insert", "A", fragment)]) + assert calls == [f"init:{fragment}\n"] + + @pytest.mark.asyncio + async def test_merges_when_graph_already_exists( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The exact bug this fix closes: a second type must be merged into + the existing schema, not applied alone (which would silently wipe + every other type β€” verified live) or re-inited (which fails + outright on an already-initialized store).""" + calls: list[str] = [] + existing = "node A {\n slug: String @key\n coco_key: String\n}\n" + + async def fake_read_schema(self: object) -> str | None: + return existing + + async def fake_init_graph(self: object, pg_fragment: str) -> None: + raise AssertionError("must not re-init an already-initialized graph") + + async def fake_apply_schema(self: object, pg_fragment: str) -> None: + calls.append(pg_fragment) + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "init_graph", fake_init_graph) + monkeypatch.setattr(_CliClient, "apply_schema", fake_apply_schema) + + b_fragment = "node B {\n slug: String @key\n coco_key: String\n}" + await ogt._apply_type_schema( + _client(), [_type_action("insert", "B", b_fragment)] + ) + (applied,) = calls + assert "node A {" in applied + assert "node B {" in applied + + @pytest.mark.asyncio + async def test_concurrent_updates_preserve_both_schema_changes( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + schema = "node A {\n slug: String @key\n coco_key: String\n}\n" + + async def fake_read_schema(self: object) -> str | None: + await asyncio.sleep(0) + return schema + + async def fake_apply_schema(self: object, pg_fragment: str) -> None: + nonlocal schema + await asyncio.sleep(0) + schema = pg_fragment + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "apply_schema", fake_apply_schema) + + b_fragment = "node B {\n slug: String @key\n coco_key: String\n}" + c_fragment = "node C {\n slug: String @key\n coco_key: String\n}" + await asyncio.gather( + ogt._apply_type_schema( + _client(), [_type_action("insert", "B", b_fragment)] + ), + ogt._apply_type_schema( + _client(), [_type_action("insert", "C", c_fragment)] + ), + ) + + assert "node A {" in schema + assert "node B {" in schema + assert "node C {" in schema + + @pytest.mark.asyncio + async def test_replace_drops_then_re_adds_instead_of_merging_in_place( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A `@key` change can't go through the ordinary single-call merge + β€” the engine rejects it outright, verified live β€” so "replace" + must drive two separate `apply_schema` calls: one with the type's + block removed, one with it re-added under its new definition. + Other types must survive both calls untouched.""" + calls: list[str] = [] + a_old = "node A {\n slug: String @key\n coco_key: String\n}\n" + b = "node B {\n slug: String @key\n coco_key: String\n}\n" + existing = a_old + "\n" + b + + async def fake_read_schema(self: object) -> str | None: + return existing + + async def fake_init_graph(self: object, pg_fragment: str) -> None: + raise AssertionError("replace must not init") + + async def fake_apply_schema(self: object, pg_fragment: str) -> None: + calls.append(pg_fragment) + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "init_graph", fake_init_graph) + monkeypatch.setattr(_CliClient, "apply_schema", fake_apply_schema) + + a_new = "node A {\n title: String @key\n coco_key: String\n}" + await ogt._apply_type_schema(_client(), [_type_action("replace", "A", a_new)]) + + assert len(calls) == 2 + drop_call, readd_call = calls + assert "node A" not in drop_call + assert "node B {" in drop_call + assert a_new in readd_call + assert "node B {" in readd_call + + @pytest.mark.asyncio + async def test_removing_an_edge_leaves_a_same_named_node_alone( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The removal half of the same collision: dropping the edge used + to delete the NODE's block and leave the edge in place.""" + calls: list[str] = [] + existing = ( + "node Link {\n slug: String @key\n coco_key: String\n}\n\n" + "edge Link: A -> B {\n coco_key: String\n}\n" + ) + + async def fake_read_schema(self: object) -> str | None: + return existing + + async def fake_apply_schema(self: object, pg_fragment: str) -> None: + calls.append(pg_fragment) + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "apply_schema", fake_apply_schema) + + await ogt._apply_type_schema( + _client(), [_type_action("delete", "Link", None, type_kind="edge")] + ) + (applied,) = calls + assert "node Link {" in applied + assert "edge Link" not in applied + + @pytest.mark.asyncio + async def test_a_name_used_as_both_kinds_is_refused( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """`schema apply` accepts `node Link` and `edge Link` together, but + a mutation resolves the name to the NODE type alone β€” `insert Link + { from: ..., to: ... }` fails with "type `Link` has no property + `from`" (verified against the binary). The edge type would exist and + be unwritable, so declaring the second one has to fail loudly.""" + + async def fake_read_schema(self: object) -> str | None: + return "node Link {\n slug: String @key\n coco_key: String\n}\n" + + async def fail_apply(self: object, pg_fragment: str) -> None: + raise AssertionError("must not write a schema with a clashing name") + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "apply_schema", fail_apply) + + frag = render_edge_type("Link", "A", "B", [], owner="og") + with pytest.raises(ValueError, match="already used by a node type"): + await ogt._apply_type_schema( + _client(), [_type_action("insert", "Link", frag, type_kind="edge")] + ) + + @pytest.mark.asyncio + async def test_a_name_used_as_both_kinds_in_one_batch_is_refused_on_a_fresh_graph( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The clash check used to sit AFTER the uninitialized-graph early + return, so on a fresh graph it was never reached at all β€” the sink + happily `init`ed a schema containing both `node Link` and `edge Link`, + creating the permanently-unwritable edge type the check exists to + prevent.""" + + async def fake_read_schema(self: object) -> str | None: + return None + + async def fail_init(self: object, pg_fragment: str) -> None: + raise AssertionError("must not init a schema with a clashing name") + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "init_graph", fail_init) + + with pytest.raises(ValueError, match="both a node and an edge"): + await ogt._apply_type_schema( + _client(), + [ + _type_action( + "insert", "Link", "node Link {\n slug: String @key\n}" + ), + _type_action( + "insert", + "Link", + render_edge_type("Link", "A", "B", [], owner="og"), + type_kind="edge", + ), + ], + ) + + @pytest.mark.asyncio + async def test_a_name_used_as_both_kinds_in_one_batch_is_refused( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The live-schema check can only see types ALREADY in the graph, so + a node and an edge of the same name arriving together in one sync slip + past it β€” neither is in `existing` yet.""" + + async def fake_read_schema(self: object) -> str | None: + return "node Other {\n slug: String @key\n coco_key: String\n}\n" + + async def fail_apply(self: object, pg_fragment: str) -> None: + raise AssertionError("must not write a schema with a clashing name") + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "apply_schema", fail_apply) + + with pytest.raises(ValueError, match="both a node and an edge"): + await ogt._apply_type_schema( + _client(), + [ + _type_action( + "insert", "Link", "node Link {\n slug: String @key\n}" + ), + _type_action( + "insert", + "Link", + render_edge_type("Link", "A", "B", [], owner="og"), + type_kind="edge", + ), + ], + ) + + @pytest.mark.asyncio + async def test_drop_removes_the_type_block( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A "drop" used to be a silent no-op: the sink returned no child + handler and wrote nothing. But when a container target state + disappears the engine emits no per-child deletes, so this apply is + the ONLY thing that removes a dropped type's rows β€” without it, the + type and every node in it persist forever, untracked and + unreachable.""" + calls: list[str] = [] + existing = ( + "node A {\n slug: String @key\n coco_key: String\n}\n\n" + "node B {\n slug: String @key\n coco_key: String\n}\n" + ) + + async def fake_read_schema(self: object) -> str | None: + return existing + + async def fake_apply_schema(self: object, pg_fragment: str) -> None: + calls.append(pg_fragment) + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "apply_schema", fake_apply_schema) + + await ogt._apply_type_schema(_client(), [_type_action("delete", "A", None)]) + (applied,) = calls + assert "node A" not in applied + assert "node B {" in applied + + @pytest.mark.asyncio + async def test_batch_reads_once_and_writes_once( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Every type action for one graph arrives in a single sink call, so + the batch costs one read-merge-write, not one per action. Folding + them in memory is also what makes the result correct: each fragment + merges into the previous fold rather than into a stale re-read.""" + reads, applies = [0], [] + existing = "node A {\n slug: String @key\n coco_key: String\n}\n" + + async def fake_read_schema(self: object) -> str | None: + reads[0] += 1 + return existing + + async def fake_apply_schema(self: object, pg_fragment: str) -> None: + applies.append(pg_fragment) + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "apply_schema", fake_apply_schema) + + await ogt._apply_type_schema( + _client(), + [ + _type_action( + "insert", + "B", + "node B {\n slug: String @key\n coco_key: String\n}", + ), + _type_action( + "insert", + "C", + "node C {\n slug: String @key\n coco_key: String\n}", + ), + _type_action("delete", "A", None), + ], + ) + assert reads == [1] + assert len(applies) == 1 + assert "node A" not in applies[0] + assert "node B {" in applies[0] and "node C {" in applies[0] + + @pytest.mark.asyncio + async def test_noop_batch_does_not_touch_the_schema( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """An all-noop batch must not even read the schema β€” an unchanged + run is the common case and should cost zero CLI invocations.""" + + async def fail(self: object, *args: object) -> None: + raise AssertionError("a noop batch must issue no schema calls") + + monkeypatch.setattr(_CliClient, "read_schema", fail) + monkeypatch.setattr(_CliClient, "apply_schema", fail) + monkeypatch.setattr(_CliClient, "init_graph", fail) + + await ogt._apply_type_schema(_client(), [_type_action(None, "A", "node A {}")]) + + @pytest.mark.asyncio + async def test_batch_with_a_replace_lands_the_removal_first( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The one case that still needs two writes: a replaced type must be + gone from a COMMITTED schema before it comes back under its new + definition. Types changing alongside it ride along in the second + apply, not a third.""" + calls: list[str] = [] + existing = ( + "node A {\n slug: String @key\n coco_key: String\n}\n\n" + "node B {\n slug: String @key\n coco_key: String\n}\n" + ) + + async def fake_read_schema(self: object) -> str | None: + return existing + + async def fake_apply_schema(self: object, pg_fragment: str) -> None: + calls.append(pg_fragment) + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "apply_schema", fake_apply_schema) + + a_new = "node A {\n title: String @key\n coco_key: String\n}" + b_new = "node B {\n slug: String @key\n note: String?\n coco_key: String\n}" + await ogt._apply_type_schema( + _client(), + [ + _type_action("replace", "A", a_new), + _type_action(None, "B", b_new, property_actions={"prop:x": "insert"}), + ], + ) + + assert len(calls) == 2 + removal, final = calls + assert "node A" not in removal + assert "node B {" in removal + assert a_new in final + assert "note: String?" in final + + +# --------------------------------------------------------------------------- +# Public API surface: NodeTarget/EdgeTarget, the twelve entry points, aliases. +# --------------------------------------------------------------------------- + + +class TestOrphanedEdgeEndpoints: + """Removing a node type that an edge type still points at. + + The engine rejects the resulting schema outright (`catalog error: edge + 'WorksAt' has an unresolved endpoint`, verified against the binary). + Every mounted type is its own processing component, and all of them + share one sink batcher β€” so when an app is dropped, the node type's + removal can run alone while the edge type's removal waits behind it in + the batcher's queue. Waiting inside the sink can therefore never see + the edge go: the node's batch has to take the connector's own edge + types along, and leave everyone else's in place. + """ + + _MARKED = " coco_key: String\n coco_managed_by_og: Bool?\n" + SCHEMA = ( + f"node Person {{\n slug: String @key\n{_MARKED}}}\n" + f"node Company {{\n slug: String @key\n{_MARKED}}}\n" + f"edge WorksAt: Person -> Company {{\n{_MARKED}}}\n" + ) + #: The same graph with an edge type the connector did not create: no + #: ownership property, though it declares `coco_key` the way a + #: `managed_by=user` type must. + SCHEMA_WITH_FOREIGN_EDGE = SCHEMA.replace( + f"edge WorksAt: Person -> Company {{\n{_MARKED}}}\n", + "edge WorksAt: Person -> Company {\n coco_key: String\n}\n", + ) + #: And with an edge type another app of this connector owns. + SCHEMA_WITH_ANOTHER_APPS_EDGE = SCHEMA.replace( + f"edge WorksAt: Person -> Company {{\n{_MARKED}}}\n", + "edge WorksAt: Person -> Company {\n coco_key: String\n" + " coco_managed_by_other: Bool?\n}\n", + ) + + @staticmethod + def _mock(monkeypatch: pytest.MonkeyPatch, schema: str) -> list[str]: + """Serve `schema` on every read; return the list every write lands in.""" + applied: list[str] = [] + + async def fake_read_schema(self: object) -> str | None: + return schema + + async def record_apply(self: object, pg_fragment: str) -> None: + applied.append(pg_fragment) + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "apply_schema", record_apply) + return applied + + @pytest.mark.asyncio + async def test_dropping_a_node_type_takes_its_own_edge_types_along( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Dropping Person while WorksAt is still declared: WorksAt declares + the connector's ownership property, so its own removal is coming (it is queued + behind this very batch when the app is dropped), and taking it + along now is exactly what that removal would do. Company stays.""" + applied = self._mock(monkeypatch, self.SCHEMA) + await ogt._apply_type_schema( + _client(), [_type_action("delete", "Person", None)] + ) + (final,) = applied + assert "node Person" not in final and "edge WorksAt" not in final + assert "node Company" in final + + @pytest.mark.asyncio + async def test_dropping_a_node_type_referenced_by_a_foreign_edge_is_refused( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """No ownership property means the edge type is a user's + (`managed_by=user`, which must declare `coco_key` too) or another + tool's. The connector + never removes those, so the drop fails naming it and writes nothing.""" + applied = self._mock(monkeypatch, self.SCHEMA_WITH_FOREIGN_EDGE) + with pytest.raises(ValueError, match=r"WorksAt.*not managed by this app"): + await ogt._apply_type_schema( + _client(), [_type_action("delete", "Person", None)] + ) + assert applied == [] + + @pytest.mark.asyncio + async def test_dropping_a_node_type_referenced_by_another_apps_edge_is_refused( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The ownership property used to say only that some app of this + connector wrote the block, and a drop took every such edge type + along β€” including one another app was still declaring, whose + removal was never coming. It names the app now: only this app's + edge types go, and the refusal names the app that owns the other.""" + applied = self._mock(monkeypatch, self.SCHEMA_WITH_ANOTHER_APPS_EDGE) + with pytest.raises(ValueError, match=r"WorksAt.*coco_managed_by_other"): + await ogt._apply_type_schema( + _client(), [_type_action("delete", "Person", None)] + ) + assert applied == [] + + @pytest.mark.asyncio + async def test_key_changing_a_referenced_node_type_is_refused( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A `@key` change is a drop-and-recreate: the FIRST of its two + applies lands the schema with the node type removed, so the dangling + endpoint appears there even though the final schema would be sound. + The edge stays declared, owned or not, so it is never taken along.""" + applied = self._mock(monkeypatch, self.SCHEMA) + with pytest.raises(ValueError, match="WorksAt"): + await ogt._apply_type_schema( + _client(), + [ + _type_action( + "replace", + "Person", + "node Person {\n email: String @key\n coco_key: String\n}", + ) + ], + ) + assert applied == [] + + @pytest.mark.asyncio + async def test_dropping_the_edge_type_too_is_allowed( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Removing both together leaves nothing dangling, so it must pass.""" + applied: list[str] = [] + + async def fake_read_schema(self: object) -> str | None: + return self_schema + + self_schema = self.SCHEMA + + async def record_apply(self: object, pg_fragment: str) -> None: + applied.append(pg_fragment) + + monkeypatch.setattr(_CliClient, "read_schema", fake_read_schema) + monkeypatch.setattr(_CliClient, "apply_schema", record_apply) + + await ogt._apply_type_schema( + _client(), + [ + _type_action("delete", "WorksAt", None, type_kind="edge"), + _type_action("delete", "Person", None), + ], + ) + assert applied and "Person" not in applied[-1] + assert "WorksAt" not in applied[-1] + + @pytest.mark.asyncio + async def test_removing_an_edge_type_already_taken_along_writes_nothing( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The edge type's own removal, arriving after the node's batch took + it along, finds nothing to remove and must not spend a `schema + apply` on an unchanged schema.""" + already_gone = remove_type_from_schema(self.SCHEMA, "edge", "WorksAt") + applied = self._mock(monkeypatch, already_gone) + await ogt._apply_type_schema( + _client(), [_type_action("delete", "WorksAt", None, type_kind="edge")] + ) + assert applied == [] + + +class TestPublicSurface: + def test_exports_one_vocabulary(self) -> None: + """Omnigraph's own schema language says `node` and `edge`, and that is + the only vocabulary this connector exports. The table/relation/record + aliases doubled every name for no behaviour of their own, and the + mount-time `key=` could only agree with the schema or be an error.""" + assert set(omnigraph.__all__) == { + "ConnectionFactory", + "EdgeSchema", + "EdgeTarget", + "NodeSchema", + "NodeTarget", + "OmnigraphType", + "PropertyDef", + "ValueEncoder", + "declare_edge_target", + "declare_node_target", + "edge_target", + "mount_edge_target", + "mount_node_target", + "node_target", + } + for alias in ( + "TableTarget", + "RelationTarget", + "TableSchema", + "ColumnDef", + "table_target", + "mount_table_target", + "declare_table_target", + "relation_target", + "mount_relation_target", + "declare_relation_target", + ): + assert not hasattr(omnigraph, alias), alias + assert not hasattr(omnigraph.NodeTarget, "declare_record") + assert not hasattr(omnigraph.EdgeTarget, "declare_relation") + + +# --------------------------------------------------------------------------- +# Endpoint key-definition wiring: the stub-and-retry recovery reads +# `_EdgeAction.from_key_property`/`to_key_property`, sourced from +# `_TypeSpec.from_key_property`/`to_key_property`. Without this wiring the +# recovery mechanism cannot preserve the endpoint schema's declared type and +# encoder. See +# `test_endpoint_metadata_is_carried_for_stubs` above for the handler-level +# half of this same contract. +# --------------------------------------------------------------------------- + + +def _in_component(fn: Callable[[], object]) -> None: + """Run `fn` inside a component of a throwaway app: `node_target` and + `edge_target` read the owning app from the component they run in.""" + + async def main() -> None: + fn() + + coco.App( + coco.AppConfig(name=f"in_component_{uuid.uuid4().hex}", environment=coco_env), + main, + ).update_blocking() + + +def test_node_target_needs_a_component() -> None: + """The block a type renders names the app that owns it, and the app is + only known inside a component.""" + schema = NodeSchema( + properties={"slug": PropertyDef("slug", "String")}, key=("slug",) + ) + with pytest.raises(RuntimeError, match="component context"): + omnigraph.node_target(_fresh_db("outside"), "Node", schema) + + +def _bare_node_target(schema: NodeSchema, type_name: str) -> omnigraph.NodeTarget[Any]: + """A NodeTarget with no real provider β€” valid for tests that only read + `.schema`/`.type_name` (spec construction, validation) and never declare + an actual node through it.""" + return omnigraph.NodeTarget(None, schema, type_name) # type: ignore[arg-type] + + +class TestNodeTargetKeyValidation: + def test_a_custom_encoder_on_a_key_property_is_refused(self) -> None: + """With `str.lower`, a declared "Mixed" was stored as "mixed"; + switching the encoder to `str.upper` then upserted a second node + "MIXED" under the same `coco_key` and left "mixed" behind β€” and an + edge declared with the raw "Mixed" would find neither. The key is + the identity: normalize it before declaring the node, never in an + encoder.""" + schema = NodeSchema( + properties={"slug": PropertyDef("slug", "String", str.lower)}, + key=("slug",), + ) + with pytest.raises( + ValueError, match=r"key property 'slug' has a custom encoder" + ): + _in_component( + lambda: omnigraph.node_target(_fresh_db("custom"), "Person", schema) + ) + + def test_a_built_in_encoder_of_another_type_is_refused_on_a_key(self) -> None: + """`datetime.date.isoformat` was accepted on a DateTime key, and on a + datetime it drops the time: noon and 14:00 on one date were stored + as one midnight node while tracked as two instants, and removing + one declaration deleted the node the other still declared. A + built-in encoder is built in for the one type it encodes.""" + for pg_type, encoder in ( + ("DateTime", datetime.date.isoformat), + ("Date", datetime.datetime.isoformat), + ("String", ogt._ENCODERS["DateTime"]), + ): + schema = NodeSchema( + properties={"k": PropertyDef("k", pg_type, encoder)}, key=("k",) + ) + + def mount(schema: NodeSchema = schema) -> None: + omnigraph.node_target(_fresh_db("other"), "T", schema) + + with pytest.raises(ValueError, match=r"key property 'k'"): + _in_component(mount) + + def test_the_built_in_date_encoder_on_a_key_is_allowed(self) -> None: + """Only the fixed `Date`/`DateTime` encoders may sit on a key: they + are a pure function of the value and never change, so raw and + encoded identity cannot drift apart.""" + schema = NodeSchema( + properties={"on": PropertyDef("on", "Date", ogt._ENCODERS["Date"])}, + key=("on",), + ) + _in_component(lambda: omnigraph.node_target(_fresh_db("date"), "Doc", schema)) + + def test_keyless_schema_is_refused(self) -> None: + """A hand-built `NodeSchema` with an empty key must be refused at + mount, not silently collapse every row onto one target state. + + `NodeSchema` is public and in `__all__`; only `from_class` enforces a + non-empty key, and the optional `key=` kwarg is only cross-checked + when supplied. Without this guard the type renders with no `@key`, + every row derives the SAME `coco_key` from the empty key tuple, all + rows share one StableKey (last write wins in tracking), and an + unkeyed Omnigraph insert is a strict insert that duplicates on every + re-run. `_validate_edge_endpoint` already has this guard β€” but only + for types used as endpoints. + """ + schema = NodeSchema(properties={"slug": PropertyDef("slug", "String")}, key=()) + with pytest.raises(ValueError, match="must declare a key"): + _in_component( + lambda: omnigraph.node_target(_fresh_db("keyless"), "Node", schema) + ) + + +class TestRecordToDict: + """`_record_to_dict` is the seam every declared row passes through, so a + field name it silently mishandles becomes an opaque engine error.""" + + @staticmethod + def _schema() -> NodeSchema: + return NodeSchema( + properties={ + "slug": PropertyDef("slug", "String"), + "title": PropertyDef("title", "String?"), + "note": PropertyDef("note", "String"), + }, + key=("slug",), + ) + + def test_unknown_dict_key_is_refused(self) -> None: + """A misspelled field used to vanish: every declared property came + back `None`, the key's `coco_key` was derived from `None`, and the + engine rejected a null into a non-nullable `@key` column with nothing + naming the typo.""" + with pytest.raises(ValueError, match="slgu"): + ogt._record_to_dict({"slgu": "a", "note": "n"}, self._schema()) + + def test_missing_non_nullable_is_refused(self) -> None: + with pytest.raises(ValueError, match="note"): + ogt._record_to_dict({"slug": "a"}, self._schema()) + + def test_missing_nullable_defaults_to_none(self) -> None: + """Omitting a nullable property is legitimate β€” that is what nullable + means β€” so only non-nullable omissions are an error.""" + assert ogt._record_to_dict({"slug": "a", "note": "n"}, self._schema()) == { + "slug": "a", + "title": None, + "note": "n", + } + + +class TestEdgeTargetKeyPropertyWiring: + def test_edge_target_populates_endpoint_key_properties(self) -> None: + # Key-only schemas β€” this test is about key-definition wiring, not the + # stub-compatibility guard (see test_mount_edge_rejects_unstubbable_ + # endpoint below for that), so neither endpoint may carry a + # non-nullable non-key property or `_build_edge_spec` would raise + # before we get there. + source = _bare_node_target( + NodeSchema( + properties={"slug": PropertyDef("slug", "String")}, key=("slug",) + ), + "Source", + ) + claim = _bare_node_target( + NodeSchema( + properties={"slug": PropertyDef("slug", "String")}, key=("slug",) + ), + "Claim", + ) + # `_build_edge_spec` is the pure seam `edge_target()` uses to build + # its `_TypeSpec` β€” asserting on it directly (rather than reaching + # into `coco.TargetState`'s private value) proves the wiring without + # depending on coco internals. + spec = ogt._build_edge_spec( + EdgeSchema(properties=_SUPPORTS_PROPS), + source, + claim, + ManagedBy.SYSTEM, + "og", + ) + assert spec.from_key_property == PropertyDef("slug", "String") + assert spec.to_key_property == PropertyDef("slug", "String") + + # Reconcile the spec through the real handler chain, exactly as the + # engine would, to prove the field actually reaches `_EdgeAction` β€” + # not just `_TypeSpec`. + out = _EdgeTypeHandler().reconcile(_EK, spec, [], False) + assert out is not None + assert not coco.is_non_existence(out.action.spec) + assert out.action.spec.from_key_property == PropertyDef("slug", "String") + assert out.action.spec.to_key_property == PropertyDef("slug", "String") + + handler = _EdgeHandler( + "Supports", + _EK, + "Source", + "Claim", + out.action.spec.from_key_property, + out.action.spec.to_key_property, + _SUPPORTS_PROPS, + ) + edge_out = handler.reconcile( + ("s1", "c1"), _EdgeValue("s1", "c1", {}), [], False + ) + assert edge_out is not None + assert edge_out.action.from_key_property == PropertyDef("slug", "String") + assert edge_out.action.to_key_property == PropertyDef("slug", "String") + + +# --------------------------------------------------------------------------- +# mount-time validation: an endpoint type with a non-nullable, non-key +# property can never be stubbed, so referencing it as an edge endpoint must +# fail at mount time, not mid-sync from an unrelated component. +# --------------------------------------------------------------------------- + + +@dataclass +class _Strict: + slug: str + required: str # non-null, outside the key -> stub would be rejected + + +@dataclass +class _StubbableClaim: + slug: str + + +async def _declare_bad_edge(db: coco.ContextKey[ConnectionFactory]) -> None: + strict = omnigraph.declare_node_target( + db, "Strict", await omnigraph.NodeSchema.from_class(_Strict, key="slug") + ) + claims = omnigraph.declare_node_target( + db, "Claim", await omnigraph.NodeSchema.from_class(_StubbableClaim, key="slug") + ) + # `strict`/`claims` are PendingS (declared, not yet synced) β€” fine for + # this test, since the guard only reads `.schema`/`.type_name`, both + # already in hand at declare time. `mount_edge_target`'s endpoint params + # default to ResolvedS, so mypy flags the mismatch even though nothing + # here depends on the endpoints having actually synced. + await omnigraph.mount_edge_target(db, "Bad", strict, claims) # type: ignore[arg-type] + + +def test_mount_edge_rejects_unstubbable_endpoint() -> None: + """Endpoint validation must fire before any I/O β€” `declare_node_target` + only registers a declaration (no store touched, no CLI invoked), so this + needs no live binary and no OMNIGRAPH_TEST_STORE gate.""" + db = coco.ContextKey[ConnectionFactory](f"test_unstubbable_{uuid.uuid4().hex}") + coco_env.context_provider.provide( + db, ConnectionFactory(store="file:///tmp/never-touched-by-this-test.omni") + ) + app = coco.App( + coco.AppConfig( + name="test_mount_edge_rejects_unstubbable_endpoint", environment=coco_env + ), + _declare_bad_edge, + db, + ) + with pytest.raises(ValueError, match="cannot be referenced as an edge endpoint"): + app.update_blocking() + + +@dataclass +class _DateKeyed: + day: datetime.date + + +async def _declare_date_keyed_edge(db: coco.ContextKey[ConnectionFactory]) -> None: + days = omnigraph.declare_node_target( + db, "Day", await omnigraph.NodeSchema.from_class(_DateKeyed, key="day") + ) + claims = omnigraph.declare_node_target( + db, "Claim", await omnigraph.NodeSchema.from_class(_StubbableClaim, key="slug") + ) + await omnigraph.mount_edge_target(db, "On", claims, days) # type: ignore[arg-type] + + +def test_mount_edge_rejects_an_endpoint_key_it_cannot_render() -> None: + """An edge addresses an endpoint by the node's `id`, a String rendering + of the key value β€” and the engine's rendering of a `Date` key is days + since the epoch (2026-01-05 -> `"20458"`), which nothing on the Python + side reproduces. Refuse the type at mount time rather than emit an + endpoint reference that silently matches no node.""" + db = coco.ContextKey[ConnectionFactory](f"test_datekey_{uuid.uuid4().hex}") + coco_env.context_provider.provide( + db, ConnectionFactory(store="file:///tmp/never-touched-by-this-test.omni") + ) + app = coco.App( + coco.AppConfig(name="test_mount_edge_date_key", environment=coco_env), + _declare_date_keyed_edge, + db, + ) + with pytest.raises(ValueError, match=r"its key 'day' is Date"): + app.update_blocking() + + +# --------------------------------------------------------------------------- +# declare_node/declare_edge behavioral coverage: nothing previously called +# these as instance methods through a real coco.App run β€” TestPublicSurface +# only checks method identity off the *class*, and every other test in this +# file bypasses them by hand-building `_NodeAction`/`_EdgeAction` directly. +# So `_record_to_dict` (dict-vs-dataclass extraction) and `declare_node`'s +# key-tuple construction from `self._schema.key` had zero coverage. These +# drive the real declarative pipeline (mount -> declare -> reconcile -> +# sink) with the CLI client's I/O methods mocked out, and assert on the +# actual `Query` that would be sent β€” the concrete, observable form of +# "what reaches the target state" β€” rather than on any internal method. +# --------------------------------------------------------------------------- + + +async def _noop_schema_call(self: object, schema_pg: str) -> None: + return None + + +async def _noop_read_schema(self: object) -> str | None: + # `None` reads as "graph not yet initialized", routing every one of + # these tests through the (also mocked) `init_graph` path β€” irrelevant + # to what's under test here, which is the declarative layer, not the + # create-vs-alter transport decision. + return None + + +def _patch_cli_schema_calls(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(_CliClient, "init_graph", _noop_schema_call) + monkeypatch.setattr(_CliClient, "apply_schema", _noop_schema_call) + monkeypatch.setattr(_CliClient, "read_schema", _noop_read_schema) + + +def _capture_mutations(monkeypatch: pytest.MonkeyPatch) -> list[Query]: + captured: list[Query] = [] + + async def fake_mutate(self: object, mutation: Query, *, branch: str) -> None: + captured.append(mutation) + + monkeypatch.setattr(_CliClient, "mutate", fake_mutate) + return captured + + +def _patch_cli_branch_calls(monkeypatch: pytest.MonkeyPatch) -> None: + """A brand-new declaration through a real `coco.App` run reports + `prev_may_be_missing=True` (there's genuinely no tracked history yet β€” + see `test_dicts_data_together_insert` in test_component_target_states.py + for the same behavior on an unrelated target), so a first-ever edge + goes through `_EdgeHandler`'s "replace" path (delete-then-insert, two + commits applied via a scratch branch) rather than a bare insert. These + no-op the scratch-branch machinery so it doesn't touch a real process; + `_capture_mutations` still sees both commits.""" + + async def noop_branch_create(self: object, name: str, *, frm: str) -> None: + return None + + async def noop_branch_merge(self: object, name: str, *, into: str) -> None: + return None + + async def noop_branch_delete(self: object, name: str) -> None: + return None + + monkeypatch.setattr(_CliClient, "branch_create", noop_branch_create) + monkeypatch.setattr(_CliClient, "branch_merge", noop_branch_merge) + monkeypatch.setattr(_CliClient, "branch_delete", noop_branch_delete) + + +def _fresh_db(label: str) -> coco.ContextKey[ConnectionFactory]: + db = coco.ContextKey[ConnectionFactory](f"test_{label}_{uuid.uuid4().hex}") + coco_env.context_provider.provide( + db, ConnectionFactory(store="file:///tmp/never-touched-by-this-test.omni") + ) + return db + + +@dataclass +class _Article: + slug: str + title: str + published: datetime.date + + +async def _declare_article_from_dataclass( + db: coco.ContextKey[ConnectionFactory], +) -> None: + schema = await NodeSchema.from_class(_Article, key="slug") + target = await omnigraph.mount_node_target(db, "Article", schema) + target.declare_node( + node=_Article(slug="a1", title="Hello", published=datetime.date(2026, 1, 1)) + ) + + +def test_declare_node_from_dataclass_reaches_the_target_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Exercises the encoder path too: `published` is a `datetime.date`, + which would blow up at `json.dumps` if `_record_to_dict` skipped the + schema-driven extraction and the raw `date` object reached the mutation + unencoded.""" + _patch_cli_schema_calls(monkeypatch) + captured = _capture_mutations(monkeypatch) + db = _fresh_db("declare_node_dataclass") + + app = coco.App( + coco.AppConfig( + name="test_declare_node_from_dataclass_reaches_the_target_state", + environment=coco_env, + ), + _declare_article_from_dataclass, + db, + ) + app.update_blocking() + + assert len(captured) == 1 + m = captured[0] + assert "insert Article" in m.expr + assert m.params["s0_p_slug"] == "a1" + assert m.params["s0_p_title"] == "Hello" + assert m.params["s0_p_published"] == "2026-01-01" # encoder applied + assert m.params["s0_p_coco_key"] == derive_coco_key(("a1",)) + + +async def _declare_article_from_dict(db: coco.ContextKey[ConnectionFactory]) -> None: + schema = await NodeSchema.from_class(_Article, key="slug") + target = await omnigraph.mount_node_target(db, "Article", schema) + target.declare_node( + node={"slug": "a2", "title": "World", "published": datetime.date(2026, 2, 2)} + ) + + +def test_declare_node_from_dict_reaches_the_target_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Same schema, same assertions, a plain dict instead of a dataclass + instance β€” `_record_to_dict` must handle both shapes equivalently.""" + _patch_cli_schema_calls(monkeypatch) + captured = _capture_mutations(monkeypatch) + db = _fresh_db("declare_node_dict") + + app = coco.App( + coco.AppConfig( + name="test_declare_node_from_dict_reaches_the_target_state", + environment=coco_env, + ), + _declare_article_from_dict, + db, + ) + app.update_blocking() + + assert len(captured) == 1 + m = captured[0] + assert "insert Article" in m.expr + assert m.params["s0_p_slug"] == "a2" + assert m.params["s0_p_title"] == "World" + assert m.params["s0_p_published"] == "2026-02-02" + assert m.params["s0_p_coco_key"] == derive_coco_key(("a2",)) + + +@dataclass +class _KeyOnlyNode: + slug: str + + +@dataclass +class _EdgeProps: + weight: int + + +async def _declare_supports_edge(db: coco.ContextKey[ConnectionFactory]) -> None: + node_schema = await NodeSchema.from_class(_KeyOnlyNode, key="slug") + source = await omnigraph.mount_node_target(db, "Source", node_schema) + claim = await omnigraph.mount_node_target(db, "Claim", node_schema) + edge_schema = await EdgeSchema.from_class(_EdgeProps) + edge = await omnigraph.mount_edge_target(db, "Supports", source, claim, edge_schema) + edge.declare_edge(from_id="s1", to_id="c1", record=_EdgeProps(weight=7)) + + +def test_declare_edge_reaches_the_target_state(monkeypatch: pytest.MonkeyPatch) -> None: + """`coco_key` for an edge is derived from `(from_id, to_id)`, not from + any property β€” verifies `declare_edge`'s own key construction, not just + `_EdgeHandler.reconcile`'s (already covered elsewhere). + + A first-ever edge declaration through a real `coco.App` run reports + `prev_may_be_missing=True`, so this lands as a "replace" (delete then + insert, two commits) rather than a bare insert β€” see + `_patch_cli_branch_calls`. Only the insert commit is asserted on.""" + _patch_cli_schema_calls(monkeypatch) + _patch_cli_branch_calls(monkeypatch) + captured = _capture_mutations(monkeypatch) + db = _fresh_db("declare_edge") + + app = coco.App( + coco.AppConfig( + name="test_declare_edge_reaches_the_target_state", environment=coco_env + ), + _declare_supports_edge, + db, + ) + app.update_blocking() + + inserts = [m for m in captured if "insert Supports" in m.expr] + assert len(inserts) == 1, captured + m = inserts[0] + assert m.params["s0_e_from"] == "s1" + assert m.params["s0_e_to"] == "c1" + assert m.params["s0_p_weight"] == 7 + assert m.params["s0_p_coco_key"] == derive_coco_key(("s1", "c1")) + + +@dataclass +class _Reading: + sensor: str + at: str + v: float + + +async def _declare_reading(db: coco.ContextKey[ConnectionFactory]) -> None: + schema = await NodeSchema.from_class(_Reading, key="at") + target = await omnigraph.mount_node_target(db, "Reading", schema) + target.declare_node(node=_Reading(sensor="s1", at="2026-01-01", v=1.5)) + + +def test_declare_node_keys_on_the_schema_key_not_the_first_field( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """`declare_node`'s `key = tuple(properties[k] for k in self._schema.key)` + reads the key out of the record by the schema's declared key field. A + wrong implementation β€” taking the first field, say β€” would still produce + *a* key, just the wrong one, so `_Reading` is deliberately keyed on its + SECOND field: `derive_coco_key(("2026-01-01",))` and + `derive_coco_key(("s1",))` are different values.""" + _patch_cli_schema_calls(monkeypatch) + captured = _capture_mutations(monkeypatch) + db = _fresh_db("declare_node_key_field") + + app = coco.App( + coco.AppConfig(name="test_declare_node_key_field", environment=coco_env), + _declare_reading, + db, + ) + app.update_blocking() + + assert len(captured) == 1 + m = captured[0] + assert "insert Reading" in m.expr + assert m.params["s0_p_sensor"] == "s1" + assert m.params["s0_p_at"] == "2026-01-01" + assert m.params["s0_p_v"] == 1.5 + assert m.params["s0_p_coco_key"] == derive_coco_key(("2026-01-01",)) + assert m.params["s0_p_coco_key"] != derive_coco_key(("s1",)) + + +@dataclass +class _Day: + on: datetime.date + note: str | None + + +async def _declare_a_day(db: coco.ContextKey[ConnectionFactory]) -> None: + days = await omnigraph.mount_node_target( + db, "Day", await NodeSchema.from_class(_Day, key="on") + ) + days.declare_node(node=_Day(on=datetime.date(2026, 1, 5), note="n")) + + +class TestDateTimeKeyIdentity: + """Omnigraph identifies a `DateTime`-keyed node by the instant, as epoch + milliseconds: `12:00+00:00` and `14:00+02:00` are one node, sub-millisecond + precision is dropped, and a naive value is read as UTC (all verified + against the engine). Tracking by the ISO spelling gave one graph node two + tracking keys; removing one declaration then deleted the node the other + still declared, and its later unchanged updates left it missing.""" + + AT = PropertyDef("at", "DateTime", ogt._isoformat) + + def test_two_offsets_of_one_instant_share_a_tracking_key(self) -> None: + utc = datetime.datetime(2026, 1, 1, 12, 0, tzinfo=datetime.UTC) + plus_two = datetime.datetime( + 2026, 1, 1, 14, 0, tzinfo=datetime.timezone(datetime.timedelta(hours=2)) + ) + assert ogt._tracking_key_value(self.AT, utc) == ogt._tracking_key_value( + self.AT, plus_two + ) + assert ogt._tracking_key_value(self.AT, utc) == 1767268800000 + + def test_a_naive_value_is_read_as_utc(self) -> None: + naive = datetime.datetime(2026, 1, 1, 12, 0) # noqa: DTZ001 + assert ogt._tracking_key_value(self.AT, naive) == 1767268800000 + + def test_sub_millisecond_precision_is_dropped(self) -> None: + a = datetime.datetime(2026, 1, 1, 12, 0, 0, 100, tzinfo=datetime.UTC) + b = datetime.datetime(2026, 1, 1, 12, 0, 0, 900, tzinfo=datetime.UTC) + c = datetime.datetime(2026, 1, 1, 12, 0, 0, 500000, tzinfo=datetime.UTC) + assert ogt._tracking_key_value(self.AT, a) == ogt._tracking_key_value( + self.AT, b + ) + assert ogt._tracking_key_value(self.AT, a) != ogt._tracking_key_value( + self.AT, c + ) + + def test_a_string_is_refused_for_a_datetime_key(self) -> None: + """A hand-built `PropertyDef("at", "DateTime")` lets a dict row carry + an ISO string, and a string is not normalised: `12:00+00:00` and + `14:00+02:00` got separate tracking keys for one graph node again. + The instant can only be derived from a `datetime`, so anything else + is refused before any state is declared.""" + bare = PropertyDef("at", "DateTime") + with pytest.raises(TypeError, match=r"key property 'at'.*datetime\.datetime"): + ogt._tracking_key_value(bare, "2026-01-01T12:00:00+00:00") + with pytest.raises(TypeError, match=r"key property 'at'.*datetime\.datetime"): + ogt._tracking_key_value(bare, 1767268800000) + + def test_only_a_date_is_accepted_for_a_date_key(self) -> None: + on = PropertyDef("on", "Date") + with pytest.raises(TypeError, match=r"key property 'on'.*datetime\.date"): + ogt._tracking_key_value(on, "2026-01-05") + # A datetime is a date subclass, but its ISO form carries a time. + with pytest.raises(TypeError, match=r"key property 'on'.*datetime\.date"): + ogt._tracking_key_value(on, datetime.datetime(2026, 1, 5, 12)) # noqa: DTZ001 + + def test_declare_node_refuses_a_string_for_a_datetime_key(self) -> None: + target = _bare_node_target( + NodeSchema(properties={"at": PropertyDef("at", "DateTime")}, key=("at",)), + "Event", + ) + with pytest.raises(TypeError, match=r"key property 'at'.*datetime\.datetime"): + target.declare_node(node={"at": "2026-01-01T12:00:00+00:00"}) + + def test_a_date_key_is_still_tracked_by_its_iso_form(self) -> None: + on = PropertyDef("on", "Date", ogt._isoformat) + assert ogt._tracking_key_value(on, datetime.date(2026, 1, 5)) == "2026-01-05" + + +def test_declare_node_keyed_by_a_date_tracks_the_encoded_key( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A `Date` key is legal in `.pg`, and mounting one was allowed β€” but + declaring a node passed the raw `datetime.date` into the target-state + key, which the engine refused (`Unsupported StableKey Python type`). + The key is tracked in its encoded (ISO) form, the same value the + graph is sent, so identity and stored key can never disagree.""" + _patch_cli_schema_calls(monkeypatch) + captured = _capture_mutations(monkeypatch) + db = _fresh_db("date_key") + + app = coco.App( + coco.AppConfig(name="test_declare_node_date_key", environment=coco_env), + _declare_a_day, + db, + ) + app.update_blocking() + + (m,) = captured + assert m.params["s0_p_on"] == "2026-01-05" + assert m.params["s0_p_coco_key"] == derive_coco_key(("2026-01-05",)) + + +@dataclass +class _AttendedRel: + is_organizer: bool + + +def test_declare_edge_rejects_a_non_scalar_endpoint_id() -> None: + """An endpoint is addressed by the node's single key value. A tuple (the + shape a composite key would have had) has no `.pg` type at all, so it + used to surface as a bare `no Omnigraph type mapping for ` + out of `plan_commits`, with nothing naming the declaration.""" + target: omnigraph.EdgeTarget[Any, Any] = omnigraph.EdgeTarget( + cast(Any, None), + None, + "ATTENDED", + _bare_node_target( + NodeSchema( + properties={"slug": PropertyDef("slug", "String")}, key=("slug",) + ), + "Person", + ), + _bare_node_target( + NodeSchema( + properties={"slug": PropertyDef("slug", "String")}, key=("slug",) + ), + "Meeting", + ), + ) + with pytest.raises(TypeError, match=r"to_id=\('m', 1\) is not usable"): + target.declare_edge(from_id="p", to_id=("m", 1)) + # A date maps to a `.pg` type but not to one the endpoint id rendering + # can reproduce, so it is refused here for the same reason + # `mount_edge_target` refuses a Date-keyed endpoint type. + with pytest.raises(TypeError, match="single string or integer"): + target.declare_edge(from_id="p", to_id=datetime.date(2026, 1, 5)) + + +def test_declare_edge_rejects_an_endpoint_id_of_the_wrong_key_type() -> None: + """An endpoint id is rendered with `str()` into the node's `id`. An + integer given for a String-keyed endpoint would silently address + whichever node's slug happens to be those digits, and a string given + for an integer-keyed endpoint can never match a node at all β€” both are + a mismatch against the endpoint type's declared key, caught here where + the `declare_edge` call is on the stack.""" + target: omnigraph.EdgeTarget[Any, Any] = omnigraph.EdgeTarget( + cast(Any, None), + None, + "ATTENDED", + _bare_node_target( + NodeSchema( + properties={"slug": PropertyDef("slug", "String")}, key=("slug",) + ), + "Person", + ), + _bare_node_target( + NodeSchema( + properties={"meeting_id": PropertyDef("meeting_id", "I64")}, + key=("meeting_id",), + ), + "Meeting", + ), + ) + with pytest.raises(TypeError, match=r"from_id=7 .*'Person'.*'slug'.*String"): + target.declare_edge(from_id=7, to_id=1) + with pytest.raises(TypeError, match=r"to_id='m1' .*'Meeting'.*'meeting_id'.*I64"): + target.declare_edge(from_id="p", to_id="m1") + + +def test_declare_edge_without_a_schema_rejects_a_record() -> None: + """A schema-less edge type declares no properties in `.pg`, and the + engine refuses an insert naming one it doesn't have ("type `X` has no + property `y`") β€” so a record here can never be written. It used to reach + `_encode_properties` and die on a bare `KeyError` deep in the sink, long + after the mistake was made. + + This is the one genuine semantic gap between Omnigraph and the Neo4j + connector a ported app hits: Neo4j's MERGE creates relationship + properties on the fly, so the same call works there. + """ + target: omnigraph.EdgeTarget[_AttendedRel, Any] = omnigraph.EdgeTarget( + cast(Any, None), + None, + "ATTENDED", + _bare_node_target( + NodeSchema( + properties={"slug": PropertyDef("slug", "String")}, key=("slug",) + ), + "Person", + ), + _bare_node_target( + NodeSchema( + properties={"slug": PropertyDef("slug", "String")}, key=("slug",) + ), + "Meeting", + ), + ) + with pytest.raises(TypeError, match="mounted without a schema"): + target.declare_edge( + from_id="p", to_id="m", record=_AttendedRel(is_organizer=True) + ) + + +def test_declare_edge_without_a_record_rejects_a_non_nullable_schema() -> None: + """The mirror of the guard above, and the gap it left open: a schema with + no record. + + `declare_edge` checked "record but no schema" and not "schema but no + record", so an edge type declaring non-nullable properties -- exactly the + example app's `AttendedRel(is_organizer: bool)` -- planned `insert + ATTENDED { from: ..., to: ..., coco_key: ... }` with every declared + property missing. Omnigraph rejects that ("must provide non-nullable + property"), naming neither this call nor the omitted argument. + """ + schema = EdgeSchema( + properties={"is_organizer": PropertyDef("is_organizer", "Bool")} + ) + target: omnigraph.EdgeTarget[_AttendedRel, Any] = omnigraph.EdgeTarget( + cast(Any, None), + schema, + "ATTENDED", + _bare_node_target( + NodeSchema( + properties={"slug": PropertyDef("slug", "String")}, key=("slug",) + ), + "Person", + ), + _bare_node_target( + NodeSchema( + properties={"slug": PropertyDef("slug", "String")}, key=("slug",) + ), + "Meeting", + ), + ) + with pytest.raises(TypeError, match="is_organizer"): + target.declare_edge(from_id="p", to_id="m") + + +# --------------------------------------------------------------------------- +# One vocabulary: node/edge. No `row=`, no `declare_record`, and no mount-time +# `key=` β€” the schema already carries the key. +# --------------------------------------------------------------------------- + + +def test_declare_node_takes_only_node() -> None: + schema = NodeSchema( + properties={"slug": PropertyDef("slug", "String")}, key=("slug",) + ) + target = _bare_node_target(schema, "Foo") + with pytest.raises(TypeError): + target.declare_node() # type: ignore[call-arg] + with pytest.raises(TypeError): + target.declare_node(row={"slug": "a"}) # type: ignore[call-arg] + + +async def _mount_with_redundant_key(db: coco.ContextKey[ConnectionFactory]) -> None: + schema = await NodeSchema.from_class(_Article, key="slug") + await omnigraph.mount_node_target(db, "Article", schema, key="slug") # type: ignore[call-arg] + + +def test_mount_node_has_no_key_kwarg(monkeypatch: pytest.MonkeyPatch) -> None: + """The schema already names the key; a second `key=` at the mount call + site could only agree with it or be an error, so it is gone.""" + _patch_cli_schema_calls(monkeypatch) + db = _fresh_db("mount_node_no_key_kwarg") + + app = coco.App( + coco.AppConfig(name="test_mount_node_has_no_key_kwarg", environment=coco_env), + _mount_with_redundant_key, + db, + ) + with pytest.raises(TypeError, match="key"): + app.update_blocking() + + +# --------------------------------------------------------------------------- +# Live-engine tests: require the omnigraph CLI binary at `test/bin/omnigraph` +# (git-ignored). Gated on OMNIGRAPH_TEST_STORE=1 so a checkout without the +# binary still runs the rest; CI installs the pinned release and sets the flag +# (see .github/workflows/_test.yml). +# --------------------------------------------------------------------------- + +_OMNIGRAPH_BIN = str(Path(__file__).resolve().parents[3] / "test" / "bin" / "omnigraph") + +_live = pytest.mark.skipif( + os.environ.get("OMNIGRAPH_TEST_STORE") != "1", + reason="requires the omnigraph CLI binary; set OMNIGRAPH_TEST_STORE=1", +) + + +def _live_conn(store_dir: Path) -> ConnectionFactory: + return ConnectionFactory(store=f"file://{store_dir}", cli=_OMNIGRAPH_BIN) + + +def _live_context(conn: ConnectionFactory) -> tuple[ContextProvider, str]: + db = ContextKey[ConnectionFactory](f"test_db_{uuid.uuid4().hex}") + cp = ContextProvider() + cp.provide(db, conn) + return cp, db.key + + +def _init_live(store_dir: Path, schema_pg: str) -> subprocess.CompletedProcess[str]: + """Run `omnigraph init` with `schema_pg` and hand back the raw result, so + a test can assert on the engine's own refusal rather than on ours.""" + with tempfile.NamedTemporaryFile("w", suffix=".pg", encoding="utf-8") as f: + f.write(schema_pg) + f.flush() + return subprocess.run( + [ + _OMNIGRAPH_BIN, + "init", + "--schema", + f.name, + "--quiet", + f"file://{store_dir}", + ], + capture_output=True, + text=True, + check=False, + ) + + +def _mutate_live( + store: str, mutation: Query, *, check: bool = True +) -> subprocess.CompletedProcess[str]: + r = subprocess.run( + [ + _OMNIGRAPH_BIN, + "mutate", + "--store", + store, + "--branch", + "main", + "--json", + "--quiet", + "-e", + mutation.expr, + "--params", + json.dumps(mutation.params), + ], + capture_output=True, + text=True, + check=False, + ) + if check: + assert r.returncode == 0, r.stderr + return r + + +@_live +class TestEngineSchemaLimitsLive: + """The engine-side facts three Python-side guards rest on. Each guard + exists only because the engine refuses the schema outright, so if a + future Omnigraph release starts accepting one of these, the guard is the + thing that's now wrong β€” and these tests are what say so.""" + + def test_two_keys_are_rejected(self, tmp_path: Path) -> None: + r = _init_live( + tmp_path / "g.omni", + ( + "node Reading {\n" + " sensor: String @key\n" + " at: String @key\n" + " coco_key: String\n" + "}\n" + ), + ) + assert r.returncode != 0 + assert "multiple @key constraints" in r.stderr + + def test_id_property_is_rejected_on_a_node(self, tmp_path: Path) -> None: + r = _init_live( + tmp_path / "g.omni", + ("node Meeting {\n slug: String @key\n id: I64\n coco_key: String\n}\n"), + ) + assert r.returncode != 0 + assert "exactly one top-level `id` field" in r.stderr + + def test_a_node_and_an_edge_may_share_a_name_but_only_one_is_addressable( + self, tmp_path: Path + ) -> None: + """Both halves of the collision, from the engine itself. + + `schema apply` ACCEPTS `node Link` alongside `edge Link` β€” which is + why `_find_type_block` must match on kind, or an edge's fragment + overwrites the node's block. But a mutation resolves the bare name + to the NODE type, so the edge is unwritable β€” which is why + `_check_no_kind_clash` refuses to create the situation at all. If a + future release changes either half, this is what says so. + """ + store_dir = tmp_path / "g.omni" + assert ( + _init_live( + store_dir, + ( + "node A {\n slug: String @key\n coco_key: String\n}\n\n" + "node B {\n slug: String @key\n coco_key: String\n}\n\n" + "node Link {\n slug: String @key\n coco_key: String\n}\n\n" + "edge Link: A -> B {\n coco_key: String\n}\n" + ), + ).returncode + == 0 + ) + + store = f"file://{store_dir}" + _mutate_live( + store, + render_query( + [ + build_node_upsert( + "A", (PropertyValue("slug", "String", "a1"),), "ck-a" + ), + build_node_upsert( + "B", (PropertyValue("slug", "String", "b1"),), "ck-b" + ), + ] + ), + ) + r = _mutate_live( + store, + render_query( + [ + build_edge_insert( + "Link", + PropertyValue("ref", "String", "a1"), + PropertyValue("ref", "String", "b1"), + (), + "ck-e", + ) + ] + ), + check=False, + ) + assert r.returncode != 0 + assert "type `Link` has no property `from`" in r.stderr + + def test_key_only_stub_needs_every_non_nullable_property( + self, tmp_path: Path + ) -> None: + """The premise behind `_validate_edge_endpoint`: the key-only stub + the sink inserts for a missing edge endpoint cannot satisfy a + non-nullable property outside the key, so referencing such a type as + an endpoint has to be refused at mount time.""" + store_dir = tmp_path / "g.omni" + assert ( + _init_live( + store_dir, + ( + "node Meeting {\n" + " slug: String @key\n" + " note_file: String\n" + " coco_key: String\n" + "}\n" + ), + ).returncode + == 0 + ) + r = subprocess.run( + [ + _OMNIGRAPH_BIN, + "mutate", + "--store", + f"file://{store_dir}", + "--branch", + "main", + "--json", + "--quiet", + "-e", + ( + "query m($p_slug: String, $p_coco_key: String) " + "{ insert Meeting { slug: $p_slug, coco_key: $p_coco_key } }" + ), + "--params", + '{"p_slug": "m1", "p_coco_key": "ck1"}', + ], + capture_output=True, + text=True, + check=False, + ) + assert r.returncode != 0 + assert "must provide non-nullable property `note_file`" in r.stderr + + +@_live +class TestFullCapCommitLive: + """A commit at the full `_MAX_ENTITIES_PER_TYPE` must actually reach the + engine. + + Nothing exercised this: every chunking test monkeypatches the cap DOWN + (to 4), so the suite only ever built commits three orders of magnitude + smaller than the one a real component produces. At the real cap the + payload is ~1.3 MB of GQ plus ~1.0 MB of params, and passing those + inline put both in argv β€” `OSError: [Errno 7] Argument list too long` + on darwin, and on Linux the 128 KiB-per-argument limit would have + failed the expression alone somewhere near 800 entities. + + Deliberately the real cap, not a scaled-down stand-in: the number that + matters is the one the connector ships with, and the whole defect was + that nobody multiplied it by the bytes per entity. + """ + + def test_a_commit_at_the_entity_cap_applies(self, tmp_path: Path) -> None: + store = f"file://{tmp_path / 'cap.omni'}" + assert ( + _init_live( + tmp_path / "cap.omni", + ( + "node Source {\n slug: String @key\n title: String?\n" + " coco_key: String\n}\n" + ), + ).returncode + == 0 + ) + + n = ogt._MAX_ENTITIES_PER_TYPE + commit = render_query( + [ + build_node_upsert( + "Source", + ( + PropertyValue("slug", "String", f"s{i:06d}"), + PropertyValue("title", "String?", f"Title {i}"), + ), + derive_coco_key((f"s{i:06d}",)), + ) + for i in range(n) + ] + ) + # The payload really is past the argv limits, so this is not a + # hypothetical: if it ever shrinks below them the test stops + # covering what it exists to cover. + assert len(commit.expr) > 131072, "expression no longer exceeds MAX_ARG_STRLEN" + assert len(commit.expr) + len(json.dumps(commit.params)) > 1048576 + + conn = ConnectionFactory(store=store, cli=_OMNIGRAPH_BIN) + asyncio.run(_CliClient(conn).mutate(commit, branch="main")) + + rows = _export_rows(store) + assert len([r for r in rows if r.get("type") == "Source"]) == n + + +@_live +class TestMutationAtomicityLive: + """A failed multi-statement `mutate` must apply NOTHING. + + Everything above rests on this and nothing pinned it. `render_query` + merges N statements into one query precisely to get one commit, and + `_mutate_with_endpoint_retry` re-runs the WHOLE commit after stubbing a + missing endpoint β€” so if a failed invocation left its earlier statements + applied, every edge before the failure point would be inserted a second + time on the retry. Edge insert is strict and never deduplicates, and the + duplicate would carry the same `coco_key` as the original, making it + invisible to tracking and unreachable by `delete ... where coco_key = $x` + (which removes one row). The three endpoint-retry tests each carry + exactly one edge insert, so none of them can tell the two worlds apart. + """ + + @staticmethod + def _seeded(tmp_path: Path) -> str: + store = f"file://{tmp_path / 'atomic.omni'}" + assert ( + _init_live( + tmp_path / "atomic.omni", + ( + "node Source {\n slug: String @key\n title: String?\n" + " coco_key: String\n}\n\n" + "node Claim {\n slug: String @key\n coco_key: String\n}\n\n" + "edge Supports: Source -> Claim {\n coco_key: String\n}\n" + ), + ).returncode + == 0 + ) + _mutate_live( + store, + render_query( + [ + build_node_upsert( + "Source", (PropertyValue("slug", "String", "s1"),), "ck-s1" + ), + build_node_upsert( + "Claim", (PropertyValue("slug", "String", "c1"),), "ck-c1" + ), + ] + ), + ) + return store + + def test_earlier_statements_do_not_land_when_a_later_one_fails( + self, tmp_path: Path + ) -> None: + store = self._seeded(tmp_path) + before = _commit_count(store) + + good = build_edge_insert( + "Supports", + PropertyValue("ref", "String", "s1"), + PropertyValue("ref", "String", "c1"), + (), + "ck-e1", + ) + # Fails at execution time (not parse time) on a missing endpoint -- + # the same failure mode `_mutate_with_endpoint_retry` recovers from. + doomed = build_edge_insert( + "Supports", + PropertyValue("ref", "String", "s1"), + PropertyValue("ref", "String", "GHOST"), + (), + "ck-e2", + ) + r = _mutate_live(store, render_query([good, doomed]), check=False) + assert r.returncode != 0 + assert "not found in Claim" in r.stderr + + rows = _export_rows(store) + assert [r for r in rows if r.get("type") == "Supports"] == [] + assert _commit_count(store) == before + + def test_a_failed_upsert_batch_leaves_prior_values_untouched( + self, tmp_path: Path + ) -> None: + """The node half of the same property: a node upsert combined ahead + of a failing statement must not have overwritten anything.""" + store = self._seeded(tmp_path) + + r = _mutate_live( + store, + render_query( + [ + build_node_upsert( + "Source", + ( + PropertyValue("slug", "String", "s1"), + PropertyValue("title", "String?", "OVERWRITTEN"), + ), + "ck-s1", + ), + build_node_upsert( + "Claim", (PropertyValue("slug", "String", "c2"),), "ck-c2" + ), + build_edge_insert( + "Supports", + PropertyValue("ref", "String", "GHOST"), + PropertyValue("ref", "String", "c1"), + (), + "ck-e3", + ), + ] + ), + check=False, + ) + assert r.returncode != 0 + + rows = _export_rows(store) + assert [r["data"]["title"] for r in rows if r.get("type") == "Source"] == [None] + assert [r["data"]["slug"] for r in rows if r.get("type") == "Claim"] == ["c1"] + + +@dataclass +class _SimpleNode: + slug: str + title: str | None + + +@_live +class TestApplyTypeActionsLive: + @pytest.mark.asyncio + async def test_alter_recreates_a_graph_whose_directory_was_deleted( + self, tmp_path: Path + ) -> None: + """`read_schema()` against a graph that was never `init`'d returns + `None` β€” happens for real when a run is interrupted between a + type's pre_commit and its `create` action landing. + `_apply_type_actions` must `init_graph` in that case and succeed + rather than raising.""" + store_dir = tmp_path / "g.omni" + conn = _live_conn(store_dir) + db = ContextKey[ConnectionFactory](f"test_db_{uuid.uuid4().hex}") + cp = ContextProvider() + cp.provide(db, conn) + + schema = await NodeSchema.from_class(_SimpleNode, key="slug") + spec = _TypeSpec( + schema=schema, + key=schema.key, + from_type=None, + to_type=None, + managed_by=ManagedBy.SYSTEM, + owner="og", + ) + key = _TypeKey(db.key, "node", "Simple") + pg_fragment = schema.render("Simple", owner="og") + + create_action = ogt._TypeAction( + key, spec, pg_fragment, "insert", {}, owner="og" + ) + out = await ogt._apply_type_actions(cp, [create_action]) + assert out[0] is not None + assert store_dir.exists() + + shutil.rmtree(store_dir) # simulate a run interrupted before create landed + assert not store_dir.exists() + + # What the engine really emits after an interrupted run: the record + # still matches, but `prev_may_be_missing` forces a write anyway. + upsert_action = ogt._TypeAction( + key, spec, pg_fragment, "upsert", {}, owner="og" + ) + out2 = await ogt._apply_type_actions(cp, [upsert_action]) + assert out2[0] is not None + assert store_dir.exists() + + @pytest.mark.asyncio + async def test_a_second_type_is_merged_not_wiping_the_first( + self, tmp_path: Path + ) -> None: + """The bug this reopening exists to fix, reproduced and then + proven closed: Omnigraph's schema is applied whole-graph, not per + type, so a naive per-type apply of the *second* type's fragment + alone silently dropped the first. `init` a graph with one type, + then drive a fresh `create` action for a second type β€” which Task + 6 always decides as `"create"` regardless of whether it's the + graph's first type or its Nth β€” through the real sink, and assert + `schema show` afterward has BOTH types. Fails hard against the old + code: `init_graph` refuses an already-initialized store outright, + and a bare `apply_schema` of just the second type's fragment wipes + the first (both verified independently against the binary).""" + store_dir = tmp_path / "g.omni" + conn = _live_conn(store_dir) + client = _CliClient(conn) + db = ContextKey[ConnectionFactory](f"test_db_{uuid.uuid4().hex}") + cp = ContextProvider() + cp.provide(db, conn) + + a_schema = await NodeSchema.from_class(_SimpleNode, key="slug") + a_spec = _TypeSpec( + schema=a_schema, + key=a_schema.key, + from_type=None, + to_type=None, + managed_by=ManagedBy.SYSTEM, + owner="og", + ) + await ogt._apply_type_actions( + cp, + [ + ogt._TypeAction( + _TypeKey(db.key, "node", "A"), + a_spec, + a_schema.render("A", owner="og"), + "insert", + {}, + owner="og", + ), + ], + ) + + b_schema = await NodeSchema.from_class(_SimpleNode, key="slug") + b_spec = _TypeSpec( + schema=b_schema, + key=b_schema.key, + from_type=None, + to_type=None, + managed_by=ManagedBy.SYSTEM, + owner="og", + ) + out = await ogt._apply_type_actions( + cp, + [ + ogt._TypeAction( + _TypeKey(db.key, "node", "B"), + b_spec, + b_schema.render("B", owner="og"), + "insert", + {}, + owner="og", + ), + ], + ) + assert out[0] is not None + + schema_source = await client.read_schema() + assert schema_source is not None + assert "node A {" in schema_source + assert "node B {" in schema_source + + +def _export_rows(store_uri: str) -> list[dict[str, Any]]: + result = subprocess.run( + [_OMNIGRAPH_BIN, "export", "--store", store_uri, "--branch", "main"], + capture_output=True, + text=True, + check=True, + ) + return [json.loads(line) for line in result.stdout.splitlines() if line.strip()] + + +def _basic_pg(source_extra: str = "") -> str: + return ( + f"node Source {{\n slug: String @key\n{source_extra} coco_key: String\n}}\n\n" + "node Claim {\n slug: String @key\n coco_key: String\n}\n\n" + "edge Supports: Source -> Claim {\n weight: I64\n coco_key: String\n}" + ) + + +def _upsert( + db_key: str, type_name: str, slug: str, extra: PropertyValue | None = None +) -> ogt._NodeAction: + props = (PropertyValue("slug", "String", slug),) + ((extra,) if extra else ()) + return ogt._NodeAction( + "upsert", + _TypeKey(db_key, "node", type_name), + type_name, + props, + derive_coco_key((slug,)), + ) + + +@_live +class TestReplaceOrderingLive: + @pytest.mark.asyncio + async def test_replace_leaves_exactly_one_edge_with_new_value( + self, tmp_path: Path + ) -> None: + """The check that would have caught the original draft's ordering + bug: if phase A (the replaced edge's delete) ran after phase B (its + re-insert) instead of before, this would see zero edges survive β€” + insert-then-delete on one coco_key removes both.""" + store_dir = tmp_path / "g.omni" + conn = _live_conn(store_dir) + store_uri = f"file://{store_dir}" + client = _CliClient(conn) + cp, db_key = _live_context(conn) + await client.init_graph(_basic_pg()) + + # Both endpoints already exist, so the edge insert below needs no + # retry β€” that path is exercised separately in TestEndpointRetryLive. + await ogt._apply_entity_actions( + cp, [_upsert(db_key, "Source", "s1"), _upsert(db_key, "Claim", "c1")] + ) + + edge_key = _TypeKey(db_key, "edge", "Supports") + edge_coco_key = derive_coco_key(("s1", "c1")) + + insert_action = ogt._EdgeAction( + "insert", + edge_key, + "Supports", + edge_coco_key, + "s1", + "c1", + (PropertyValue("weight", "I64", 1),), + "Source", + "Claim", + PropertyDef("slug", "String"), + PropertyDef("slug", "String"), + ) + await ogt._apply_entity_actions(cp, [insert_action]) + + rows = _export_rows(store_uri) + edges = [r for r in rows if r.get("edge") == "Supports"] + assert len(edges) == 1 + assert edges[0]["data"]["weight"] == 1 + + replace_action = ogt._EdgeAction( + "replace", + edge_key, + "Supports", + edge_coco_key, + "s1", + "c1", + (PropertyValue("weight", "I64", 2),), + "Source", + "Claim", + PropertyDef("slug", "String"), + PropertyDef("slug", "String"), + ) + commits = plan_commits([replace_action]) + assert len(commits) == 2 # phase A (delete) + phase B (insert) + await ogt._apply_entity_actions(cp, [replace_action]) + + rows = _export_rows(store_uri) + edges = [r for r in rows if r.get("edge") == "Supports"] + assert len(edges) == 1 + assert edges[0]["data"]["weight"] == 2 + + +@_live +class TestEndpointRetryLive: + @pytest.mark.asyncio + async def test_edge_insert_retries_after_creating_missing_endpoint( + self, tmp_path: Path + ) -> None: + """An edge referencing an endpoint no other component has written + yet must still succeed: the sink builds a stub for just that + endpoint after the engine reports it missing, then retries. Only + one endpoint (Source) is missing here, so this only needs the + first of the two retries the sink budgets β€” see + test_edge_insert_retries_after_creating_both_missing_endpoints for + the case that needs both.""" + store_dir = tmp_path / "g.omni" + conn = _live_conn(store_dir) + store_uri = f"file://{store_dir}" + client = _CliClient(conn) + cp, db_key = _live_context(conn) + await client.init_graph(_basic_pg()) + + await ogt._apply_entity_actions(cp, [_upsert(db_key, "Claim", "c1")]) + + edge_key = _TypeKey(db_key, "edge", "Supports") + insert_action = ogt._EdgeAction( + "insert", + edge_key, + "Supports", + derive_coco_key(("s1", "c1")), + "s1", + "c1", + (PropertyValue("weight", "I64", 1),), + "Source", + "Claim", + PropertyDef("slug", "String"), + PropertyDef("slug", "String"), + ) + await ogt._apply_entity_actions(cp, [insert_action]) + + rows = _export_rows(store_uri) + assert len([r for r in rows if r.get("edge") == "Supports"]) == 1 + sources = [r for r in rows if r.get("type") == "Source"] + assert len(sources) == 1 + assert sources[0]["data"]["slug"] == "s1" + + @pytest.mark.asyncio + async def test_edge_insert_does_not_wipe_existing_nullable_property( + self, tmp_path: Path + ) -> None: + """The regression this redesign exists to fix. Both endpoints + already exist, so the edge insert must succeed on the FIRST try, + with no stub ever touching Source β€” and Source's own nullable + `title`, already written by its owning component, must survive. + This fails against the old unconditional-stub design, which nulled + `title` out on every edge insert (verified against the engine).""" + store_dir = tmp_path / "g.omni" + conn = _live_conn(store_dir) + store_uri = f"file://{store_dir}" + client = _CliClient(conn) + cp, db_key = _live_context(conn) + await client.init_graph(_basic_pg(source_extra=" title: String?\n")) + + await ogt._apply_entity_actions( + cp, + [ + _upsert( + db_key, + "Source", + "s1", + PropertyValue("title", "String", "Real Title"), + ), + _upsert(db_key, "Claim", "c1"), + ], + ) + + edge_key = _TypeKey(db_key, "edge", "Supports") + insert_action = ogt._EdgeAction( + "insert", + edge_key, + "Supports", + derive_coco_key(("s1", "c1")), + "s1", + "c1", + (PropertyValue("weight", "I64", 1),), + "Source", + "Claim", + PropertyDef("slug", "String"), + PropertyDef("slug", "String"), + ) + await ogt._apply_entity_actions(cp, [insert_action]) + + rows = _export_rows(store_uri) + assert len([r for r in rows if r.get("edge") == "Supports"]) == 1 + sources = [r for r in rows if r.get("type") == "Source"] + assert len(sources) == 1 + assert sources[0]["data"]["title"] == "Real Title" + + @pytest.mark.asyncio + async def test_edge_insert_retries_after_creating_both_missing_endpoints( + self, tmp_path: Path + ) -> None: + """Both endpoints absent is ordinary, not exotic: CocoIndex runs up + to 1024 components concurrently, so an edge component can easily + precede both of the components that own its endpoints. The engine + reports only the first missing endpoint per attempt (`src` before + `dst`), so a single retry isn't enough β€” this needs both of the + sink's two retries. Fails against a single-retry implementation: + the first retry fixes `src`, the second attempt then fails on + `dst`, and that would propagate instead of getting its own stub.""" + store_dir = tmp_path / "g.omni" + conn = _live_conn(store_dir) + store_uri = f"file://{store_dir}" + client = _CliClient(conn) + cp, db_key = _live_context(conn) + await client.init_graph(_basic_pg()) + + # Neither Source nor Claim has been written by any other component. + edge_key = _TypeKey(db_key, "edge", "Supports") + insert_action = ogt._EdgeAction( + "insert", + edge_key, + "Supports", + derive_coco_key(("s1", "c1")), + "s1", + "c1", + (PropertyValue("weight", "I64", 1),), + "Source", + "Claim", + PropertyDef("slug", "String"), + PropertyDef("slug", "String"), + ) + await ogt._apply_entity_actions(cp, [insert_action]) + + rows = _export_rows(store_uri) + assert len([r for r in rows if r.get("edge") == "Supports"]) == 1 + assert len([r for r in rows if r.get("type") == "Source"]) == 1 + assert len([r for r in rows if r.get("type") == "Claim"]) == 1 + + @pytest.mark.asyncio + async def test_two_edges_with_four_missing_endpoints(self, tmp_path: Path) -> None: + """One commit carrying two edges whose four endpoints are all absent: + the engine reports one missing endpoint per attempt, so this takes + four stub rounds. A fixed two-round budget failed on the third.""" + store_dir = tmp_path / "g.omni" + conn = _live_conn(store_dir) + store_uri = f"file://{store_dir}" + client = _CliClient(conn) + cp, db_key = _live_context(conn) + await client.init_graph(_basic_pg()) + + edge_key = _TypeKey(db_key, "edge", "Supports") + actions = [ + ogt._EdgeAction( + "insert", + edge_key, + "Supports", + derive_coco_key((s, c)), + s, + c, + (PropertyValue("weight", "I64", 1),), + "Source", + "Claim", + PropertyDef("slug", "String"), + PropertyDef("slug", "String"), + ) + for s, c in (("s1", "c1"), ("s2", "c2")) + ] + await ogt._apply_entity_actions(cp, actions) + + rows = _export_rows(store_uri) + assert len([r for r in rows if r.get("edge") == "Supports"]) == 2 + assert sorted(r["data"]["slug"] for r in rows if r.get("type") == "Source") == [ + "s1", + "s2", + ] + assert sorted(r["data"]["slug"] for r in rows if r.get("type") == "Claim") == [ + "c1", + "c2", + ] + assert _branch_names(store_uri) == ["main"] + + +@_live +class TestScratchBranchLivenessLive: + @pytest.mark.asyncio + async def test_a_schema_apply_through_an_equivalent_uri_waits_for_a_live_branch( + self, tmp_path: Path + ) -> None: + """Process A holds a scratch branch through `file:///.../g.omni`; + process B applies a schema through `file:///..././g.omni`. B used to + take a different lock, hit the engine's non-main-branch refusal, + reap A's live branch, and succeed β€” and A then failed its own + cleanup. B must wait for A instead, and A's cleanup must succeed.""" + store_dir = tmp_path / "g.omni" + store_uri = f"file://{store_dir}" + a = _CliClient(_live_conn(store_dir)) + b = _CliClient( + ConnectionFactory(store=f"file://{tmp_path}/./g.omni", cli=_OMNIGRAPH_BIN) + ) + await a.init_graph(_basic_pg()) + + entered = asyncio.Event() + release = asyncio.Event() + + async def hold_a_scratch_branch() -> None: + async with ogt._scratch_branch(a, frm="main"): + entered.set() + await release.wait() + + holder = asyncio.create_task(hold_a_scratch_branch()) + await entered.wait() + assert len(_branch_names(store_uri)) == 2 + + applier = asyncio.create_task( + ogt._apply_type_schema( + b, + [ + _type_action( + "insert", + "Extra", + "node Extra {\n slug: String @key\n coco_key: String\n}", + ) + ], + ) + ) + await asyncio.sleep(1.0) + assert not applier.done(), "B ran ahead instead of waiting for A's lock" + assert len(_branch_names(store_uri)) == 2 # A's branch was not reaped + + release.set() + await holder # A's own cleanup found its branch and deleted it + await applier + assert _branch_names(store_uri) == ["main"] + assert "node Extra" in _read_schema_source(store_uri) + + +# --------------------------------------------------------------------------- +# End-to-end acceptance tests: drive a real coco.App against a real store, +# the way a user actually would, rather than calling _apply_entity_actions +# or _apply_type_actions directly (as the live tests above do). Every +# scenario declares Source, Claim, and (except where noted) a Supports edge +# TOGETHER in the same run, never a single type in isolation -- the two +# worst defects found while building this connector were both invisible to +# single-type coverage: a keyed insert that turned out to be a full-record +# replace (silently wiping a node's other properties via an endpoint stub), +# and a whole-graph schema apply (a second type's apply alone wiping the +# first). Both only show up once more than one type is in play at once. +# --------------------------------------------------------------------------- + + +@dataclass +class _ScSourceNarrow: + slug: str + # `title` is nullable so Source stays usable as an edge endpoint (see + # test_mount_edge_rejects_unstubbable_endpoint above): the endpoint + # stub the sink builds when an edge races ahead of its node's own + # component can only ever populate the key, so any other non-nullable + # property makes the type unstubbable. + title: str | None + + +@dataclass +class _ScSourceWide: + slug: str + title: str | None + note: str | None + + +@dataclass +class _ScClaim: + slug: str + + +@dataclass +class _ScEdgeProps: + weight: int + + +@dataclass +class _ScMeeting: + """An INTEGER-keyed endpoint, the shape the meeting-notes example uses + (a generated numeric meeting id). Every other e2e type here is + string-keyed, and the endpoint reference bug this covers was invisible + to all of them.""" + + meeting_id: int + note: str | None + + +@dataclass +class _KcSource: + """Dedicated to test_key_change_rebuilds: `title` must be non-nullable + here so it's legal to use as the new `@key` (NodeSchema.from_class + rejects a nullable key field) -- unlike _ScSourceNarrow, this type is + never used as an edge endpoint, so it doesn't need the nullability the + stub-compatibility guard would otherwise require.""" + + slug: str + title: str + + +def _e2e_db(store: str, label: str) -> ContextKey[ConnectionFactory]: + db = ContextKey[ConnectionFactory](f"e2e_{label}_{uuid.uuid4().hex}") + coco_env.context_provider.provide( + db, ConnectionFactory(store=store, cli=_OMNIGRAPH_BIN) + ) + return db + + +def _commit_count(store: str, branch: str = "main") -> int: + out = subprocess.run( + [ + _OMNIGRAPH_BIN, + "commit", + "list", + "--store", + store, + "--branch", + branch, + "--json", + "--quiet", + ], + check=True, + capture_output=True, + text=True, + ) + return len(json.loads(out.stdout)["commits"]) + + +def _read_schema_source(store: str) -> str: + out = subprocess.run( + [_OMNIGRAPH_BIN, "schema", "show", "--store", store, "--json", "--quiet"], + check=True, + capture_output=True, + text=True, + ) + return str(json.loads(out.stdout)["schema_source"]) + + +def _branch_names(store: str) -> list[str]: + out = subprocess.run( + [_OMNIGRAPH_BIN, "branch", "list", "--store", store, "--json", "--quiet"], + check=True, + capture_output=True, + text=True, + ) + return list(json.loads(out.stdout)["branches"]) + + +@pytest.fixture +def store(tmp_path: Path) -> str: + """A `file://` URI for a graph that does not exist yet -- the + connector's own type sink creates it (via `init_graph`) on the first + sync, exactly as it would happen for a real user's first + `app.update()`.""" + return f"file://{tmp_path / 'e2e.omni'}" + + +@_live +class TestEndToEnd: + """Seven acceptance cases, one per design decision in the spec. Each + drives a real `coco.App` -- reused across repeated `update_blocking()` + calls exactly as a real app would be -- and asserts against the store + via `_export_rows`/`_commit_count`/`_branch_names`, never against the + connector's own internal types.""" + + def test_property_added_is_additive(self, store: str) -> None: + """Re-mount Source with a wider dataclass; the already-written + node's existing property must survive the alter, and Claim/Supports + -- declared in the SAME run but otherwise untouched -- must come + through unaffected, which is exactly what a whole-graph-schema + regression would break.""" + db = _e2e_db(store, "prop_added") + wide = {"on": False} + + async def main() -> None: + source_schema = await NodeSchema.from_class( + _ScSourceWide if wide["on"] else _ScSourceNarrow, key="slug" + ) + claim_schema = await NodeSchema.from_class(_ScClaim, key="slug") + edge_schema = await EdgeSchema.from_class(_ScEdgeProps) + + sources = await omnigraph.mount_node_target(db, "Source", source_schema) + claims = await omnigraph.mount_node_target(db, "Claim", claim_schema) + supports = await omnigraph.mount_edge_target( + db, "Supports", sources, claims, edge_schema + ) + node: Any = ( + _ScSourceWide(slug="a", title="A", note=None) + if wide["on"] + else _ScSourceNarrow(slug="a", title="A") + ) + sources.declare_node(node=node) + claims.declare_node(node=_ScClaim(slug="c1")) + supports.declare_edge( + from_id="a", to_id="c1", record=_ScEdgeProps(weight=1) + ) + + app = coco.App( + coco.AppConfig(name="e2e_prop_added", environment=coco_env), main + ) + app.update_blocking() + + rows = _export_rows(store) + assert [r["data"]["title"] for r in rows if r.get("type") == "Source"] == ["A"] + assert len([r for r in rows if r.get("type") == "Claim"]) == 1 + assert len([r for r in rows if r.get("edge") == "Supports"]) == 1 + + wide["on"] = True + app.update_blocking() + + rows = _export_rows(store) + (source,) = [r for r in rows if r.get("type") == "Source"] + assert source["data"]["title"] == "A" # old value survived the alter + assert source["data"]["note"] is None # new column present + assert len([r for r in rows if r.get("type") == "Claim"]) == 1 # untouched + assert len([r for r in rows if r.get("edge") == "Supports"]) == 1 # untouched + + def test_property_dropped_forces_reupsert(self, store: str) -> None: + """Dropping a property is lossy: the design says every existing + node of that type must be re-upserted, not merely the type's own + schema altered. Two Source rows, so a bug that rewrites zero (or + only one) of them is distinguishable from rewriting both -- proven + via the commit count, not just the resulting values, since + Omnigraph's soft-drop already hides a dropped property from + `export` on its own, so row contents alone can't tell "dropped and + re-upserted" apart from "dropped and left alone".""" + db = _e2e_db(store, "prop_dropped") + narrow = {"on": False} + + async def main() -> None: + source_schema = await NodeSchema.from_class( + _ScSourceNarrow if narrow["on"] else _ScSourceWide, key="slug" + ) + claim_schema = await NodeSchema.from_class(_ScClaim, key="slug") + edge_schema = await EdgeSchema.from_class(_ScEdgeProps) + + sources = await omnigraph.mount_node_target(db, "Source", source_schema) + claims = await omnigraph.mount_node_target(db, "Claim", claim_schema) + supports = await omnigraph.mount_edge_target( + db, "Supports", sources, claims, edge_schema + ) + for slug, title, note in [("a", "A", "nA"), ("b", "B", "nB")]: + node: Any = ( + _ScSourceNarrow(slug=slug, title=title) + if narrow["on"] + else _ScSourceWide(slug=slug, title=title, note=note) + ) + sources.declare_node(node=node) + claims.declare_node(node=_ScClaim(slug="c1")) + supports.declare_edge( + from_id="a", to_id="c1", record=_ScEdgeProps(weight=1) + ) + + app = coco.App( + coco.AppConfig(name="e2e_prop_dropped", environment=coco_env), main + ) + app.update_blocking() + + rows = _export_rows(store) + sources_before = { + r["data"]["slug"]: r["data"] for r in rows if r.get("type") == "Source" + } + assert sources_before["a"]["note"] == "nA" + assert sources_before["b"]["note"] == "nB" + + before = _commit_count(store) + narrow["on"] = True + app.update_blocking() + after = _commit_count(store) + + # One commit for Source's own schema alter (drops the column), one + # for the batched re-upsert of both rows -- upserts of the same + # type always land in a single commit (see + # test_upsert_only_is_one_commit), so this delta is exact: if + # either row were skipped, or if only the schema changed without a + # re-upsert, the delta would be 1, not 2. + assert after - before == 2 + + rows = _export_rows(store) + sources_after = { + r["data"]["slug"]: r["data"] for r in rows if r.get("type") == "Source" + } + assert set(sources_after) == {"a", "b"} + assert ( + sources_after["a"]["title"] == "A" + and sources_after["a"].get("note") is None + ) + assert ( + sources_after["b"]["title"] == "B" + and sources_after["b"].get("note") is None + ) + assert len([r for r in rows if r.get("type") == "Claim"]) == 1 # untouched + assert len([r for r in rows if r.get("edge") == "Supports"]) == 1 # untouched + + def test_key_change_rebuilds(self, store: str) -> None: + """`@key` changing is destructive: the type action is decided as a + `replace` specifically because an in-place `alter` can't legally + change which property is the key, and existing nodes must be + re-declared under the new key rather than left as orphaned rows + nothing tracks. Claim and a Claim -> Claim `Cites` edge are + declared alongside Source and must survive the replace untouched.""" + db = _e2e_db(store, "key_change") + keyed_by_title = {"on": False} + + async def main() -> None: + source_schema = await NodeSchema.from_class( + _KcSource, key=("title" if keyed_by_title["on"] else "slug") + ) + claim_schema = await NodeSchema.from_class(_ScClaim, key="slug") + edge_schema = await EdgeSchema.from_class(_ScEdgeProps) + + sources = await omnigraph.mount_node_target(db, "Source", source_schema) + claims = await omnigraph.mount_node_target(db, "Claim", claim_schema) + cites = await omnigraph.mount_edge_target( + db, "Cites", claims, claims, edge_schema + ) + sources.declare_node(node=_KcSource(slug="a", title="A")) + claims.declare_node(node=_ScClaim(slug="c1")) + claims.declare_node(node=_ScClaim(slug="c2")) + cites.declare_edge(from_id="c1", to_id="c2", record=_ScEdgeProps(weight=1)) + + app = coco.App( + coco.AppConfig(name="e2e_key_change", environment=coco_env), main + ) + app.update_blocking() + + rows = _export_rows(store) + (before_source,) = [r for r in rows if r.get("type") == "Source"] + assert before_source["data"]["slug"] == "a" + assert before_source["data"]["title"] == "A" + + keyed_by_title["on"] = True + app.update_blocking() + + rows = _export_rows(store) + sources_after = [r for r in rows if r.get("type") == "Source"] + assert len(sources_after) == 1 # rebuilt in place, not duplicated + (after_source,) = sources_after + assert after_source["data"]["slug"] == "a" + assert after_source["data"]["title"] == "A" + assert after_source["data"]["coco_key"] != before_source["data"]["coco_key"] + + assert len([r for r in rows if r.get("type") == "Claim"]) == 2 # untouched + assert len([r for r in rows if r.get("edge") == "Cites"]) == 1 # untouched + + def test_edge_repointed(self, store: str) -> None: + """Edges have no update mutation -- a changed edge is delete-then- + insert on the same coco_key (plan_commits' phase ordering). Point + a -> c1, then repoint the same source to c2: exactly one edge must + survive, pointing at the new target, not two edges and not the + deleted one.""" + db = _e2e_db(store, "edge_repointed") + target = {"claim": "c1"} + + async def main() -> None: + source_schema = await NodeSchema.from_class(_ScSourceNarrow, key="slug") + claim_schema = await NodeSchema.from_class(_ScClaim, key="slug") + edge_schema = await EdgeSchema.from_class(_ScEdgeProps) + + sources = await omnigraph.mount_node_target(db, "Source", source_schema) + claims = await omnigraph.mount_node_target(db, "Claim", claim_schema) + supports = await omnigraph.mount_edge_target( + db, "Supports", sources, claims, edge_schema + ) + sources.declare_node(node=_ScSourceNarrow(slug="a", title="A")) + claims.declare_node(node=_ScClaim(slug="c1")) + claims.declare_node(node=_ScClaim(slug="c2")) + supports.declare_edge( + from_id="a", to_id=target["claim"], record=_ScEdgeProps(weight=1) + ) + + app = coco.App( + coco.AppConfig(name="e2e_edge_repointed", environment=coco_env), main + ) + app.update_blocking() + + rows = _export_rows(store) + (edge,) = [r for r in rows if r.get("edge") == "Supports"] + assert edge["to"] == "c1" + + target["claim"] = "c2" + app.update_blocking() + + rows = _export_rows(store) + edges = [r for r in rows if r.get("edge") == "Supports"] + assert len(edges) == 1 + assert edges[0]["from"] == "a" and edges[0]["to"] == "c2" + assert ( + len([r for r in rows if r.get("type") == "Claim"]) == 2 + ) # both endpoints remain + + def test_node_deleted_cascades_to_its_edges(self, store: str) -> None: + """Undeclaring a node together with the edges that reference it + removes both (a node undeclared while an edge still references it + is another matter: test_undeclaring_a_referenced_node_keeps_its_edge). + This isn't ordinary CocoIndex parent-child cleanup -- Source and + Supports are independent target-state trees, so nothing + automatically knows an edge references a node from a SEPARATE + component going away. It's `plan_commits`' own edge-before-node + delete ordering (see test_edge_deletes_precede_node_deletes) that + makes this safe against the real engine, which this proves live: + two Source nodes, each with an edge to the same Claim, so deleting + one leaves the other's node and edge provably untouched.""" + db = _e2e_db(store, "node_deleted") + keep_b = {"on": True} + + async def main() -> None: + source_schema = await NodeSchema.from_class(_ScSourceNarrow, key="slug") + claim_schema = await NodeSchema.from_class(_ScClaim, key="slug") + edge_schema = await EdgeSchema.from_class(_ScEdgeProps) + + sources = await omnigraph.mount_node_target(db, "Source", source_schema) + claims = await omnigraph.mount_node_target(db, "Claim", claim_schema) + supports = await omnigraph.mount_edge_target( + db, "Supports", sources, claims, edge_schema + ) + claims.declare_node(node=_ScClaim(slug="c1")) + sources.declare_node(node=_ScSourceNarrow(slug="a", title="A")) + supports.declare_edge( + from_id="a", to_id="c1", record=_ScEdgeProps(weight=1) + ) + if keep_b["on"]: + sources.declare_node(node=_ScSourceNarrow(slug="b", title="B")) + supports.declare_edge( + from_id="b", to_id="c1", record=_ScEdgeProps(weight=2) + ) + + app = coco.App( + coco.AppConfig(name="e2e_node_deleted", environment=coco_env), main + ) + app.update_blocking() + + rows = _export_rows(store) + assert len([r for r in rows if r.get("type") == "Source"]) == 2 + assert len([r for r in rows if r.get("edge") == "Supports"]) == 2 + + keep_b["on"] = False + app.update_blocking() + + rows = _export_rows(store) + sources_after = [r for r in rows if r.get("type") == "Source"] + assert [r["data"]["slug"] for r in sources_after] == ["a"] + edges_after = [r for r in rows if r.get("edge") == "Supports"] + assert len(edges_after) == 1 + assert edges_after[0]["from"] == "a" and edges_after[0]["to"] == "c1" + + def test_undeclaring_a_referenced_node_keeps_its_edge(self, store: str) -> None: + """The edge lives in one component and its source node in another. + The node's component stops declaring the node, then declares it + again, while the edge's component declares the edge unchanged + throughout. Deleting the node cascaded to the edge in the graph + while the edge's tracking still said it was there, so when the + node came back the edge's reconcile had nothing to do, and the + graph ended with both endpoints and no edge. A node an edge still + references is reduced to a key-only stub instead β€” the shape an + edge arriving before its node leaves behind β€” and the owner's next + declaration fills it back in.""" + db = _e2e_db(store, "referenced_node") + declare_a = {"on": True} + + @coco.fn + async def declare_source(sources: omnigraph.NodeTarget[Any]) -> None: + if declare_a["on"]: + sources.declare_node(node=_ScSourceNarrow(slug="a", title="A")) + + @coco.fn + async def declare_support( + supports: omnigraph.EdgeTarget[Any], claims: omnigraph.NodeTarget[Any] + ) -> None: + claims.declare_node(node=_ScClaim(slug="c1")) + supports.declare_edge( + from_id="a", to_id="c1", record=_ScEdgeProps(weight=1) + ) + + async def main() -> None: + sources = await omnigraph.mount_node_target( + db, "Source", await NodeSchema.from_class(_ScSourceNarrow, key="slug") + ) + claims = await omnigraph.mount_node_target( + db, "Claim", await NodeSchema.from_class(_ScClaim, key="slug") + ) + supports = await omnigraph.mount_edge_target( + db, + "Supports", + sources, + claims, + await EdgeSchema.from_class(_ScEdgeProps), + ) + await coco.mount(declare_source, sources) + await coco.mount(declare_support, supports, claims) + + def graph() -> tuple[list[tuple[str, str | None]], int]: + rows = _export_rows(store) + sources = [ + (r["data"]["slug"], r["data"]["title"]) + for r in rows + if r.get("type") == "Source" + ] + return sources, len([r for r in rows if r.get("edge") == "Supports"]) + + app = coco.App( + coco.AppConfig(name="e2e_referenced_node", environment=coco_env), main + ) + app.update_blocking() + assert graph() == ([("a", "A")], 1) + + declare_a["on"] = False + app.update_blocking() + assert graph() == ([("a", None)], 1) + + declare_a["on"] = True + app.update_blocking() + assert graph() == ([("a", "A")], 1) + + def test_unchanged_issues_no_writes(self, store: str) -> None: + """The single most important case: CocoIndex's whole incrementality + promise rests on an unchanged source producing zero writes. Checked + via the commit count, not by trusting an absence of exceptions -- a + silently reissued no-op upsert would still pass a weaker check. + Three types declared together so a memoization bug in one type + can't hide behind a passing check on another.""" + db = _e2e_db(store, "unchanged") + + async def main() -> None: + source_schema = await NodeSchema.from_class(_ScSourceNarrow, key="slug") + claim_schema = await NodeSchema.from_class(_ScClaim, key="slug") + edge_schema = await EdgeSchema.from_class(_ScEdgeProps) + + sources = await omnigraph.mount_node_target(db, "Source", source_schema) + claims = await omnigraph.mount_node_target(db, "Claim", claim_schema) + supports = await omnigraph.mount_edge_target( + db, "Supports", sources, claims, edge_schema + ) + sources.declare_node(node=_ScSourceNarrow(slug="a", title="A")) + claims.declare_node(node=_ScClaim(slug="c1")) + supports.declare_edge( + from_id="a", to_id="c1", record=_ScEdgeProps(weight=1) + ) + + app = coco.App(coco.AppConfig(name="e2e_unchanged", environment=coco_env), main) + app.update_blocking() + + before = _commit_count(store) + app.update_blocking() + after = _commit_count(store) + + assert after == before + + def test_oversized_component_uses_branch( + self, store: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Force the per-type chunking cap down so a batch of new Source + rows plans into multiple upsert commits (plan_commits' chunking via + _MAX_ENTITIES_PER_TYPE). More than one commit for a single sync + must land on `main` as ONE commit, via a scratch branch merged in + and cleaned up afterward -- not as several commits applied directly + to `main`, and not with a stray coco_scratch_* branch left behind. + Warms the store up with 2 Source rows (under the cap) first, so the + commit-count delta measured across the second run isolates the + chunked entity phase from the type-creation commits, which are a + separate, unconditional code path covered elsewhere. A sibling + Claim row, declared and left unchanged in both runs, proves the + scratch-branch path doesn't disturb it.""" + monkeypatch.setattr(ogt, "_MAX_ENTITIES_PER_TYPE", 4) + db = _e2e_db(store, "oversized") + row_count = {"n": 2} + + async def main() -> None: + source_schema = await NodeSchema.from_class(_ScSourceNarrow, key="slug") + claim_schema = await NodeSchema.from_class(_ScClaim, key="slug") + + sources = await omnigraph.mount_node_target(db, "Source", source_schema) + claims = await omnigraph.mount_node_target(db, "Claim", claim_schema) + for i in range(row_count["n"]): + sources.declare_node(node=_ScSourceNarrow(slug=f"s{i}", title=f"S{i}")) + claims.declare_node(node=_ScClaim(slug="c1")) + + app = coco.App(coco.AppConfig(name="e2e_oversized", environment=coco_env), main) + app.update_blocking() # warm-up: creates both types, 2 Source rows, 1 Claim row + + before = _commit_count(store) + row_count["n"] = 10 # 8 new rows -> 2 upsert chunks of 4 at cap=4 + app.update_blocking() + after = _commit_count(store) + + assert after - before == 1 # one merge commit, not two direct commits + assert _branch_names(store) == ["main"] # no coco_scratch_* branch left behind + + rows = _export_rows(store) + assert len([r for r in rows if r.get("type") == "Source"]) == 10 + assert len([r for r in rows if r.get("type") == "Claim"]) == 1 + + def test_user_managed_first_run_writes_rows_without_touching_the_schema( + self, store: str + ) -> None: + """`managed_by=USER` on a graph this app has never tracked. The type + already exists (created here the way a user's own `omnigraph init` + would), and the app's job is only to keep rows in sync. + + On run one there are no tracking records, so the tracked diff is + empty β€” which used to raise, making the mode unusable at exactly the + moment it's meant to be used. The docs advertise it as working. + """ + store_dir = Path(store[len("file://") :]) + assert ( + _init_live( + store_dir, + ( + "node Source {\n slug: String @key\n title: String?\n" + " coco_key: String\n}\n" + ), + ).returncode + == 0 + ) + schema_before = _read_schema_source(store) + + db = _e2e_db(store, "user_managed") + rows_to_write = {"n": 2} + + async def main() -> None: + sources = await omnigraph.mount_node_target( + db, + "Source", + await NodeSchema.from_class(_ScSourceNarrow, key="slug"), + managed_by=ManagedBy.USER, + ) + for i in range(rows_to_write["n"]): + sources.declare_node(node=_ScSourceNarrow(slug=f"s{i}", title=f"S{i}")) + + app = coco.App( + coco.AppConfig(name="e2e_user_managed", environment=coco_env), main + ) + app.update_blocking() + + rows = _export_rows(store) + assert sorted(r["data"]["slug"] for r in rows if r.get("type") == "Source") == [ + "s0", + "s1", + ] + # Never rewritten: byte-identical to what the user applied. + assert _read_schema_source(store) == schema_before + + # Run two: rows still reconcile normally, schema still untouched. + rows_to_write["n"] = 1 + app.update_blocking() + rows = _export_rows(store) + assert [r["data"]["slug"] for r in rows if r.get("type") == "Source"] == ["s0"] + assert _read_schema_source(store) == schema_before + + def test_undeclared_type_is_dropped_with_its_rows(self, store: str) -> None: + """Stop mounting a node type and it must disappear from the graph, + rows included. + + This is the only thing that removes them: the engine emits no + per-child deletes when a container target state goes away, so a + "drop" that wrote nothing left the type and every node in it behind + forever β€” still in `schema show`, still in `export`, and no longer + tracked by anything that could ever clean them up. + """ + db = _e2e_db(store, "type_dropped") + keep_extra = {"on": True} + + async def main() -> None: + claims = await omnigraph.mount_node_target( + db, "Claim", await NodeSchema.from_class(_ScClaim, key="slug") + ) + claims.declare_node(node=_ScClaim(slug="c1")) + if keep_extra["on"]: + sources = await omnigraph.mount_node_target( + db, + "Source", + await NodeSchema.from_class(_ScSourceNarrow, key="slug"), + ) + sources.declare_node(node=_ScSourceNarrow(slug="s1", title="S1")) + + app = coco.App( + coco.AppConfig(name="e2e_type_dropped", environment=coco_env), main + ) + app.update_blocking() + + rows = _export_rows(store) + assert len([r for r in rows if r.get("type") == "Source"]) == 1 + assert "node Source" in _read_schema_source(store) + + keep_extra["on"] = False + app.update_blocking() + + rows = _export_rows(store) + assert [r for r in rows if r.get("type") == "Source"] == [] + assert "node Source" not in _read_schema_source(store) + # The type that stayed declared is untouched by the other's drop. + assert len([r for r in rows if r.get("type") == "Claim"]) == 1 + assert "node Claim" in _read_schema_source(store) + + def test_undeclared_user_managed_type_survives_with_its_rows( + self, store: str + ) -> None: + """The `managed_by=USER` counterpart of the test above: stop mounting + the type and BOTH its schema block and its rows must survive. + + Same app-level edit as the drop test, opposite required outcome. The + connector did not create this type, so it does not get to delete it, + and deleting it takes the user's own rows with it -- unrecoverable, + and the exact opposite of what the docs promise. + """ + store_dir = Path(store[len("file://") :]) + assert ( + _init_live( + store_dir, + ( + "node Source {\n slug: String @key\n title: String?\n" + " coco_key: String\n}\n" + ), + ).returncode + == 0 + ) + + db = _e2e_db(store, "user_managed_undeclared") + keep_extra = {"on": True} + + async def main() -> None: + claims = await omnigraph.mount_node_target( + db, "Claim", await NodeSchema.from_class(_ScClaim, key="slug") + ) + claims.declare_node(node=_ScClaim(slug="c1")) + if keep_extra["on"]: + sources = await omnigraph.mount_node_target( + db, + "Source", + await NodeSchema.from_class(_ScSourceNarrow, key="slug"), + managed_by=ManagedBy.USER, + ) + sources.declare_node(node=_ScSourceNarrow(slug="s1", title="S1")) + + app = coco.App( + coco.AppConfig(name="e2e_user_managed_undeclared", environment=coco_env), + main, + ) + app.update_blocking() + + rows = _export_rows(store) + assert [r["data"]["slug"] for r in rows if r.get("type") == "Source"] == ["s1"] + assert "node Source" in _read_schema_source(store) + + keep_extra["on"] = False + app.update_blocking() + + # The user owns this type: the block stays, and so does its row. + rows = _export_rows(store) + assert [r["data"]["slug"] for r in rows if r.get("type") == "Source"] == ["s1"] + assert "node Source" in _read_schema_source(store) + + def test_integer_keyed_endpoint(self, store: str) -> None: + """An edge whose endpoint node is keyed on an int. + + The endpoint reference is the node's `id`, which is a String + rendering of the key value β€” never the key's own type. Inferring + `I64` from the Python value instead made EVERY edge insert into an + int-keyed endpoint fail with "cannot assign/compare I64 with String + for property `to`", which is exactly what the shipped example does. + """ + db = _e2e_db(store, "int_endpoint") + + async def main() -> None: + people = await omnigraph.mount_node_target( + db, "Person", await NodeSchema.from_class(_ScClaim, key="slug") + ) + meetings = await omnigraph.mount_node_target( + db, "Meeting", await NodeSchema.from_class(_ScMeeting, key="meeting_id") + ) + attended = await omnigraph.mount_edge_target( + db, "Attended", people, meetings + ) + people.declare_node(node=_ScClaim(slug="Ada")) + meetings.declare_node(node=_ScMeeting(meeting_id=7, note="Kickoff")) + attended.declare_edge(from_id="Ada", to_id=7) + + app = coco.App( + coco.AppConfig(name="e2e_int_endpoint", environment=coco_env), main + ) + app.update_blocking() + + rows = _export_rows(store) + (edge,) = [r for r in rows if r.get("edge") == "Attended"] + assert (edge["from"], edge["to"]) == ("Ada", "7") + # Re-running must not duplicate it, and must not re-stub the endpoint + # over the real node (a keyed insert is a full-record replace). + app.update_blocking() + rows = _export_rows(store) + assert len([r for r in rows if r.get("edge") == "Attended"]) == 1 + assert [r["data"]["note"] for r in rows if r.get("type") == "Meeting"] == [ + "Kickoff" + ] + + def test_list_of_dates_round_trips(self, store: str) -> None: + """`list[datetime.date]` renders as `[Date]`, which the engine accepts + at `init` β€” and then every write of such a node died in `json.dumps` + because only the bare `Date`/`DateTime` scalars had an encoder.""" + db = _e2e_db(store, "list_of_dates") + + @dataclass + class _Holiday: + slug: str + days: list[datetime.date] + + async def main() -> None: + holidays = await omnigraph.mount_node_target( + db, "Holiday", await NodeSchema.from_class(_Holiday, key="slug") + ) + holidays.declare_node( + node=_Holiday( + slug="xmas", + days=[datetime.date(2026, 12, 25), datetime.date(2026, 12, 26)], + ) + ) + + app = coco.App( + coco.AppConfig(name="e2e_list_of_dates", environment=coco_env), main + ) + app.update_blocking() + + (row,) = [r for r in _export_rows(store) if r.get("type") == "Holiday"] + assert len(row["data"]["days"]) == 2 + + def test_encoder_change_rewrites_the_stored_value(self, store: str) -> None: + """Same raw record, different `PropertyDef.encoder`: the value the + graph holds must follow the encoder, which only happens if change + detection fingerprints the encoded value rather than the raw one.""" + db = _e2e_db(store, "encoder_change") + encoder = {"fn": str.lower} + + async def main() -> None: + schema = NodeSchema( + properties={ + "slug": PropertyDef("slug", "String"), + "name": PropertyDef("name", "String", encoder["fn"]), + }, + key=("slug",), + ) + people = await omnigraph.mount_node_target(db, "Person", schema) + people.declare_node(node={"slug": "ada", "name": "Ada Lovelace"}) + + app = coco.App( + coco.AppConfig(name="e2e_encoder_change", environment=coco_env), main + ) + app.update_blocking() + assert [r["data"]["name"] for r in _export_rows(store) if r.get("type")] == [ + "ada lovelace" + ] + + encoder["fn"] = str.upper + app.update_blocking() + assert [r["data"]["name"] for r in _export_rows(store) if r.get("type")] == [ + "ADA LOVELACE" + ] + + def test_encoder_change_reaches_a_memoized_component(self, store: str) -> None: + """The test above declares its node from an unmemoized main. With the + declaration inside a `@coco.fn(memo=True)` component, the second + update skipped that component outright: neither its memo key nor + the type's tracking record knew about the encoder, so no reconcile + ever saw the new encoding and the lowercase value stayed.""" + db = _e2e_db(store, "encoder_change_memo") + encoder = {"fn": str.lower} + + @coco.fn(memo=True) + async def declare_ada(people: omnigraph.NodeTarget[Any]) -> None: + people.declare_node(node={"slug": "ada", "name": "Ada Lovelace"}) + + async def main() -> None: + schema = NodeSchema( + properties={ + "slug": PropertyDef("slug", "String"), + "name": PropertyDef("name", "String", encoder["fn"]), + }, + key=("slug",), + ) + people = await omnigraph.mount_node_target(db, "Person", schema) + await coco.mount(declare_ada, people) + + app = coco.App( + coco.AppConfig(name="e2e_encoder_change_memo", environment=coco_env), main + ) + app.update_blocking() + assert [r["data"]["name"] for r in _export_rows(store) if r.get("type")] == [ + "ada lovelace" + ] + + encoder["fn"] = str.upper + app.update_blocking() + assert [r["data"]["name"] for r in _export_rows(store) if r.get("type")] == [ + "ADA LOVELACE" + ] + + def test_abandoned_scratch_branch_is_reaped_before_a_schema_change( + self, store: str + ) -> None: + """An interrupted update leaves its `coco_scratch_*` branch behind, + and Omnigraph refuses every later schema change on the store while + it exists. The next schema change must recover on its own: reap the + abandoned branch, apply, and leave only `main` behind.""" + db = _e2e_db(store, "abandoned_scratch") + wide = {"on": False} + + async def main() -> None: + schema = await NodeSchema.from_class( + _ScSourceWide if wide["on"] else _ScSourceNarrow, key="slug" + ) + sources = await omnigraph.mount_node_target(db, "Source", schema) + node: Any = ( + _ScSourceWide(slug="a", title="A", note=None) + if wide["on"] + else _ScSourceNarrow(slug="a", title="A") + ) + sources.declare_node(node=node) + + app = coco.App( + coco.AppConfig(name="e2e_abandoned_scratch", environment=coco_env), main + ) + app.update_blocking() + + # What a process killed mid-sync leaves behind. + subprocess.run( + [ + _OMNIGRAPH_BIN, + "branch", + "create", + "coco_scratch_deadbeef", + "--from", + "main", + "--store", + store, + "--json", + "--quiet", + ], + check=True, + capture_output=True, + ) + assert "coco_scratch_deadbeef" in _branch_names(store) + + wide["on"] = True + app.update_blocking() + + assert _branch_names(store) == ["main"] + assert "note: String?" in _read_schema_source(store) + + def test_user_managed_type_follows_an_external_migration(self, store: str) -> None: + """`managed_by="user"` means the schema is the user's: they migrate + it with `omnigraph schema apply`, then declare the wider dataclass. + The connector used to compare the new declaration against what it + had tracked and refuse β€” advising exactly that migration, which it + then rejected again on every run because it never looked at the + live schema. A user-managed type must never be validated against + tracking history; it must simply write rows.""" + store_dir = Path(store[len("file://") :]) + v1 = "node Doc {\n slug: String @key\n coco_key: String\n}\n" + v2 = ( + "node Doc {\n slug: String @key\n title: String?\n coco_key: String\n}\n" + ) + assert _init_live(store_dir, v1).returncode == 0 + + @dataclass + class _DocV1: + slug: str + + @dataclass + class _DocV2: + slug: str + title: str | None + + db = _e2e_db(store, "user_managed_migration") + wide = {"on": False} + + async def main() -> None: + docs = await omnigraph.mount_node_target( + db, + "Doc", + await NodeSchema.from_class( + _DocV2 if wide["on"] else _DocV1, key="slug" + ), + managed_by=ManagedBy.USER, + ) + node: Any = ( + _DocV2(slug="d1", title="T") if wide["on"] else _DocV1(slug="d1") + ) + docs.declare_node(node=node) + + app = coco.App( + coco.AppConfig(name="e2e_user_managed_migration", environment=coco_env), + main, + ) + app.update_blocking() + + # The user's own migration, applied outside CocoIndex. + with tempfile.NamedTemporaryFile("w", suffix=".pg", encoding="utf-8") as f: + f.write(v2) + f.flush() + subprocess.run( + [ + _OMNIGRAPH_BIN, + "schema", + "apply", + "--schema", + f.name, + "--store", + store, + "--json", + "--quiet", + ], + check=True, + capture_output=True, + ) + + wide["on"] = True + app.update_blocking() + + (row,) = [r for r in _export_rows(store) if r.get("type") == "Doc"] + assert row["data"]["title"] == "T" + # The connector still never wrote the schema: the user's block is verbatim. + assert _read_schema_source(store) == v2 + + def test_hand_formatted_schema_survives_a_merge(self, store: str) -> None: + """`schema show` hands back the source verbatim, formatting and all, + so a schema somebody wrote by hand β€” indented, commented, two + declarations on one line β€” is what the merger has to edit. Each of + those used to produce a schema the engine refused.""" + store_dir = Path(store[len("file://") :]) + hand_written = ( + " node Source {\n" + " slug: String @key // the key } not a block end\n" + " title: String?\n" + " coco_key: String\n" + " }\n" + "node Other { slug: String @key coco_key: String } " + "node Extra { slug: String @key coco_key: String }\n" + ) + assert _init_live(store_dir, hand_written).returncode == 0 + db = _e2e_db(store, "hand_formatted") + + async def main() -> None: + sources = await omnigraph.mount_node_target( + db, "Source", await NodeSchema.from_class(_ScSourceWide, key="slug") + ) + sources.declare_node(node=_ScSourceWide(slug="a", title="A", note="n")) + + app = coco.App( + coco.AppConfig(name="e2e_hand_formatted", environment=coco_env), main + ) + app.update_blocking() + + source = _read_schema_source(store) + assert source.count("node Source") == 1 + assert "note: String?" in source + assert "node Other {" in source and "node Extra {" in source + (row,) = [r for r in _export_rows(store) if r.get("type") == "Source"] + assert row["data"]["note"] == "n" + + def test_dropping_a_connected_app_removes_every_type(self, store: str) -> None: + """Source, Claim and Supports are three processing components, torn + down concurrently by `drop()`. Whichever node type's removal ran + ahead of the edge's used to hit the endpoint guard and stay behind, + so the first drop failed and only a second one finished the job. + One drop must leave nothing.""" + db = _e2e_db(store, "drop_connected") + + async def main() -> None: + sources = await omnigraph.mount_node_target( + db, "Source", await NodeSchema.from_class(_ScSourceNarrow, key="slug") + ) + claims = await omnigraph.mount_node_target( + db, "Claim", await NodeSchema.from_class(_ScClaim, key="slug") + ) + supports = await omnigraph.mount_edge_target( + db, + "Supports", + sources, + claims, + await EdgeSchema.from_class(_ScEdgeProps), + ) + sources.declare_node(node=_ScSourceNarrow(slug="a", title="A")) + claims.declare_node(node=_ScClaim(slug="c1")) + supports.declare_edge( + from_id="a", to_id="c1", record=_ScEdgeProps(weight=1) + ) + + app = coco.App( + coco.AppConfig(name="e2e_drop_connected", environment=coco_env), main + ) + app.update_blocking() + assert "edge Supports" in _read_schema_source(store) + + app.drop_blocking() + + assert _read_schema_source(store).strip() == "" + assert _export_rows(store) == [] + + def test_handing_an_edge_type_to_the_user_protects_it_from_a_drop( + self, store: str + ) -> None: + """Nodes and an edge type created system-managed, then the edge type + declared `managed_by="user"`, then the app dropped. The drop used to + succeed and delete the user's edge type with its edges: its block + still carried the ownership marker, and the node drops read that as + current ownership. Now the handoff releases the marker, so the drop + refuses to remove the nodes the user's edge still points at and + leaves everything in place.""" + db = _e2e_db(store, "handoff_drop") + edge_owner = {"m": ManagedBy.SYSTEM} + + async def main() -> None: + sources = await omnigraph.mount_node_target( + db, "Source", await NodeSchema.from_class(_ScSourceNarrow, key="slug") + ) + claims = await omnigraph.mount_node_target( + db, "Claim", await NodeSchema.from_class(_ScClaim, key="slug") + ) + supports = await omnigraph.mount_edge_target( + db, + "Supports", + sources, + claims, + await EdgeSchema.from_class(_ScEdgeProps), + managed_by=edge_owner["m"], + ) + sources.declare_node(node=_ScSourceNarrow(slug="a", title="A")) + claims.declare_node(node=_ScClaim(slug="c1")) + supports.declare_edge( + from_id="a", to_id="c1", record=_ScEdgeProps(weight=1) + ) + + app = coco.App( + coco.AppConfig(name="e2e_handoff_drop", environment=coco_env), main + ) + app.update_blocking() + assert _read_schema_source(store).count("coco_managed") == 3 + + edge_owner["m"] = ManagedBy.USER + app.update_blocking() + source = _read_schema_source(store) + assert source.count("coco_managed") == 2 + assert ( + "edge Supports: Source -> Claim {\n weight: I64\n coco_key: String\n}" + in source + ) + + with pytest.raises(ValueError, match=r"Supports.*not managed by this app"): + app.drop_blocking() + assert _read_schema_source(store) == source + rows = _export_rows(store) + assert len([r for r in rows if r.get("edge") == "Supports"]) == 1 + assert len([r for r in rows if r.get("type")]) == 2 + + def test_a_comment_in_a_user_owned_edge_does_not_authorize_its_deletion( + self, store: str + ) -> None: + """The user's edge type mentions `coco_managed` in a comment only. + Ownership detection read the raw block text, so an app drop took + the edge type along with the nodes and deleted the user's edges.""" + store_dir = Path(store[len("file://") :]) + owned = " coco_key: String\n coco_managed_by_e2e_comment_ownership: Bool?\n" + hand_written = ( + f"node Source {{\n slug: String @key\n title: String?\n{owned}}}\n\n" + f"node Claim {{\n slug: String @key\n{owned}}}\n\n" + "edge Supports: Source -> Claim {\n" + " // coco_managed_by_e2e_comment_ownership: Bool? was deliberately omitted\n" + " coco_key: String\n" + "}\n" + ) + assert _init_live(store_dir, hand_written).returncode == 0 + db = _e2e_db(store, "comment_ownership") + + async def main() -> None: + sources = await omnigraph.mount_node_target( + db, "Source", await NodeSchema.from_class(_ScSourceNarrow, key="slug") + ) + claims = await omnigraph.mount_node_target( + db, "Claim", await NodeSchema.from_class(_ScClaim, key="slug") + ) + supports = await omnigraph.mount_edge_target( + db, "Supports", sources, claims, managed_by=ManagedBy.USER + ) + sources.declare_node(node=_ScSourceNarrow(slug="a", title="A")) + claims.declare_node(node=_ScClaim(slug="c1")) + supports.declare_edge(from_id="a", to_id="c1") + + app = coco.App( + coco.AppConfig(name="e2e_comment_ownership", environment=coco_env), main + ) + app.update_blocking() + assert len([r for r in _export_rows(store) if r.get("edge") == "Supports"]) == 1 + + with pytest.raises(ValueError, match=r"Supports.*not managed by this app"): + app.drop_blocking() + assert "was deliberately omitted" in _read_schema_source(store) + assert len([r for r in _export_rows(store) if r.get("edge") == "Supports"]) == 1 + + def test_dropping_an_app_leaves_another_apps_edge_type_alone( + self, store: str + ) -> None: + """App A creates Source. App B adopts Source as user-managed and + creates Claim and Supports: Source -> Claim. Dropping A removed B's + Supports with every edge in it: the drop took every referencing + edge type carrying the ownership property along, and that property + only said some app of this connector wrote the block. It names the + owning app now, and a drop refuses to remove a node type another + app's edge type still points at, naming that app.""" + db_a = _e2e_db(store, "two_apps_a") + db_b = _e2e_db(store, "two_apps_b") + + async def main_a() -> None: + sources = await omnigraph.mount_node_target( + db_a, "Source", await NodeSchema.from_class(_ScSourceNarrow, key="slug") + ) + sources.declare_node(node=_ScSourceNarrow(slug="a", title="A")) + + async def main_b() -> None: + sources = await omnigraph.mount_node_target( + db_b, + "Source", + await NodeSchema.from_class(_ScSourceNarrow, key="slug"), + managed_by=ManagedBy.USER, + ) + claims = await omnigraph.mount_node_target( + db_b, "Claim", await NodeSchema.from_class(_ScClaim, key="slug") + ) + supports = await omnigraph.mount_edge_target( + db_b, + "Supports", + sources, + claims, + await EdgeSchema.from_class(_ScEdgeProps), + ) + claims.declare_node(node=_ScClaim(slug="c1")) + supports.declare_edge( + from_id="a", to_id="c1", record=_ScEdgeProps(weight=1) + ) + + app_a = coco.App( + coco.AppConfig(name="e2e_two_apps_a", environment=coco_env), main_a + ) + app_b = coco.App( + coco.AppConfig(name="e2e_two_apps_b", environment=coco_env), main_b + ) + app_a.update_blocking() + app_b.update_blocking() + source = _read_schema_source(store) + assert "coco_managed_by_e2e_two_apps_a" in source + assert "coco_managed_by_e2e_two_apps_b" in source + + def supports_edges() -> int: + return len([r for r in _export_rows(store) if r.get("edge") == "Supports"]) + + assert supports_edges() == 1 + with pytest.raises( + ValueError, match=r"Supports.*coco_managed_by_e2e_two_apps_b" + ): + app_a.drop_blocking() + assert _read_schema_source(store) == source + assert supports_edges() == 1 + app_b.update_blocking() + assert supports_edges() == 1 + + def test_reclaiming_an_edge_type_restores_its_ownership(self, store: str) -> None: + """system -> user -> system with the schema unchanged throughout. + The second handoff wrote nothing, so the edge block stayed without + `coco_managed`, and dropping the connected app then failed because + its own edge type looked user-owned.""" + db = _e2e_db(store, "reclaim") + edge_owner = {"m": ManagedBy.SYSTEM} + + async def main() -> None: + sources = await omnigraph.mount_node_target( + db, "Source", await NodeSchema.from_class(_ScSourceNarrow, key="slug") + ) + claims = await omnigraph.mount_node_target( + db, "Claim", await NodeSchema.from_class(_ScClaim, key="slug") + ) + supports = await omnigraph.mount_edge_target( + db, + "Supports", + sources, + claims, + await EdgeSchema.from_class(_ScEdgeProps), + managed_by=edge_owner["m"], + ) + sources.declare_node(node=_ScSourceNarrow(slug="a", title="A")) + claims.declare_node(node=_ScClaim(slug="c1")) + supports.declare_edge( + from_id="a", to_id="c1", record=_ScEdgeProps(weight=1) + ) + + app = coco.App(coco.AppConfig(name="e2e_reclaim", environment=coco_env), main) + app.update_blocking() + edge_owner["m"] = ManagedBy.USER + app.update_blocking() + assert _read_schema_source(store).count("coco_managed") == 2 + edge_owner["m"] = ManagedBy.SYSTEM + app.update_blocking() + assert _read_schema_source(store).count("coco_managed") == 3 + + app.drop_blocking() + assert _read_schema_source(store).strip() == "" + assert _export_rows(store) == [] + + def test_a_date_keyed_node_round_trips(self, store: str) -> None: + """Declared twice, a `Date`-keyed node must be one node: its identity + is the ISO form of the key, both in tracking and in the graph.""" + db = _e2e_db(store, "date_key") + + async def main() -> None: + await _declare_a_day(db) + + app = coco.App(coco.AppConfig(name="e2e_date_key", environment=coco_env), main) + app.update_blocking() + app.update_blocking() + + (row,) = [r for r in _export_rows(store) if r.get("type") == "Day"] + assert row["data"]["note"] == "n" + + def test_two_spellings_of_one_instant_are_one_node(self, store: str) -> None: + """The regression: `12:00+00:00` and `14:00+02:00` used to get two + tracking keys for one graph node. Declared together they are now the + same target state, which the engine refuses as a duplicate; declared + one after the other they are one node whose removal is one delete.""" + db = _e2e_db(store, "datetime_key") + + @dataclass + class _Event: + at: datetime.datetime + note: str | None + + utc = datetime.datetime(2026, 1, 1, 12, 0, tzinfo=datetime.UTC) + plus_two = datetime.datetime( + 2026, 1, 1, 14, 0, tzinfo=datetime.timezone(datetime.timedelta(hours=2)) + ) + declare = {"ats": [utc, plus_two]} + + async def main() -> None: + events = await omnigraph.mount_node_target( + db, "Event", await NodeSchema.from_class(_Event, key="at") + ) + for at in declare["ats"]: + events.declare_node(node=_Event(at=at, note=at.isoformat())) + + app = coco.App( + coco.AppConfig(name="e2e_datetime_key", environment=coco_env), main + ) + with pytest.raises(ValueError, match="already declared"): + app.update_blocking() + + declare["ats"] = [utc] + app.update_blocking() + declare["ats"] = [plus_two] # same instant, respelled: still that node + app.update_blocking() + (row,) = [r for r in _export_rows(store) if r.get("type") == "Event"] + assert row["data"]["at"] == 1767268800000 + assert row["data"]["note"] == plus_two.isoformat() + + declare["ats"] = [] + app.update_blocking() + assert [r for r in _export_rows(store) if r.get("type") == "Event"] == [] + + def test_an_explicit_datetime_keyed_schema_round_trips(self, store: str) -> None: + """Positive case for a hand-built schema: no dataclass, no encoder + given, `datetime` values in dict rows. Declared twice, one node.""" + db = _e2e_db(store, "explicit_datetime") + schema = NodeSchema( + properties={ + "at": PropertyDef("at", "DateTime"), + "note": PropertyDef("note", "String?"), + }, + key=("at",), + ) + at = datetime.datetime(2026, 1, 1, 12, 0, tzinfo=datetime.UTC) + + async def main() -> None: + events = await omnigraph.mount_node_target(db, "Event", schema) + events.declare_node(node={"at": at, "note": "n"}) + + app = coco.App( + coco.AppConfig(name="e2e_explicit_datetime", environment=coco_env), main + ) + app.update_blocking() + app.update_blocking() + (row,) = [r for r in _export_rows(store) if r.get("type") == "Event"] + assert row["data"]["at"] == 1767268800000 + assert row["data"]["note"] == "n" + + def test_an_explicit_date_keyed_schema_round_trips(self, store: str) -> None: + db = _e2e_db(store, "explicit_date") + schema = NodeSchema( + properties={ + "on": PropertyDef("on", "Date"), + "note": PropertyDef("note", "String?"), + }, + key=("on",), + ) + + async def main() -> None: + days = await omnigraph.mount_node_target(db, "Day", schema) + days.declare_node(node={"on": datetime.date(2026, 1, 5), "note": "n"}) + + app = coco.App( + coco.AppConfig(name="e2e_explicit_date", environment=coco_env), main + ) + app.update_blocking() + app.update_blocking() + (row,) = [r for r in _export_rows(store) if r.get("type") == "Day"] + assert row["data"]["note"] == "n" diff --git a/rust/py/src/context.rs b/rust/py/src/context.rs index 00dbc9254..334d54aac 100644 --- a/rust/py/src/context.rs +++ b/rust/py/src/context.rs @@ -28,6 +28,12 @@ impl PyComponentProcessorContext { PyStablePath(self.0.stable_path().clone()) } + /// Name of the app this component runs under (`AppConfig.name`). + #[getter] + fn app_name(&self) -> &str { + self.0.app_ctx().app_reg().name() + } + #[getter] fn live(&self) -> bool { self.0.live()