From 5aadce3c6327785188df143e0bf8ba70fb0f0e2b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=90=AA=E9=9C=84?= Date: Thu, 27 Aug 2026 17:47:49 +0800 Subject: [PATCH 1/7] feat(backend): collection import/export pipelines with snapshot mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Export via zvec 0.7 iter_docs: GET /collections/{name}/documents:export streams JSONL with constant memory (includeVector/outputFields/format; non-finite values fail fast with 422). - Batched import: POST /collections/{name}/documents:import with a per-row report (mode replace|insert, onError abort|skip). SDK write-batch limits are absorbed internally, so >1024-doc requests no longer fail with a 500. - Snapshot mode: mode=snapshot streams manifest.json + documents.jsonl as a .tar.gz; POST /collections:import rebuilds a whole collection from it — schema pre-checked before any write, rolled back if the load fails midway. - Document wire representation moved to storage/doc_repr.py: the primary key rides the reserved chain id -> $id -> $$id... whenever a column occupies "id", fixing silent key loss on such schemas. - Open collection names are globally unique: create/open/import answer 409 COLLECTION_ALREADY_EXISTS with the blocking path; close/destroy release it. - Error taxonomy: DOCUMENT_CONFLICT(409), INVALID_DOCUMENT(422), MAINTENANCE_BLOCKED(409) while an export is open, UNSUPPORTED_VECTOR_DATA_TYPE(422), precise 400s for corrupt sources. Signed-off-by: 琪霄 --- apps/backend/pyproject.toml | 5 +- .../integration/test_collections_router.py | 80 ++- .../tests/integration/test_export_pipeline.py | 359 ++++++++++ .../tests/integration/test_fs_router.py | 92 +++ .../integration/test_id_column_collections.py | 219 ++++++ .../integration/test_import_collection.py | 284 ++++++++ .../tests/integration/test_import_endpoint.py | 266 +++++++ .../tests/integration/test_import_pipeline.py | 302 ++++++++ .../tests/integration/test_snapshot_io.py | 230 ++++++ .../tests/integration/test_write_limits.py | 138 ++++ .../tests/unit/test_collections_schemas.py | 47 +- apps/backend/tests/unit/test_doc_repr.py | 269 +++++++ .../backend/tests/unit/test_export_formats.py | 97 +++ .../tests/unit/test_import_batch_abort.py | 116 +++ .../backend/tests/unit/test_import_formats.py | 98 +++ .../tests/unit/test_sdk_query_params.py | 18 + apps/backend/tests/unit/test_snapshot.py | 230 ++++++ apps/backend/tests/unit/test_write_batches.py | 152 ++++ apps/backend/uv.lock | 96 +-- apps/backend/zvec_studio/exceptions.py | 94 +++ .../zvec_studio/routers/collections.py | 83 +++ apps/backend/zvec_studio/routers/documents.py | 194 ++++- apps/backend/zvec_studio/routers/fs.py | 105 ++- apps/backend/zvec_studio/schemas/__init__.py | 12 + .../backend/zvec_studio/schemas/collection.py | 83 ++- apps/backend/zvec_studio/schemas/document.py | 81 ++- apps/backend/zvec_studio/storage/doc_repr.py | 176 +++++ apps/backend/zvec_studio/storage/formats.py | 212 ++++++ apps/backend/zvec_studio/storage/import_.py | 104 +++ apps/backend/zvec_studio/storage/sdk.py | 676 +++++++++++++++--- apps/backend/zvec_studio/storage/snapshot.py | 325 +++++++++ 31 files changed, 5030 insertions(+), 213 deletions(-) create mode 100644 apps/backend/tests/integration/test_export_pipeline.py create mode 100644 apps/backend/tests/integration/test_id_column_collections.py create mode 100644 apps/backend/tests/integration/test_import_collection.py create mode 100644 apps/backend/tests/integration/test_import_endpoint.py create mode 100644 apps/backend/tests/integration/test_import_pipeline.py create mode 100644 apps/backend/tests/integration/test_snapshot_io.py create mode 100644 apps/backend/tests/integration/test_write_limits.py create mode 100644 apps/backend/tests/unit/test_doc_repr.py create mode 100644 apps/backend/tests/unit/test_export_formats.py create mode 100644 apps/backend/tests/unit/test_import_batch_abort.py create mode 100644 apps/backend/tests/unit/test_import_formats.py create mode 100644 apps/backend/tests/unit/test_snapshot.py create mode 100644 apps/backend/tests/unit/test_write_batches.py create mode 100644 apps/backend/zvec_studio/storage/doc_repr.py create mode 100644 apps/backend/zvec_studio/storage/formats.py create mode 100644 apps/backend/zvec_studio/storage/import_.py create mode 100644 apps/backend/zvec_studio/storage/snapshot.py diff --git a/apps/backend/pyproject.toml b/apps/backend/pyproject.toml index 570d6d2..409d70b 100644 --- a/apps/backend/pyproject.toml +++ b/apps/backend/pyproject.toml @@ -31,7 +31,10 @@ dependencies = [ "orjson>=3.10,<4.0", "typing-extensions>=4.12", # Zvec embedded vector database SDK (core runtime dependency). - "zvec>=0.6,<0.7", + # 0.7 is the floor: ``Collection.iter_docs()`` (streaming full-collection + # traversal) only exists from 0.7 onwards and the export feature is built + # on it. + "zvec>=0.7,<0.8", # Required by zvec for numpy-backed vector payloads. "numpy>=1.23,<2.5", ] diff --git a/apps/backend/tests/integration/test_collections_router.py b/apps/backend/tests/integration/test_collections_router.py index e486165..4de2db0 100644 --- a/apps/backend/tests/integration/test_collections_router.py +++ b/apps/backend/tests/integration/test_collections_router.py @@ -46,6 +46,67 @@ async def _create(client: AsyncClient, path: Path, name: str) -> dict: return resp.json() +class TestNameUniqueness: + """Open collection names are globally unique. + + Nearly every document API resolves a collection by name alone, so two + same-named collections can never be open at once — otherwise requests + would silently resolve to whichever registered first (observed in the + wild as "deleted in A, B's browse changed too"). + """ + + async def test_create_same_name_different_path_is_409( + self, client: AsyncClient, tmp_path: Path + ) -> None: + await _create(client, tmp_path / "first", "demo") + + resp = await client.post( + f"{API}/collections", + json={"path": str(tmp_path / "second"), "schema": _payload("demo")}, + ) + + assert resp.status_code == 409, resp.text + body = resp.json() + assert body["code"] == "COLLECTION_ALREADY_EXISTS" + # The message points at the blocking collection and says what to do. + assert "demo" in body["detail"] + assert str(tmp_path / "first") in body["detail"] + + async def test_open_same_name_different_path_is_409( + self, client: AsyncClient, tmp_path: Path + ) -> None: + await _create(client, tmp_path / "first", "demo") + # A second on-disk collection with the same recorded name. + import zvec + + other = tmp_path / "second" + zvec.create_and_open( + path=str(other), + schema=zvec.CollectionSchema( + name="demo", + vectors=[zvec.VectorSchema("embedding", zvec.DataType.VECTOR_FP32, 4)], + fields=[zvec.FieldSchema("title", zvec.DataType.STRING)], + ), + ).flush() + + resp = await client.post(f"{API}/collections/open", json={"path": str(other)}) + + assert resp.status_code == 409, resp.text + assert resp.json()["code"] == "COLLECTION_ALREADY_EXISTS" + + async def test_destroy_frees_the_name( + self, client: AsyncClient, tmp_path: Path + ) -> None: + await _create(client, tmp_path / "first", "demo") + await client.delete(f"{API}/collections/demo?path={tmp_path / 'first'}") + + second = await client.post( + f"{API}/collections", + json={"path": str(tmp_path / "second"), "schema": _payload("demo")}, + ) + assert second.status_code == 201, second.text + + class TestCreate: async def test_create_returns_201_with_summary( self, client: AsyncClient, tmp_path: Path @@ -60,25 +121,6 @@ async def test_create_returns_201_with_summary( assert body["schema"]["vectors"][0]["indexParam"]["metric"] == "COSINE" assert body["stats"]["documentCount"] == 0 - async def test_duplicate_name_different_path_both_open( - self, client: AsyncClient, tmp_path: Path - ) -> None: - await _create(client, tmp_path / "a", "same") - resp = await client.post( - f"{API}/collections", - json={"path": str(tmp_path / "b"), "schema": _payload("same")}, - ) - assert resp.status_code == 201 - body = resp.json() - assert body["name"] == "same" - assert body["path"] == str(tmp_path / "b") - # Both should be listed as open - list_resp = await client.get(f"{API}/collections") - items = list_resp.json()["items"] - paths = {i["path"] for i in items} - assert str(tmp_path / "a") in paths - assert str(tmp_path / "b") in paths - async def test_duplicate_path_returns_409( self, client: AsyncClient, tmp_path: Path ) -> None: diff --git a/apps/backend/tests/integration/test_export_pipeline.py b/apps/backend/tests/integration/test_export_pipeline.py new file mode 100644 index 0000000..c5333ca --- /dev/null +++ b/apps/backend/tests/integration/test_export_pipeline.py @@ -0,0 +1,359 @@ +"""Integration tests for the export pipeline. + +Two layers under test: + +* ``SdkBackend.iter_documents`` — the read path. The critical contract from + PR #597: while an iterator is open, maintenance ops (create_index, + optimize, ...) are rejected. Studio therefore MUST release the iterator on + every exit path — normal completion, early break (client disconnect), and + exceptions — or the user's collection stays "locked". +* ``GET /collections/{name}/documents:export`` — the HTTP surface: streaming + JSONL, headers, format resolution, and the export->import roundtrip. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +from httpx import AsyncClient + +zvec = pytest.importorskip("zvec") + +from zvec_studio.schemas.collection import ( # noqa: E402 + CollectionSchema, + IndexType, + MetricType, + VectorIndexParam, +) +from zvec_studio.storage.sdk import SdkBackend # noqa: E402 + +pytestmark = pytest.mark.integration + +API = "/api/v1" + + +def _schema(name: str = "export_source", *, with_id_field: bool = False) -> CollectionSchema: + fields = [{"name": "title", "dataType": "STRING"}] + if with_id_field: + fields.insert(0, {"name": "id", "dataType": "STRING"}) + return CollectionSchema.model_validate( + { + "name": name, + "vectors": [ + { + "name": "embedding", + "dataType": "VECTOR_FP32", + "dimension": 4, + "indexParam": {"indexType": "FLAT", "metric": "L2"}, + } + ], + "fields": fields, + } + ) + + +def _make_collection( + tmp_path: Path, *, count: int = 5, name: str = "export_source", with_id_field: bool = False +) -> SdkBackend: + backend = SdkBackend() + backend.create(path=tmp_path / name, schema=_schema(name, with_id_field=with_id_field)) + docs = [] + for i in range(count): + doc = { + "id": f"doc-{i:03d}", + "title": f"t{i}", + "embedding": [float(i), 0.0, 0.0, 0.0], + } + if with_id_field: + doc["id"] = f"USER-{i}" + doc["$id"] = f"doc-{i:03d}" + docs.append(doc) + backend.insert_documents(name, docs) + return backend + + +class TestIterDocuments: + def test_streams_all_rows_with_fields(self, tmp_path: Path) -> None: + backend = _make_collection(tmp_path, count=5) + + rows = list(backend.iter_documents("export_source", include_vector=False)) + + assert len(rows) == 5 + assert rows[0] == {"id": "doc-000", "title": "t0"} + assert "embedding" not in rows[0] + + def test_include_vector_adds_vectors(self, tmp_path: Path) -> None: + backend = _make_collection(tmp_path, count=2) + + rows = list(backend.iter_documents("export_source", include_vector=True)) + + assert rows[1]["embedding"] == [1.0, 0.0, 0.0, 0.0] + + def test_output_fields_filters_columns(self, tmp_path: Path) -> None: + backend = _make_collection(tmp_path, count=2) + + rows = list( + backend.iter_documents( + "export_source", include_vector=False, output_fields=["title"] + ) + ) + + assert rows[0] == {"id": "doc-000", "title": "t0"} + + def test_id_column_schema_uses_reserved_pk_key(self, tmp_path: Path) -> None: + backend = _make_collection(tmp_path, count=1, with_id_field=True) + + rows = list(backend.iter_documents("export_source", include_vector=False)) + + assert rows[0] == {"$id": "doc-000", "id": "USER-0", "title": "t0"} + + +class TestIteratorLifecycle: + def test_maintenance_allowed_after_full_iteration(self, tmp_path: Path) -> None: + backend = _make_collection(tmp_path, count=3) + + list(backend.iter_documents("export_source", include_vector=False)) + # Would raise if the iterator had leaked. + backend.create_index( + "export_source", + vector_field="embedding", + index_param=VectorIndexParam(indexType=IndexType.FLAT, metric=MetricType.L2), + ) + + def test_maintenance_allowed_after_early_close(self, tmp_path: Path) -> None: + """Simulates a client disconnect: the generator is closed mid-stream.""" + backend = _make_collection(tmp_path, count=10) + + gen = backend.iter_documents("export_source", include_vector=False) + next(gen) # read a single row, then abandon the iterator + gen.close() + + backend.create_index( + "export_source", + vector_field="embedding", + index_param=VectorIndexParam(indexType=IndexType.FLAT, metric=MetricType.L2), + ) + + def test_maintenance_allowed_after_iteration_error(self, tmp_path: Path) -> None: + """An exception while consuming must still release the iterator.""" + backend = _make_collection(tmp_path, count=5) + + with pytest.raises(RuntimeError, match="test abort"): + for _ in backend.iter_documents("export_source", include_vector=False): + raise RuntimeError("test abort") + + backend.create_index( + "export_source", + vector_field="embedding", + index_param=VectorIndexParam(indexType=IndexType.FLAT, metric=MetricType.L2), + ) + + def test_maintenance_blocked_while_iterator_open_is_409(self, tmp_path: Path) -> None: + """While an export iterator is open, maintenance is a transient + conflict (409 MAINTENANCE_BLOCKED), not an opaque 5xx / 400.""" + from zvec_studio.exceptions import MaintenanceBlockedError + + backend = _make_collection(tmp_path, count=3) + gen = backend.iter_documents("export_source", include_vector=False) + next(gen) + try: + with pytest.raises(MaintenanceBlockedError) as exc: + backend.optimize("export_source") + assert exc.value.status_code == 409 + with pytest.raises(MaintenanceBlockedError): + backend.create_index( + "export_source", + vector_field="embedding", + index_param=VectorIndexParam( + indexType=IndexType.FLAT, metric=MetricType.L2 + ), + ) + finally: + gen.close() + + # Released: maintenance succeeds again. + backend.optimize("export_source") + + +class TestMaintenanceBlockedHttp: + async def test_optimize_during_export_is_409( + self, client: AsyncClient, tmp_path: Path + ) -> None: + await _seed(client, tmp_path, "demo", 3) + app = client._transport.app # type: ignore[attr-defined] + gen = app.state.backend.iter_documents("demo", include_vector=False) + next(gen) + try: + resp = await client.post(f"{API}/collections/demo:optimize") + assert resp.status_code == 409, resp.text + assert resp.json()["code"] == "MAINTENANCE_BLOCKED" + finally: + gen.close() + + resp = await client.post(f"{API}/collections/demo:optimize") + assert resp.status_code == 200, resp.text + + +# --------------------------------------------------------------------------- +# HTTP surface +# --------------------------------------------------------------------------- + + +def _collection_payload(name: str, dim: int = 4) -> dict: + return { + "name": name, + "vectors": [ + { + "name": "embedding", + "dataType": "VECTOR_FP32", + "dimension": dim, + "indexParam": {"indexType": "FLAT", "metric": "L2"}, + } + ], + "fields": [ + {"name": "title", "dataType": "STRING"}, + {"name": "score", "dataType": "INT64"}, + ], + } + + +async def _seed(client: AsyncClient, tmp_path: Path, name: str, count: int) -> None: + resp = await client.post( + f"{API}/collections", + json={"path": str(tmp_path / name), "schema": _collection_payload(name)}, + ) + assert resp.status_code == 201, resp.text + docs = [ + { + "id": f"doc-{i:03d}", + "title": f"t{i}", + "score": i, + "embedding": [float(i), 0.0, 0.0, 0.0], + } + for i in range(count) + ] + if not docs: + return + resp = await client.post(f"{API}/collections/{name}/documents", json={"documents": docs}) + assert resp.status_code == 201, resp.text + + +class TestExportEndpoint: + async def test_exports_jsonl_with_headers(self, client: AsyncClient, tmp_path: Path) -> None: + await _seed(client, tmp_path, "demo", 3) + + resp = await client.get(f"{API}/collections/demo/documents:export") + + assert resp.status_code == 200, resp.text + assert resp.headers["content-type"].startswith("application/x-ndjson") + disposition = resp.headers["content-disposition"] + assert "attachment" in disposition + assert ".jsonl" in disposition + + rows = [json.loads(line) for line in resp.text.splitlines() if line] + assert len(rows) == 3 + assert rows[0]["id"] == "doc-000" + # Default export includes vectors. + assert rows[0]["embedding"] == [0.0, 0.0, 0.0, 0.0] + + async def test_include_vector_false_omits_vectors( + self, client: AsyncClient, tmp_path: Path + ) -> None: + await _seed(client, tmp_path, "demo", 1) + + resp = await client.get( + f"{API}/collections/demo/documents:export", + params={"includeVector": "false"}, + ) + + row = json.loads(resp.text.strip()) + assert "embedding" not in row + + async def test_output_fields_filters_columns( + self, client: AsyncClient, tmp_path: Path + ) -> None: + await _seed(client, tmp_path, "demo", 1) + + resp = await client.get( + f"{API}/collections/demo/documents:export", + params={"outputFields": "title", "includeVector": "false"}, + ) + + row = json.loads(resp.text.strip()) + assert row == {"id": "doc-000", "title": "t0"} + + async def test_empty_collection_yields_empty_body( + self, client: AsyncClient, tmp_path: Path + ) -> None: + await _seed(client, tmp_path, "demo", 0) + + resp = await client.get(f"{API}/collections/demo/documents:export") + + assert resp.status_code == 200 + assert resp.text == "" + + async def test_unknown_format_is_400(self, client: AsyncClient, tmp_path: Path) -> None: + await _seed(client, tmp_path, "demo", 1) + + resp = await client.get( + f"{API}/collections/demo/documents:export", params={"format": "csv"} + ) + + assert resp.status_code == 400, resp.text + assert resp.json()["code"] == "EXPORT_UNSUPPORTED_FORMAT" + + async def test_unknown_collection_is_404(self, client: AsyncClient) -> None: + resp = await client.get(f"{API}/collections/ghost/documents:export") + + assert resp.status_code == 404 + assert resp.json()["code"] == "COLLECTION_NOT_FOUND" + + +class TestExportImportRoundtrip: + async def test_export_then_import_is_lossless( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """The core promise: export -> file -> import reproduces every row.""" + await _seed(client, tmp_path, "source", 5) + + exported = await client.get( + f"{API}/collections/source/documents:export", + params={"includeVector": "true"}, + ) + assert exported.status_code == 200 + + # Write the exported bytes to a file and import into a fresh collection. + file_path = tmp_path / "roundtrip.jsonl" + file_path.write_bytes(exported.content) + resp = await client.post( + f"{API}/collections", + json={"path": str(tmp_path / "target"), "schema": _collection_payload("target")}, + ) + assert resp.status_code == 201, resp.text + imported = await client.post( + f"{API}/collections/target/documents:import", + json={"source": {"kind": "localPath", "path": str(file_path)}}, + ) + assert imported.status_code == 200, imported.text + assert imported.json()["imported"] == 5 + assert imported.json()["failed"] == 0 + + # Compare row by row through the browse API. + source_rows = ( + await client.post( + f"{API}/collections/source/documents:browse", + json={"limit": 100, "includeVector": True}, + ) + ).json()["items"] + target_rows = ( + await client.post( + f"{API}/collections/target/documents:browse", + json={"limit": 100, "includeVector": True}, + ) + ).json()["items"] + + assert sorted(source_rows, key=lambda r: r["id"]) == sorted( + target_rows, key=lambda r: r["id"] + ) diff --git a/apps/backend/tests/integration/test_fs_router.py b/apps/backend/tests/integration/test_fs_router.py index aadf456..d356703 100644 --- a/apps/backend/tests/integration/test_fs_router.py +++ b/apps/backend/tests/integration/test_fs_router.py @@ -104,3 +104,95 @@ async def test_root_has_no_parent(client: AsyncClient) -> None: body = response.json() assert body["path"] == "/" assert body["parent"] is None + + +# ---- file listing (import file picker support) ---- + + +def _picker_sandbox(tmp_path: Path) -> Path: + sandbox = tmp_path / "picker" + sandbox.mkdir() + (sandbox / "subdir").mkdir() + (sandbox / "data.jsonl").write_text('{"id": "a"}\n') + (sandbox / "archive.tar.gz").write_bytes(b"\x1f\x8b\x08\x00") + (sandbox / "notes.txt").write_text("hello") + return sandbox + + +@pytest.mark.integration +async def test_files_hidden_without_include_files(tmp_path: Path, client: AsyncClient) -> None: + """Backward compatibility: the default listing keeps directories only.""" + sandbox = _picker_sandbox(tmp_path) + + response = await client.get("/api/v1/fs/list", params={"path": str(sandbox)}) + + body = response.json() + assert [e["name"] for e in body["entries"]] == ["subdir"] + assert [e["kind"] for e in body["entries"]] == ["dir"] + assert body["entries"][0]["size"] is None + + +@pytest.mark.integration +async def test_files_listed_with_include_files(tmp_path: Path, client: AsyncClient) -> None: + sandbox = _picker_sandbox(tmp_path) + + response = await client.get( + "/api/v1/fs/list", + params={"path": str(sandbox), "includeFiles": "true"}, + ) + + assert response.status_code == 200 + entries = response.json()["entries"] + by_name = {e["name"]: e for e in entries} + assert sorted(by_name) == ["archive.tar.gz", "data.jsonl", "notes.txt", "subdir"] + assert by_name["data.jsonl"]["kind"] == "file" + assert by_name["data.jsonl"]["size"] == len('{"id": "a"}\n') + assert by_name["subdir"]["kind"] == "dir" + assert by_name["subdir"]["size"] is None + + +@pytest.mark.integration +async def test_files_filtered_by_extension(tmp_path: Path, client: AsyncClient) -> None: + sandbox = _picker_sandbox(tmp_path) + + response = await client.get( + "/api/v1/fs/list", + params={"path": str(sandbox), "includeFiles": "true", "extensions": ".jsonl,.tar.gz"}, + ) + + entries = response.json()["entries"] + names = [e["name"] for e in entries] + # Directories always stay (needed for navigation); files are filtered. + assert names == ["archive.tar.gz", "data.jsonl", "subdir"] + + +@pytest.mark.integration +async def test_extension_filter_is_case_insensitive( + tmp_path: Path, client: AsyncClient +) -> None: + sandbox = tmp_path / "picker" + sandbox.mkdir() + (sandbox / "UPPER.JSONL").write_text("{}\n") + + response = await client.get( + "/api/v1/fs/list", + params={"path": str(sandbox), "includeFiles": "true", "extensions": ".jsonl"}, + ) + + assert [e["name"] for e in response.json()["entries"]] == ["UPPER.JSONL"] + + +@pytest.mark.integration +async def test_extensions_without_include_files_ignored( + tmp_path: Path, client: AsyncClient +) -> None: + sandbox = _picker_sandbox(tmp_path) + + response = await client.get( + "/api/v1/fs/list", + params={"path": str(sandbox), "extensions": ".jsonl"}, + ) + + # includeFiles is off -> extensions is a no-op; directories-only as before. + assert [e["name"] for e in response.json()["entries"]] == ["subdir"] + diff --git a/apps/backend/tests/integration/test_id_column_collections.py b/apps/backend/tests/integration/test_id_column_collections.py new file mode 100644 index 0000000..3385106 --- /dev/null +++ b/apps/backend/tests/integration/test_id_column_collections.py @@ -0,0 +1,219 @@ +"""Integration tests for collections whose schema declares an ``id`` column. + +Zvec allows a scalar field (or a vector) named ``id`` to coexist with the +built-in ``Doc.id`` primary key. Studio used to reject that name, which meant: + +* ``POST /collections/open`` returned 500 for any SDK-created collection with + such a column -- it could not be managed through Studio at all; +* the flat row representation silently dropped the real primary key. + +Both are fixed by moving the primary key to the ``$id`` row key whenever a +column takes ``id`` (see ``storage/doc_repr.py``). +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import zvec +from httpx import AsyncClient + +API = "/api/v1" +pytestmark = pytest.mark.integration + + +def _create_via_sdk(path: Path, *, id_is_vector: bool = False) -> None: + """Create a collection with an ``id`` column using the raw SDK.""" + if id_is_vector: + vectors = [zvec.VectorSchema("id", zvec.DataType.VECTOR_FP32, 4)] + fields = [zvec.FieldSchema("title", zvec.DataType.STRING)] + else: + vectors = [zvec.VectorSchema("embedding", zvec.DataType.VECTOR_FP32, 4)] + fields = [ + zvec.FieldSchema("id", zvec.DataType.STRING), + zvec.FieldSchema("title", zvec.DataType.STRING), + ] + schema = zvec.CollectionSchema(name=path.name, vectors=vectors, fields=fields) + collection = zvec.create_and_open(path=str(path), schema=schema) + if id_is_vector: + collection.insert( + [zvec.Doc(id="PK-001", vectors={"id": [0.1, 0.2, 0.3, 0.4]}, fields={"title": "t"})] + ) + else: + collection.insert( + [ + zvec.Doc( + id="PK-001", + vectors={"embedding": [0.1, 0.2, 0.3, 0.4]}, + fields={"id": "USER-999", "title": "t"}, + ) + ] + ) + collection.flush() + + +class TestOpenExternalCollection: + async def test_open_collection_with_id_field( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """Regression: this used to fail with HTTP 500.""" + path = tmp_path / "legacy_field" + _create_via_sdk(path) + + resp = await client.post(f"{API}/collections/open", json={"path": str(path)}) + + assert resp.status_code == 200, resp.text + names = [f["name"] for f in resp.json()["schema"]["fields"]] + assert "id" in names + + async def test_open_collection_with_id_vector( + self, client: AsyncClient, tmp_path: Path + ) -> None: + path = tmp_path / "legacy_vector" + _create_via_sdk(path, id_is_vector=True) + + resp = await client.post(f"{API}/collections/open", json={"path": str(path)}) + + assert resp.status_code == 200, resp.text + assert [v["name"] for v in resp.json()["schema"]["vectors"]] == ["id"] + + +class TestPrimaryKeyIsPreserved: + async def test_get_document_keeps_both_values( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """The primary key lands on ``$id``; the ``id`` field keeps its own value.""" + path = tmp_path / "legacy_field" + _create_via_sdk(path) + await client.post(f"{API}/collections/open", json={"path": str(path)}) + + got = await client.get(f"{API}/collections/{path.name}/documents/PK-001") + + assert got.status_code == 200, got.text + body = got.json() + assert body["$id"] == "PK-001" + assert body["id"] == "USER-999" + assert body["title"] == "t" + + async def test_browse_keeps_both_values( + self, client: AsyncClient, tmp_path: Path + ) -> None: + path = tmp_path / "legacy_field" + _create_via_sdk(path) + await client.post(f"{API}/collections/open", json={"path": str(path)}) + + resp = await client.post( + f"{API}/collections/{path.name}/documents:browse", json={"limit": 10} + ) + + assert resp.status_code == 200, resp.text + rows = resp.json()["items"] + assert len(rows) == 1 + assert rows[0]["$id"] == "PK-001" + assert rows[0]["id"] == "USER-999" + + async def test_search_keeps_both_values( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """The search read path goes through the same row representation.""" + path = tmp_path / "legacy_field" + _create_via_sdk(path) + await client.post(f"{API}/collections/open", json={"path": str(path)}) + + resp = await client.post( + f"{API}/collections/{path.name}/searches", + json={ + "vector": [0.1, 0.2, 0.3, 0.4], + "vectorField": "embedding", + "topK": 5, + }, + ) + + assert resp.status_code == 200, resp.text + results = resp.json()["results"] + assert len(results) == 1 + assert results[0]["id"] == "PK-001" # dedicated hit pk + assert results[0]["fields"]["$id"] == "PK-001" + assert results[0]["fields"]["id"] == "USER-999" + + +class TestWriteWithIdColumn: + async def test_insert_requires_reserved_key( + self, client: AsyncClient, tmp_path: Path + ) -> None: + path = tmp_path / "legacy_field" + _create_via_sdk(path) + await client.post(f"{API}/collections/open", json={"path": str(path)}) + + resp = await client.post( + f"{API}/collections/{path.name}/documents", + json={ + "documents": [ + { + "$id": "PK-002", + "id": "USER-2", + "title": "t2", + "embedding": [0.5, 0.5, 0.5, 0.5], + } + ] + }, + ) + assert resp.status_code == 201, resp.text + assert resp.json()["ids"] == ["PK-002"] + + got = await client.get(f"{API}/collections/{path.name}/documents/PK-002") + assert got.json()["$id"] == "PK-002" + assert got.json()["id"] == "USER-2" + + async def test_ambiguous_plain_id_is_rejected( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """A bare ``id`` cannot be resolved when the schema declares one.""" + path = tmp_path / "legacy_field" + _create_via_sdk(path) + await client.post(f"{API}/collections/open", json={"path": str(path)}) + + resp = await client.post( + f"{API}/collections/{path.name}/documents", + json={"documents": [{"id": "???", "title": "t", "embedding": [0.1, 0.2, 0.3, 0.4]}]}, + ) + + assert resp.status_code == 400, resp.text + assert "$id" in resp.json()["detail"] + + +class TestOrdinarySchemaUnchanged: + async def test_plain_collection_still_uses_id( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """Collections without an ``id`` column keep the pre-existing contract.""" + resp = await client.post( + f"{API}/collections", + json={ + "path": str(tmp_path / "plain"), + "schema": { + "name": "plain", + "vectors": [ + { + "name": "embedding", + "dataType": "VECTOR_FP32", + "dimension": 4, + "indexParam": {"indexType": "FLAT", "metric": "L2"}, + } + ], + "fields": [{"name": "title", "dataType": "STRING"}], + }, + }, + ) + assert resp.status_code == 201, resp.text + + await client.post( + f"{API}/collections/plain/documents", + json={"documents": [{"id": "doc-1", "title": "t", "embedding": [0.1, 0.2, 0.3, 0.4]}]}, + ) + got = await client.get(f"{API}/collections/plain/documents/doc-1") + + assert got.status_code == 200 + assert got.json()["id"] == "doc-1" + assert "$id" not in got.json() diff --git a/apps/backend/tests/integration/test_import_collection.py b/apps/backend/tests/integration/test_import_collection.py new file mode 100644 index 0000000..e246d2a --- /dev/null +++ b/apps/backend/tests/integration/test_import_collection.py @@ -0,0 +1,284 @@ +"""Integration tests for ``POST /collections:import``. + +Snapshot restore is a *collection-level* lifecycle operation (create the +collection from the manifest schema, then load the data) — deliberately a +sibling of create/open rather than a variant of document import. +""" + +from __future__ import annotations + +import io +import tarfile +from pathlib import Path + +import pytest +from httpx import AsyncClient + +pytestmark = pytest.mark.integration + +API = "/api/v1" + + +def _payload(name: str = "demo") -> dict: + return { + "name": name, + "vectors": [ + { + "name": "embedding", + "dataType": "VECTOR_FP32", + "dimension": 4, + "indexParam": {"indexType": "FLAT", "metric": "L2"}, + } + ], + "fields": [{"name": "title", "dataType": "STRING"}], + } + + +async def _make_snapshot(client: AsyncClient, tmp_path: Path) -> Path: + """Create + seed a collection, then export it as a snapshot package.""" + resp = await client.post( + f"{API}/collections", + json={"path": str(tmp_path / "source"), "schema": _payload()}, + ) + assert resp.status_code == 201, resp.text + docs = [ + { + "id": f"doc-{i}", + "title": f"t{i}", + "embedding": [float(i), 0.0, 0.0, 0.0], + } + for i in range(3) + ] + resp = await client.post(f"{API}/collections/demo/documents", json={"documents": docs}) + assert resp.status_code == 201, resp.text + + resp = await client.get(f"{API}/collections/demo/documents:export?mode=snapshot") + assert resp.status_code == 200, resp.text + snapshot = tmp_path / "demo.tar.gz" + snapshot.write_bytes(resp.content) + return snapshot + + +def _restore_body(snapshot: Path, target: Path, **overrides: object) -> dict: + body: dict = { + "source": {"kind": "localPath", "path": str(snapshot)}, + "targetPath": str(target), + } + body.update(overrides) + return body + + +class TestRestoreSnapshot: + async def test_restores_collection_from_snapshot( + self, client: AsyncClient, tmp_path: Path + ) -> None: + snapshot = await _make_snapshot(client, tmp_path) + + # The exporter itself is still open under the manifest name, so the + # restore must rename (the default-name clash has its own test). + resp = await client.post( + f"{API}/collections:import", + json=_restore_body(snapshot, tmp_path / "restored", name="restored_demo"), + ) + + assert resp.status_code == 201, resp.text + body = resp.json() + assert body["collection"]["name"] == "restored_demo" + assert body["collection"]["path"].endswith("restored") + assert body["report"]["imported"] == 3 + assert body["report"]["failed"] == 0 + + # The restored collection is live and browsable. + got = await client.get(f"{API}/collections/restored_demo/documents/doc-1") + assert got.status_code == 200, got.text + assert got.json()["title"] == "t1" + + async def test_restore_with_name_override( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """The manifest name may clash with an open collection; renaming + the restored one is always possible.""" + snapshot = await _make_snapshot(client, tmp_path) + + resp = await client.post( + f"{API}/collections:import", + json=_restore_body(snapshot, tmp_path / "copy", name="demo_copy"), + ) + + assert resp.status_code == 201, resp.text + assert resp.json()["collection"]["name"] == "demo_copy" + got = await client.get(f"{API}/collections/demo_copy/documents/doc-0") + assert got.status_code == 200 + + async def test_restore_default_name_clash_is_409_with_guidance( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """The exporter's own collection is still open under the manifest + name — restoring with the default name must conflict with a message + that says to close it or pick another name (the exact scenario that + once produced two same-named open collections).""" + snapshot = await _make_snapshot(client, tmp_path) # 'demo' stays open + + resp = await client.post( + f"{API}/collections:import", + json=_restore_body(snapshot, tmp_path / "restored"), + ) + + assert resp.status_code == 409, resp.text + body = resp.json() + assert body["code"] == "COLLECTION_ALREADY_EXISTS" + assert "close" in body["detail"].lower() + assert "another name" in body["detail"].lower() + + async def test_missing_snapshot_is_404(self, client: AsyncClient, tmp_path: Path) -> None: + resp = await client.post( + f"{API}/collections:import", + json=_restore_body(tmp_path / "nope.tar.gz", tmp_path / "out"), + ) + assert resp.status_code == 404, resp.text + assert resp.json()["code"] == "IMPORT_FILE_NOT_FOUND" + + async def test_corrupt_snapshot_is_400(self, client: AsyncClient, tmp_path: Path) -> None: + broken = tmp_path / "broken.tar.gz" + broken.write_bytes(b"definitely not gzip") + + resp = await client.post( + f"{API}/collections:import", + json=_restore_body(broken, tmp_path / "out"), + ) + + assert resp.status_code == 400, resp.text + assert resp.json()["code"] == "IMPORT_MANIFEST_INVALID" + + async def test_tar_without_manifest_is_400( + self, client: AsyncClient, tmp_path: Path + ) -> None: + tar_path = tmp_path / "nomani.tar.gz" + with tarfile.open(tar_path, "w:gz") as tar: + data = b'{"id": "x"}\n' + info = tarfile.TarInfo("documents.jsonl") + info.size = len(data) + tar.addfile(info, io.BytesIO(data)) + + resp = await client.post( + f"{API}/collections:import", + json=_restore_body(tar_path, tmp_path / "out"), + ) + + assert resp.status_code == 400, resp.text + assert resp.json()["code"] == "IMPORT_MANIFEST_INVALID" + + async def test_existing_target_directory_is_409( + self, client: AsyncClient, tmp_path: Path + ) -> None: + snapshot = await _make_snapshot(client, tmp_path) + occupied = tmp_path / "occupied" + occupied.mkdir() + (occupied / "sentinel.txt").write_text("keep me") + + resp = await client.post( + f"{API}/collections:import", + json=_restore_body(snapshot, occupied), + ) + + assert resp.status_code == 409, resp.text + assert resp.json()["code"] == "COLLECTION_ALREADY_EXISTS" + assert (occupied / "sentinel.txt").exists() # untouched + + async def test_existing_empty_target_directory_is_409( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """The Zvec engine refuses any existing path, empty dirs included.""" + snapshot = await _make_snapshot(client, tmp_path) + empty = tmp_path / "empty-dir" + empty.mkdir() + + resp = await client.post( + f"{API}/collections:import", + json=_restore_body(snapshot, empty), + ) + + assert resp.status_code == 409, resp.text + assert resp.json()["code"] == "COLLECTION_ALREADY_EXISTS" + + async def test_invalid_override_name_is_422( + self, client: AsyncClient, tmp_path: Path + ) -> None: + snapshot = await _make_snapshot(client, tmp_path) + + resp = await client.post( + f"{API}/collections:import", + json=_restore_body(snapshot, tmp_path / "out", name="x"), # too short + ) + + assert resp.status_code == 422 + + async def test_vectorless_snapshot_rejected_before_anything_is_created( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """A snapshot exported with includeVector=false cannot repopulate a + vector-bearing collection. The mismatch must be detected BEFORE the + target directory is created — otherwise the failed import leaves a + ghost collection that permanently occupies ``targetPath``.""" + source_dir = tmp_path / "src" + resp = await client.post( + f"{API}/collections", + json={"path": str(source_dir), "schema": _payload()}, + ) + assert resp.status_code == 201, resp.text + await client.post( + f"{API}/collections/demo/documents", + json={"documents": [{"id": "d1", "title": "t", "embedding": [0.1, 0.2, 0.3, 0.4]}]}, + ) + + # Snapshot WITHOUT vectors. + resp = await client.get( + f"{API}/collections/demo/documents:export?mode=snapshot&includeVector=false" + ) + assert resp.status_code == 200, resp.text + snapshot = tmp_path / "novectors.tar.gz" + snapshot.write_bytes(resp.content) + + target = tmp_path / "out" + resp = await client.post(f"{API}/collections:import", json=_restore_body(snapshot, target)) + + assert resp.status_code == 409, resp.text + assert resp.json()["code"] == "IMPORT_SCHEMA_MISMATCH" + # Nothing was created: no directory, no open collection. + assert not target.exists() + listing = (await client.get(f"{API}/collections")).json()["items"] + assert all(i["name"] != "demo" or i["path"] != str(target) for i in listing) + + async def test_mid_load_failure_rolls_back_fresh_collection( + self, client: AsyncClient, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """If the data load blows up mid-stream after rows were written (a + corrupt member decompressing half-way is the real-world trigger — + see the truncation cases in :mod:`formats` tests and the guarded + reader in ``sdk._open_import_stream``), the freshly created + collection must be rolled back instead of being left half-imported + behind an error response. + + The load itself is stubbed to fail deterministically; what this test + pins is the rollback orchestration on the route. + """ + from zvec_studio.exceptions import InvalidDocumentError + from zvec_studio.storage.sdk import SdkBackend + + snapshot = await _make_snapshot(client, tmp_path) + + def boom(self, *args, **kwargs): + raise InvalidDocumentError("simulated mid-stream failure") + + monkeypatch.setattr(SdkBackend, "import_documents", boom) + + target = tmp_path / "rolled-back" + resp = await client.post( + f"{API}/collections:import", + json=_restore_body(snapshot, target, name="imported_demo"), + ) + + assert resp.status_code == 422, resp.text + assert not target.exists(), "rollback must remove the collection directory" + names = {i["name"] for i in (await client.get(f"{API}/collections")).json()["items"]} + assert names == {"demo"} diff --git a/apps/backend/tests/integration/test_import_endpoint.py b/apps/backend/tests/integration/test_import_endpoint.py new file mode 100644 index 0000000..983784b --- /dev/null +++ b/apps/backend/tests/integration/test_import_endpoint.py @@ -0,0 +1,266 @@ +"""Integration tests for ``POST /collections/{name}/documents:import``. + +Covers the HTTP contract from design doc §6.4: request validation, format +resolution (name / extension / default), source errors mapped to 4xx, and +row-level failures reported in the 200 body rather than as HTTP errors. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +from httpx import AsyncClient + +from .test_documents_router import _doc, _make_collection + +API = "/api/v1" +pytestmark = pytest.mark.integration + + +def _jsonl(tmp_path: Path, rows: list[dict], name: str = "data.jsonl") -> Path: + p = tmp_path / name + p.write_text("".join(json.dumps(r) + "\n" for r in rows), encoding="utf-8") + return p + + +class TestImportHappyPath: + async def test_import_jsonl_file(self, client: AsyncClient, tmp_path: Path) -> None: + name = await _make_collection(client, tmp_path) + source = _jsonl(tmp_path, [_doc(i) for i in range(3)]) + + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={"source": {"kind": "localPath", "path": str(source)}}, + ) + + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["imported"] == 3 + assert body["failed"] == 0 + assert body["totalLines"] == 3 + assert body["aborted"] is False + assert body["errors"] == [] + assert body["errorsTruncated"] is False + assert body["durationMs"] >= 0 + + got = await client.get(f"{API}/collections/{name}/documents/doc-001") + assert got.status_code == 200 + + async def test_replace_is_the_default_mode( + self, client: AsyncClient, tmp_path: Path + ) -> None: + name = await _make_collection(client, tmp_path) + await client.post( + f"{API}/collections/{name}/documents", json={"documents": [_doc(0)]} + ) + source = _jsonl(tmp_path, [{**_doc(0), "title": "new-title"}]) + + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={"source": {"kind": "localPath", "path": str(source)}}, + ) + + assert resp.status_code == 200, resp.text + assert resp.json()["imported"] == 1 + got = await client.get(f"{API}/collections/{name}/documents/doc-000") + assert got.json()["title"] == "new-title" + + async def test_insert_mode_conflicts_reported_in_body( + self, client: AsyncClient, tmp_path: Path + ) -> None: + name = await _make_collection(client, tmp_path) + await client.post( + f"{API}/collections/{name}/documents", json={"documents": [_doc(0)]} + ) + source = _jsonl(tmp_path, [_doc(0), _doc(1)]) + + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={ + "source": {"kind": "localPath", "path": str(source)}, + "mode": "insert", + }, + ) + + # Row-level failures live in the 200 body, not the status code. + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["aborted"] is True + assert body["failed"] == 1 + assert body["errors"][0]["code"] == "DOCUMENT_CONFLICT" + assert body["errors"][0]["line"] == 1 + + +class TestFormatResolution: + async def test_format_inferred_from_extension( + self, client: AsyncClient, tmp_path: Path + ) -> None: + name = await _make_collection(client, tmp_path) + source = _jsonl(tmp_path, [_doc(0)], name="payload.ndjson") + + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={"source": {"kind": "localPath", "path": str(source)}}, + ) + + assert resp.status_code == 200, resp.text + assert resp.json()["imported"] == 1 + + async def test_explicit_format_wins_over_extension( + self, client: AsyncClient, tmp_path: Path + ) -> None: + name = await _make_collection(client, tmp_path) + source = _jsonl(tmp_path, [_doc(0)], name="payload.txt") + + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={ + "source": {"kind": "localPath", "path": str(source)}, + "format": "jsonl", + }, + ) + + assert resp.status_code == 200, resp.text + + async def test_unknown_extension_is_400( + self, client: AsyncClient, tmp_path: Path + ) -> None: + name = await _make_collection(client, tmp_path) + source = _jsonl(tmp_path, [_doc(0)], name="data.csv") + + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={"source": {"kind": "localPath", "path": str(source)}}, + ) + + assert resp.status_code == 400, resp.text + assert resp.json()["code"] == "IMPORT_UNSUPPORTED_FORMAT" + + async def test_extensionless_file_defaults_to_jsonl( + self, client: AsyncClient, tmp_path: Path + ) -> None: + name = await _make_collection(client, tmp_path) + source = _jsonl(tmp_path, [_doc(0)], name="datafile") + + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={"source": {"kind": "localPath", "path": str(source)}}, + ) + + assert resp.status_code == 200, resp.text + + +class TestSourceErrors: + async def test_missing_file_is_404(self, client: AsyncClient, tmp_path: Path) -> None: + name = await _make_collection(client, tmp_path) + + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={"source": {"kind": "localPath", "path": str(tmp_path / "nope.jsonl")}}, + ) + + assert resp.status_code == 404, resp.text + assert resp.json()["code"] == "IMPORT_FILE_NOT_FOUND" + + async def test_directory_is_404(self, client: AsyncClient, tmp_path: Path) -> None: + name = await _make_collection(client, tmp_path) + + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={"source": {"kind": "localPath", "path": str(tmp_path)}}, + ) + + assert resp.status_code == 404, resp.text + assert resp.json()["code"] == "IMPORT_FILE_NOT_FOUND" + + async def test_unknown_collection_is_404( + self, client: AsyncClient, tmp_path: Path + ) -> None: + source = _jsonl(tmp_path, [_doc(0)]) + + resp = await client.post( + f"{API}/collections/ghost/documents:import", + json={"source": {"kind": "localPath", "path": str(source)}}, + ) + + assert resp.status_code == 404 + assert resp.json()["code"] == "COLLECTION_NOT_FOUND" + + async def test_corrupt_snapshot_is_400_not_500( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """A file named ``*.tar.gz`` that is not actually gzip is a user + input error (400), not a server failure.""" + name = await _make_collection(client, tmp_path) + source = tmp_path / "broken.tar.gz" + source.write_bytes(b"this is definitely not a gzip stream") + + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={"source": {"kind": "localPath", "path": str(source)}}, + ) + + assert resp.status_code == 400, resp.text + assert resp.json()["code"] == "IMPORT_MANIFEST_INVALID" + + async def test_truncated_snapshot_is_400_not_500( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """A valid gzip header followed by garbage (mid-stream corruption) + fails while reading members — still a user input error.""" + name = await _make_collection(client, tmp_path) + source = tmp_path / "truncated.tar.gz" + source.write_bytes(b"\x1f\x8b\x08\x00" + b"\x00" * 16 + b"garbage tail") + + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={"source": {"kind": "localPath", "path": str(source)}}, + ) + + assert resp.status_code == 400, resp.text + assert resp.json()["code"] == "IMPORT_MANIFEST_INVALID" + + +class TestRequestValidation: + async def test_batch_size_bounds(self, client: AsyncClient, tmp_path: Path) -> None: + name = await _make_collection(client, tmp_path) + source = _jsonl(tmp_path, [_doc(0)]) + + for bad in (0, 1025): + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={ + "source": {"kind": "localPath", "path": str(source)}, + "batchSize": bad, + }, + ) + assert resp.status_code == 422, resp.text + + async def test_invalid_mode_is_422(self, client: AsyncClient, tmp_path: Path) -> None: + name = await _make_collection(client, tmp_path) + source = _jsonl(tmp_path, [_doc(0)]) + + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={ + "source": {"kind": "localPath", "path": str(source)}, + "mode": "merge", + }, + ) + + assert resp.status_code == 422 + + async def test_unknown_source_kind_is_422( + self, client: AsyncClient, tmp_path: Path + ) -> None: + name = await _make_collection(client, tmp_path) + source = _jsonl(tmp_path, [_doc(0)]) + + resp = await client.post( + f"{API}/collections/{name}/documents:import", + json={"source": {"kind": "upload", "path": str(source)}}, + ) + + assert resp.status_code == 422 diff --git a/apps/backend/tests/integration/test_import_pipeline.py b/apps/backend/tests/integration/test_import_pipeline.py new file mode 100644 index 0000000..e02ffeb --- /dev/null +++ b/apps/backend/tests/integration/test_import_pipeline.py @@ -0,0 +1,302 @@ +"""Integration tests for the import pipeline (`SdkBackend.import_documents`). + +Exercises the real SDK: file parsing (JSONL), internal batching, write modes +(`insert` / `replace`), error policies (`abort` / `skip`), and the report +shape defined in design doc §6.4. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +zvec = pytest.importorskip("zvec") + +from zvec_studio.exceptions import ( # noqa: E402 (import after skip) + ImportFileNotFoundError, + ImportFileNotReadableError, +) +from zvec_studio.schemas.collection import CollectionSchema # noqa: E402 +from zvec_studio.storage.formats import JsonlFormat # noqa: E402 +from zvec_studio.storage.import_ import ImportMode, OnErrorMode # noqa: E402 +from zvec_studio.storage.sdk import SdkBackend # noqa: E402 + +pytestmark = pytest.mark.integration + +DIM = 4 + + +def _schema(name: str = "import_target") -> CollectionSchema: + return CollectionSchema.model_validate( + { + "name": name, + "vectors": [ + { + "name": "embedding", + "dataType": "VECTOR_FP32", + "dimension": DIM, + "indexParam": {"indexType": "FLAT", "metric": "L2"}, + } + ], + "fields": [ + {"name": "title", "dataType": "STRING"}, + {"name": "score", "dataType": "INT64"}, + ], + } + ) + + +def _make_backend_and_collection(tmp_path: Path, name: str = "import_target") -> SdkBackend: + backend = SdkBackend() + backend.create(path=tmp_path / name, schema=_schema(name)) + return backend + + +def _write_jsonl(tmp_path: Path, rows: list[dict], file_name: str = "data.jsonl") -> Path: + path = tmp_path / file_name + path.write_text("".join(json.dumps(r) + "\n" for r in rows), encoding="utf-8") + return path + + +def _row(i: int, **overrides: object) -> dict: + base: dict = { + "id": f"doc-{i:04d}", + "title": f"t{i}", + "score": i, + "embedding": [float(i), 0.0, 0.0, 0.0], + } + base.update(overrides) + return base + + +class TestBasicImport: + def test_imports_rows_and_reports(self, tmp_path: Path) -> None: + backend = _make_backend_and_collection(tmp_path) + source = _write_jsonl(tmp_path, [_row(i) for i in range(3)]) + + report = backend.import_documents( + "import_target", + source_path=str(source), + fmt=JsonlFormat(), + ) + + assert report.imported == 3 + assert report.failed == 0 + assert report.total_lines == 3 + assert report.aborted is False + assert report.errors == [] + assert backend.get_document("import_target", "doc-0001")["title"] == "t1" + + def test_blank_lines_are_skipped_and_numbering_kept(self, tmp_path: Path) -> None: + backend = _make_backend_and_collection(tmp_path) + source = tmp_path / "gaps.jsonl" + source.write_text( + json.dumps(_row(0)) + "\n\n" + json.dumps(_row(1)) + "\n", + encoding="utf-8", + ) + + report = backend.import_documents( + "import_target", source_path=str(source), fmt=JsonlFormat() + ) + + assert (report.imported, report.total_lines) == (2, 2) + + def test_batches_larger_than_sdk_limit(self, tmp_path: Path) -> None: + """600 rows crosses the 512-doc batch boundary.""" + backend = _make_backend_and_collection(tmp_path) + source = _write_jsonl(tmp_path, [_row(i) for i in range(600)]) + + report = backend.import_documents( + "import_target", source_path=str(source), fmt=JsonlFormat() + ) + + assert report.imported == 600 + assert backend.stats("import_target").documentCount == 600 + + def test_empty_file(self, tmp_path: Path) -> None: + backend = _make_backend_and_collection(tmp_path) + source = tmp_path / "empty.jsonl" + source.write_text("", encoding="utf-8") + + report = backend.import_documents( + "import_target", source_path=str(source), fmt=JsonlFormat() + ) + + assert (report.imported, report.failed, report.total_lines) == (0, 0, 0) + + +class TestWriteModes: + def test_replace_mode_overwrites_whole_document(self, tmp_path: Path) -> None: + backend = _make_backend_and_collection(tmp_path) + backend.insert_documents("import_target", [_row(0, title="original")]) + + source = _write_jsonl(tmp_path, [_row(0, title="replaced", score=99)]) + report = backend.import_documents( + "import_target", + source_path=str(source), + fmt=JsonlFormat(), + mode=ImportMode.REPLACE, + ) + + assert report.imported == 1 + doc = backend.get_document("import_target", "doc-0000") + assert doc["title"] == "replaced" + assert doc["score"] == 99 + assert backend.stats("import_target").documentCount == 1 + + def test_insert_mode_rejects_duplicate_and_aborts(self, tmp_path: Path) -> None: + backend = _make_backend_and_collection(tmp_path) + backend.insert_documents("import_target", [_row(0)]) + + source = _write_jsonl(tmp_path, [_row(0, title="dup"), _row(1)]) + report = backend.import_documents( + "import_target", + source_path=str(source), + fmt=JsonlFormat(), + mode=ImportMode.INSERT, + on_error=OnErrorMode.ABORT, + ) + + assert report.aborted is True + assert report.failed == 1 + assert report.errors[0].line == 1 + assert report.errors[0].code == "DOCUMENT_CONFLICT" + # Abort granularity is the batch: both rows were already submitted to + # the SDK, whose per-doc statuses are independent — the valid row in + # the same batch lands. No further rows are read afterwards. + assert report.imported == 1 + assert backend.stats("import_target").documentCount == 2 + + def test_insert_mode_skip_records_and_continues(self, tmp_path: Path) -> None: + backend = _make_backend_and_collection(tmp_path) + backend.insert_documents("import_target", [_row(0)]) + + source = _write_jsonl(tmp_path, [_row(0, title="dup"), _row(1), _row(2)]) + report = backend.import_documents( + "import_target", + source_path=str(source), + fmt=JsonlFormat(), + mode=ImportMode.INSERT, + on_error=OnErrorMode.SKIP, + ) + + assert report.aborted is False + assert (report.imported, report.failed) == (2, 1) + assert [e.line for e in report.errors] == [1] + assert backend.stats("import_target").documentCount == 3 + + +class TestErrorPolicies: + def test_abort_on_invalid_row_keeps_prior_rows(self, tmp_path: Path) -> None: + backend = _make_backend_and_collection(tmp_path) + # Row 2 has a dimension mismatch. + rows = [_row(0), _row(1, embedding=[0.0, 0.0]), _row(2)] + source = _write_jsonl(tmp_path, rows) + + report = backend.import_documents( + "import_target", + source_path=str(source), + fmt=JsonlFormat(), + on_error=OnErrorMode.ABORT, + ) + + assert report.aborted is True + assert (report.imported, report.failed) == (1, 1) + assert report.errors[0].line == 2 + assert report.errors[0].code == "DIMENSION_MISMATCH" + # Row 0 landed and was flushed; row 2 never written. + assert backend.stats("import_target").documentCount == 1 + + def test_skip_invalid_rows_and_continue(self, tmp_path: Path) -> None: + backend = _make_backend_and_collection(tmp_path) + rows = [ + _row(0), + _row(1, embedding=[0.0, 0.0]), # dimension mismatch + _row(2, unknownColumn="x"), # unknown column + _row(3), + ] + source = _write_jsonl(tmp_path, rows) + + report = backend.import_documents( + "import_target", + source_path=str(source), + fmt=JsonlFormat(), + on_error=OnErrorMode.SKIP, + ) + + assert report.aborted is False + assert (report.imported, report.failed) == (2, 2) + assert [(e.line, e.code) for e in report.errors] == [ + (2, "DIMENSION_MISMATCH"), + (3, "INVALID_SCHEMA"), + ] + + def test_skip_malformed_json_line(self, tmp_path: Path) -> None: + backend = _make_backend_and_collection(tmp_path) + source = tmp_path / "bad.jsonl" + source.write_text( + json.dumps(_row(0)) + "\n" + "not json" + "\n" + json.dumps(_row(1)) + "\n", + encoding="utf-8", + ) + + report = backend.import_documents( + "import_target", + source_path=str(source), + fmt=JsonlFormat(), + on_error=OnErrorMode.SKIP, + ) + + assert (report.imported, report.failed) == (2, 1) + assert report.errors[0].line == 2 + assert report.errors[0].code == "INVALID_DOCUMENT" + + def test_error_report_is_capped(self, tmp_path: Path) -> None: + backend = _make_backend_and_collection(tmp_path) + rows = [_row(i, embedding=[0.0]) for i in range(150)] + source = _write_jsonl(tmp_path, rows) + + report = backend.import_documents( + "import_target", + source_path=str(source), + fmt=JsonlFormat(), + on_error=OnErrorMode.SKIP, + ) + + assert report.failed == 150 + assert len(report.errors) == 100 + assert report.errors_truncated is True + + +class TestSourceValidation: + def test_missing_file(self, tmp_path: Path) -> None: + backend = _make_backend_and_collection(tmp_path) + with pytest.raises(ImportFileNotFoundError) as exc: + backend.import_documents( + "import_target", + source_path=str(tmp_path / "nope.jsonl"), + fmt=JsonlFormat(), + ) + assert exc.value.status_code == 404 + + def test_directory_is_not_a_valid_source(self, tmp_path: Path) -> None: + backend = _make_backend_and_collection(tmp_path) + with pytest.raises(ImportFileNotFoundError) as exc: + backend.import_documents( + "import_target", source_path=str(tmp_path), fmt=JsonlFormat() + ) + assert exc.value.status_code == 404 + + def test_unreadable_file(self, tmp_path: Path) -> None: + backend = _make_backend_and_collection(tmp_path) + source = _write_jsonl(tmp_path, [_row(0)]) + source.chmod(0o000) + try: + with pytest.raises(ImportFileNotReadableError) as exc: + backend.import_documents( + "import_target", source_path=str(source), fmt=JsonlFormat() + ) + assert exc.value.status_code == 403 + finally: + source.chmod(0o644) diff --git a/apps/backend/tests/integration/test_snapshot_io.py b/apps/backend/tests/integration/test_snapshot_io.py new file mode 100644 index 0000000..f97ad1f --- /dev/null +++ b/apps/backend/tests/integration/test_snapshot_io.py @@ -0,0 +1,230 @@ +"""Integration tests for snapshot-mode export/import (``mode=snapshot``). + +A snapshot is ``manifest.json`` + ``documents.jsonl`` in a tar.gz. Covered: +* snapshot export streams a readable gzip tar with both members; +* snapshot import reads the manifest, pre-checks schema compatibility, and + only then writes rows; +* schema mismatch fails fast (409) with zero rows written; +* export -> snapshot file -> import roundtrip is lossless. +""" + +from __future__ import annotations + +import io +import json +import tarfile +from pathlib import Path + +import pytest +from httpx import AsyncClient + +zvec = pytest.importorskip("zvec") + +pytestmark = pytest.mark.integration + +API = "/api/v1" + + +def _collection_payload(name: str, dim: int = 4) -> dict: + return { + "name": name, + "vectors": [ + { + "name": "embedding", + "dataType": "VECTOR_FP32", + "dimension": dim, + "indexParam": {"indexType": "FLAT", "metric": "L2"}, + } + ], + "fields": [ + {"name": "title", "dataType": "STRING"}, + {"name": "score", "dataType": "INT64"}, + ], + } + + +async def _seed(client: AsyncClient, tmp_path: Path, name: str, count: int) -> None: + resp = await client.post( + f"{API}/collections", + json={"path": str(tmp_path / name), "schema": _collection_payload(name)}, + ) + assert resp.status_code == 201, resp.text + docs = [ + { + "id": f"doc-{i:03d}", + "title": f"t{i}", + "score": i, + "embedding": [float(i), 0.0, 0.0, 0.0], + } + for i in range(count) + ] + if docs: + resp = await client.post( + f"{API}/collections/{name}/documents", json={"documents": docs} + ) + assert resp.status_code == 201, resp.text + + +class TestSnapshotExport: + async def test_exports_tar_gz_with_manifest_and_data( + self, client: AsyncClient, tmp_path: Path + ) -> None: + await _seed(client, tmp_path, "demo", 3) + + resp = await client.get( + f"{API}/collections/demo/documents:export", params={"mode": "snapshot"} + ) + + assert resp.status_code == 200, resp.text + assert resp.headers["content-type"].startswith("application/gzip") + assert ".tar.gz" in resp.headers["content-disposition"] + + with tarfile.open(fileobj=io.BytesIO(resp.content), mode="r:gz") as tar: + names = tar.getnames() + assert "manifest.json" in names + assert "documents.jsonl" in names + + manifest = json.load(tar.extractfile("manifest.json")) + assert manifest["format"] == "zvec-studio.export/1" + assert manifest["collection"]["name"] == "demo" + assert manifest["options"]["includeVector"] is True + + data = tar.extractfile("documents.jsonl").read().decode("utf-8") + rows = [json.loads(line) for line in data.splitlines() if line] + assert len(rows) == 3 + assert rows[0]["id"] == "doc-000" + + async def test_snapshot_honours_include_vector_false( + self, client: AsyncClient, tmp_path: Path + ) -> None: + await _seed(client, tmp_path, "demo", 1) + + resp = await client.get( + f"{API}/collections/demo/documents:export", + params={"mode": "snapshot", "includeVector": "false"}, + ) + + with tarfile.open(fileobj=io.BytesIO(resp.content), mode="r:gz") as tar: + manifest = json.load(tar.extractfile("manifest.json")) + data = tar.extractfile("documents.jsonl").read().decode("utf-8") + + assert manifest["options"]["includeVector"] is False + assert "embedding" not in json.loads(data.strip()) + + +class TestSnapshotImport: + async def _export_to_file( + self, client: AsyncClient, tmp_path: Path, source: str, count: int + ) -> Path: + await _seed(client, tmp_path, source, count) + resp = await client.get( + f"{API}/collections/{source}/documents:export", params={"mode": "snapshot"} + ) + assert resp.status_code == 200, resp.text + path = tmp_path / "snapshot.tar.gz" + path.write_bytes(resp.content) + return path + + async def test_import_snapshot_roundtrip( + self, client: AsyncClient, tmp_path: Path + ) -> None: + path = await self._export_to_file(client, tmp_path, "source", 4) + resp = await client.post( + f"{API}/collections", + json={"path": str(tmp_path / "target"), "schema": _collection_payload("target")}, + ) + assert resp.status_code == 201, resp.text + + imported = await client.post( + f"{API}/collections/target/documents:import", + json={"source": {"kind": "localPath", "path": str(path)}}, + ) + + assert imported.status_code == 200, imported.text + assert imported.json()["imported"] == 4 + assert imported.json()["failed"] == 0 + + rows = ( + await client.post( + f"{API}/collections/target/documents:browse", + json={"limit": 100, "includeVector": True}, + ) + ).json()["items"] + assert {r["id"] for r in rows} == {f"doc-{i:03d}" for i in range(4)} + + async def test_schema_mismatch_is_409_and_writes_nothing( + self, client: AsyncClient, tmp_path: Path + ) -> None: + # Snapshot was exported with dimension 4; target declares dimension 8. + path = await self._export_to_file(client, tmp_path, "source", 3) + resp = await client.post( + f"{API}/collections", + json={"path": str(tmp_path / "target"), "schema": _collection_payload("target", dim=8)}, + ) + assert resp.status_code == 201, resp.text + + imported = await client.post( + f"{API}/collections/target/documents:import", + json={"source": {"kind": "localPath", "path": str(path)}}, + ) + + assert imported.status_code == 409, imported.text + assert imported.json()["code"] == "IMPORT_SCHEMA_MISMATCH" + assert imported.json()["mismatches"] + + # No rows may have been written. + stats = (await client.get(f"{API}/collections/target/stats")).json() + assert stats["documentCount"] == 0 + + async def test_vectors_not_exported_but_required_is_409( + self, client: AsyncClient, tmp_path: Path + ) -> None: + await _seed(client, tmp_path, "source", 2) + resp = await client.get( + f"{API}/collections/source/documents:export", + params={"mode": "snapshot", "includeVector": "false"}, + ) + path = tmp_path / "novectors.tar.gz" + path.write_bytes(resp.content) + + resp = await client.post( + f"{API}/collections", + json={"path": str(tmp_path / "target"), "schema": _collection_payload("target")}, + ) + assert resp.status_code == 201, resp.text + + imported = await client.post( + f"{API}/collections/target/documents:import", + json={"source": {"kind": "localPath", "path": str(path)}}, + ) + + assert imported.status_code == 409, imported.text + assert imported.json()["code"] == "IMPORT_SCHEMA_MISMATCH" + + async def test_corrupt_manifest_is_400( + self, client: AsyncClient, tmp_path: Path + ) -> None: + await _seed(client, tmp_path, "source", 1) + # Build a tar.gz whose manifest.json is not valid JSON. + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tar: + info = tarfile.TarInfo("manifest.json") + payload = b"not json" + info.size = len(payload) + tar.addfile(info, io.BytesIO(payload)) + path = tmp_path / "bad.tar.gz" + path.write_bytes(buf.getvalue()) + + resp = await client.post( + f"{API}/collections", + json={"path": str(tmp_path / "target"), "schema": _collection_payload("target")}, + ) + assert resp.status_code == 201, resp.text + + imported = await client.post( + f"{API}/collections/target/documents:import", + json={"source": {"kind": "localPath", "path": str(path)}}, + ) + + assert imported.status_code == 400, imported.text + assert imported.json()["code"] == "IMPORT_MANIFEST_INVALID" diff --git a/apps/backend/tests/integration/test_write_limits.py b/apps/backend/tests/integration/test_write_limits.py new file mode 100644 index 0000000..be6454d --- /dev/null +++ b/apps/backend/tests/integration/test_write_limits.py @@ -0,0 +1,138 @@ +"""Boundary regressions for the document write surface. + +Two shipped defects are pinned here: + +1. ``insert`` / ``upsert`` / ``update`` accepted up to 10,000 documents by + contract but passed them straight to the SDK, whose (undocumented) write + batch limit is 1024 -- anything larger answered HTTP 500. +2. SDK-level document validation errors and duplicate-id writes fell through + the generic handler as HTTP 500 instead of 4xx. + +The service layer absorbs the SDK batch constraint internally (see +``_write_in_batches``) and classifies errors per the RFC 7807 catalog. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from httpx import AsyncClient + +from .test_documents_router import _doc, _make_collection + +API = "/api/v1" +pytestmark = pytest.mark.integration + + +class TestBatchBoundary: + async def test_insert_1025_docs_succeeds_via_internal_batching( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """Regression: this used to return HTTP 500 (INTERNAL_ERROR).""" + name = await _make_collection(client, tmp_path) + + resp = await client.post( + f"{API}/collections/{name}/documents", + json={"documents": [_doc(i) for i in range(1025)]}, + ) + + assert resp.status_code == 201, resp.text + body = resp.json() + assert body["inserted"] == 1025 + assert len(body["ids"]) == 1025 + + # The documents really landed: fetch one from the second internal batch. + got = await client.get(f"{API}/collections/{name}/documents/doc-1024") + assert got.status_code == 200 + + async def test_insert_1024_docs_control(self, client: AsyncClient, tmp_path: Path) -> None: + name = await _make_collection(client, tmp_path) + resp = await client.post( + f"{API}/collections/{name}/documents", + json={"documents": [_doc(i) for i in range(1024)]}, + ) + assert resp.status_code == 201, resp.text + assert resp.json()["inserted"] == 1024 + + async def test_update_1025_docs_succeeds_via_internal_batching( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """Regression: PATCH with >1024 documents used to return HTTP 500.""" + name = await _make_collection(client, tmp_path) + await client.post( + f"{API}/collections/{name}/documents", + json={"documents": [_doc(i) for i in range(1025)]}, + ) + + resp = await client.patch( + f"{API}/collections/{name}/documents", + json={"documents": [{"id": f"doc-{i:03d}", "title": "updated"} for i in range(1025)]}, + ) + + assert resp.status_code == 200, resp.text + assert resp.json()["updated"] == 1025 + + async def test_upsert_1025_docs_succeeds_via_internal_batching( + self, client: AsyncClient, tmp_path: Path + ) -> None: + name = await _make_collection(client, tmp_path) + resp = await client.post( + f"{API}/collections/{name}/documents:upsert", + json={"documents": [_doc(i) for i in range(1025)]}, + ) + assert resp.status_code == 200, resp.text + + +class TestWriteErrorClassification: + async def test_insert_duplicate_id_returns_409( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """The SDK reports the duplicate through a per-doc Status, not an exception.""" + name = await _make_collection(client, tmp_path) + await client.post( + f"{API}/collections/{name}/documents", json={"documents": [_doc(1)]} + ) + + resp = await client.post( + f"{API}/collections/{name}/documents", json={"documents": [_doc(1)]} + ) + + assert resp.status_code == 409, resp.text + assert resp.json()["code"] == "DOCUMENT_CONFLICT" + + async def test_insert_missing_required_field_returns_422( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """Regression: SDK ValueError used to fall through to HTTP 500.""" + name = await _make_collection(client, tmp_path) + + resp = await client.post( + f"{API}/collections/{name}/documents", + json={"documents": [{"id": "x", "title": "t", "embedding": [0.1, 0.2, 0.3, 0.4]}]}, + ) + + assert resp.status_code == 422, resp.text + assert resp.json()["code"] == "INVALID_DOCUMENT" + + +class TestPrimaryKeyContract: + async def test_explicit_null_id_auto_generates( + self, client: AsyncClient, tmp_path: Path + ) -> None: + """Historical write contract: ``"id": null`` behaves like omitting + the key — a fresh ULID is minted, not a 422.""" + name = await _make_collection(client, tmp_path) + + resp = await client.post( + f"{API}/collections/{name}/documents", + json={ + "documents": [ + {"id": None, "title": "t", "score": 1, "embedding": [0.1, 0.2, 0.3, 0.4]} + ] + }, + ) + + assert resp.status_code == 201, resp.text + ids = resp.json()["ids"] + assert len(ids) == 1 and ids[0] # a non-empty generated key diff --git a/apps/backend/tests/unit/test_collections_schemas.py b/apps/backend/tests/unit/test_collections_schemas.py index a4e77ea..1c57ad4 100644 --- a/apps/backend/tests/unit/test_collections_schemas.py +++ b/apps/backend/tests/unit/test_collections_schemas.py @@ -3,9 +3,13 @@ The v0.2.0 contract drops ``isPrimary`` / ``description`` / ``JSON`` / collection-level ``indexParams`` and keeps: - per-vector ``indexParam`` ({indexType, metric, params}) -- reserved field names ``{"id", "_id"}`` -- stricter collection-name regex (start with letter, len ≥ 3). +- stricter collection-name regex (start with letter, len >= 3). - Zvec 0.6.x FTS scalar indexes and DiskANN vector params. + +Column names ``id`` / ``_id`` are *accepted*: Zvec allows them, and rejecting +them stopped Studio from opening SDK-created collections. When a column takes +``id``, the primary key moves to the ``$id`` row key (see +``storage/doc_repr.py``). """ from __future__ import annotations @@ -60,11 +64,17 @@ def test_invalid_name(self, bad_name: str) -> None: with pytest.raises(ValidationError): FieldSchema(name=bad_name, dataType=ScalarDataType.INT64) - @pytest.mark.parametrize("reserved", ["id", "_id"]) - def test_reserved_field_names_rejected(self, reserved: str) -> None: - with pytest.raises(ValidationError) as exc: - FieldSchema(name=reserved, dataType=ScalarDataType.STRING) - assert "reserved" in str(exc.value) + @pytest.mark.parametrize("name", ["id", "_id"]) + def test_columns_may_be_named_id(self, name: str) -> None: + """Zvec allows these names; rejecting them blocked opening SDK collections.""" + assert FieldSchema(name=name, dataType=ScalarDataType.STRING).name == name + + def test_dollar_prefixed_scalar_field_name_rejected(self) -> None: + """Pin the guard the reserved pk chain relies on: Studio (and Zvec) + reject ``$`` in scalar field names, so ``$id``/``$$id`` can never be + user columns.""" + with pytest.raises(ValidationError): + FieldSchema(name="$id", dataType=ScalarDataType.STRING) def test_extra_forbidden(self) -> None: with pytest.raises(ValidationError): @@ -115,13 +125,30 @@ def test_dimension_bounds(self, bad_dim: int) -> None: with pytest.raises(ValidationError): VectorSchema(name="v", dataType=VectorDataType.VECTOR_FP32, dimension=bad_dim) - @pytest.mark.parametrize("reserved", ["id", "_id"]) - def test_reserved_vector_names_rejected(self, reserved: str) -> None: + @pytest.mark.parametrize("name", ["id", "_id"]) + def test_vectors_may_be_named_id(self, name: str) -> None: + """Zvec accepts a vector named ``id`` (it just cannot clash with a field).""" + v = VectorSchema(name=name, dataType=VectorDataType.VECTOR_FP32, dimension=4) + assert v.name == name + + def test_vector_fp64_is_rejected_at_validation(self) -> None: + """Regression: zvec rejects FP64 dense vectors at create time with an + opaque error; refuse the value up front instead (defect §4.3 in the + import/export design doc).""" with pytest.raises(ValidationError): VectorSchema( - name=reserved, dataType=VectorDataType.VECTOR_FP32, dimension=4 + name="v", dataType="VECTOR_FP64", dimension=4 # type: ignore[arg-type] ) + def test_supported_vector_types(self) -> None: + assert VectorDataType.__members__.keys() == { + "VECTOR_FP32", + "VECTOR_FP16", + "VECTOR_INT8", + "SPARSE_VECTOR_FP32", + "SPARSE_VECTOR_FP16", + } + def test_index_param_attached(self) -> None: v = VectorSchema.model_validate( { diff --git a/apps/backend/tests/unit/test_doc_repr.py b/apps/backend/tests/unit/test_doc_repr.py new file mode 100644 index 0000000..8636ca2 --- /dev/null +++ b/apps/backend/tests/unit/test_doc_repr.py @@ -0,0 +1,269 @@ +"""Unit tests for the document wire representation (``storage/doc_repr.py``). + +These cover the primary-key representation rules from +``design/import-export-design.md`` §6.2, including the regression that +motivated them: a collection whose schema declares its own scalar field +named ``id`` used to lose the real primary key. +""" + +from __future__ import annotations + +import pytest + +from zvec_studio.exceptions import InvalidSchemaError +from zvec_studio.schemas import ( + CollectionSchema, + FieldSchema, + ScalarDataType, + VectorDataType, + VectorSchema, +) +from zvec_studio.storage.doc_repr import ( + PK_KEY, + RESERVED_PK_KEY, + doc_to_row, + pk_key, + split_pk, +) + + +class _FakeDoc: + """Stand-in for ``zvec.Doc`` (id / fields / vectors triple).""" + + def __init__( + self, + doc_id: str, + fields: dict[str, object] | None = None, + vectors: dict[str, object] | None = None, + ) -> None: + self.id = doc_id + self.fields = fields + self.vectors = vectors + + +def _schema(*, with_id_field: bool) -> CollectionSchema: + """Build a schema, optionally declaring a scalar field named ``id``. + + ``FieldSchema`` currently rejects the name ``id`` at validation time, but + Zvec itself allows it — so collections created through the SDK can and do + have such a field. ``model_construct`` bypasses validation to model that + reality, keeping these tests independent of whatever policy the create + path adopts. + """ + fields = [FieldSchema(name="title", dataType=ScalarDataType.STRING)] + if with_id_field: + fields.insert( + 0, FieldSchema.model_construct(name="id", dataType=ScalarDataType.STRING) + ) + return CollectionSchema.model_construct( + name="demo", + vectors=[ + VectorSchema(name="embedding", dataType=VectorDataType.VECTOR_FP32, dimension=2) + ], + fields=fields, + ) + + +def _schema_with_id_vector() -> CollectionSchema: + """Schema whose *vector* is named ``id`` (Zvec allows this too). + + Zvec rejects a schema where a scalar field and a vector share a name, so + ``id`` is either a scalar field or a vector, never both. + """ + return CollectionSchema.model_construct( + name="demo", + vectors=[ + VectorSchema.model_construct( + name="id", dataType=VectorDataType.VECTOR_FP32, dimension=2 + ) + ], + fields=[FieldSchema(name="title", dataType=ScalarDataType.STRING)], + ) + + +class TestPkKey: + def test_plain_schema_uses_id(self) -> None: + assert pk_key(_schema(with_id_field=False)) == PK_KEY == "id" + + def test_schema_with_id_field_uses_reserved_key(self) -> None: + assert pk_key(_schema(with_id_field=True)) == RESERVED_PK_KEY == "$id" + + def test_schema_with_id_vector_uses_reserved_key(self) -> None: + """A vector named ``id`` shadows the primary key just as a field does.""" + assert pk_key(_schema_with_id_vector()) == RESERVED_PK_KEY + + +class TestDocToRow: + def test_plain_schema_flattens_under_id(self) -> None: + doc = _FakeDoc("doc-1", fields={"title": "hello"}, vectors={"embedding": [0.1, 0.2]}) + row = doc_to_row(doc, schema=_schema(with_id_field=False), include_vector=True) + assert row == {"id": "doc-1", "title": "hello", "embedding": [0.1, 0.2]} + + def test_include_vector_false_omits_vectors(self) -> None: + doc = _FakeDoc("doc-1", fields={"title": "hello"}, vectors={"embedding": [0.1, 0.2]}) + row = doc_to_row(doc, schema=_schema(with_id_field=False), include_vector=False) + assert row == {"id": "doc-1", "title": "hello"} + + def test_id_field_does_not_shadow_primary_key(self) -> None: + """Regression: the primary key used to be overwritten by the field.""" + doc = _FakeDoc("PK-001", fields={"id": "USER-999", "title": "t"}) + row = doc_to_row(doc, schema=_schema(with_id_field=True), include_vector=False) + assert row == {"$id": "PK-001", "id": "USER-999", "title": "t"} + + def test_doc_without_fields_or_vectors(self) -> None: + row = doc_to_row(_FakeDoc("doc-1"), schema=_schema(with_id_field=False), include_vector=True) + assert row == {"id": "doc-1"} + + def test_id_vector_does_not_shadow_primary_key(self) -> None: + doc = _FakeDoc("PK-001", fields={"title": "t"}, vectors={"id": [0.1, 0.2]}) + row = doc_to_row(doc, schema=_schema_with_id_vector(), include_vector=True) + assert row == {"$id": "PK-001", "title": "t", "id": [0.1, 0.2]} + + +class TestSplitPk: + def test_reserved_key_wins(self) -> None: + row = {"$id": "PK-001", "id": "USER-999", "title": "t"} + found, rest = split_pk(row, schema=_schema(with_id_field=True)) + assert found == "PK-001" + assert rest == {"id": "USER-999", "title": "t"} + + def test_reserved_key_accepted_even_without_id_field(self) -> None: + """Keeps exports portable across schemas that do/don't declare ``id``.""" + found, rest = split_pk({"$id": "PK-1", "title": "t"}, schema=_schema(with_id_field=False)) + assert found == "PK-1" + assert rest == {"title": "t"} + + def test_plain_id_is_the_primary_key(self) -> None: + found, rest = split_pk({"id": "doc-1", "title": "t"}, schema=_schema(with_id_field=False)) + assert found == "doc-1" + assert rest == {"title": "t"} + + def test_ambiguous_id_is_rejected(self) -> None: + """schema has an ``id`` field and the row carries only ``id``.""" + with pytest.raises(InvalidSchemaError) as exc: + split_pk({"id": "???", "title": "t"}, schema=_schema(with_id_field=True)) + assert "$id" in str(exc.value) + + def test_ambiguous_id_vector_is_rejected(self) -> None: + with pytest.raises(InvalidSchemaError): + split_pk({"id": [0.1, 0.2], "title": "t"}, schema=_schema_with_id_vector()) + + def test_missing_primary_key_returns_none(self) -> None: + found, rest = split_pk({"title": "t"}, schema=_schema(with_id_field=False)) + assert found is None + assert rest == {"title": "t"} + + def test_non_string_primary_key_is_rejected(self) -> None: + with pytest.raises(InvalidSchemaError): + split_pk({"id": 123, "title": "t"}, schema=_schema(with_id_field=False)) + + def test_non_string_reserved_primary_key_is_rejected(self) -> None: + with pytest.raises(InvalidSchemaError): + split_pk({"$id": 123}, schema=_schema(with_id_field=True)) + + def test_input_row_is_not_mutated(self) -> None: + row = {"id": "doc-1", "title": "t"} + split_pk(row, schema=_schema(with_id_field=False)) + assert row == {"id": "doc-1", "title": "t"} + + def test_explicit_null_id_means_generate(self) -> None: + """Legacy contract: an explicit ``"id": null`` behaves like omitting + the key — the caller mints a fresh ULID (not a 422).""" + found, rest = split_pk({"id": None, "title": "t"}, schema=_schema(with_id_field=False)) + assert found is None + assert rest == {"title": "t"} + + def test_explicit_null_reserved_id_means_generate(self) -> None: + found, rest = split_pk({"$id": None, "id": "U", "title": "t"}, schema=_schema(with_id_field=True)) + assert found is None + assert rest == {"id": "U", "title": "t"} + + def test_explicit_null_id_is_still_ambiguous_with_id_column(self) -> None: + """A bare ``id`` — even null — stays ambiguous when the schema + declares its own ``id`` column (could be the nullable field).""" + with pytest.raises(InvalidSchemaError): + split_pk({"id": None, "title": "t"}, schema=_schema(with_id_field=True)) + + +def _schema_with_dollar_id_vector(*, with_id_field: bool = False) -> CollectionSchema: + """A schema whose *vector* is named ``$id``. + + The Zvec engine rejects such names at create time today, but a + hand-crafted manifest (or a future engine) could carry one — the chain + must stay total either way. ``model_construct`` bypasses Studio's own + validation to exercise exactly that. + """ + fields = [FieldSchema(name="title", dataType=ScalarDataType.STRING)] + if with_id_field: + fields.insert( + 0, FieldSchema.model_construct(name="id", dataType=ScalarDataType.STRING) + ) + return CollectionSchema.model_construct( + name="demo", + vectors=[ + VectorSchema.model_construct( + name="$id", dataType=VectorDataType.VECTOR_FP32, dimension=2 + ) + ], + fields=fields, + ) + + +class TestReservedChain: + """The reserved key chain ``id`` -> ``$id`` -> ``$$id`` -> ...""" + + def test_dollar_id_vector_alone_does_not_take_the_plain_key(self) -> None: + """A vector named ``$id`` occupies ``$id`` only; ``id`` stays free.""" + assert pk_key(_schema_with_dollar_id_vector()) == "id" + + def test_id_field_plus_dollar_id_vector_pushes_pk_two_levels_down(self) -> None: + schema = _schema_with_dollar_id_vector(with_id_field=True) + assert pk_key(schema) == "$$id" + + def test_doc_to_row_keeps_pk_and_dollar_id_vector_apart(self) -> None: + schema = _schema_with_dollar_id_vector(with_id_field=True) + doc = _FakeDoc( + "PK-001", fields={"id": "USER-9", "title": "t"}, vectors={"$id": [0.1, 0.2]} + ) + row = doc_to_row(doc, schema=schema, include_vector=True) + assert row == {"$$id": "PK-001", "id": "USER-9", "title": "t", "$id": [0.1, 0.2]} + + def test_split_pk_separates_chain_pk_from_dollar_id_vector(self) -> None: + schema = _schema_with_dollar_id_vector(with_id_field=True) + row = {"$$id": "PK-001", "id": "USER-9", "$id": [0.1, 0.2], "title": "t"} + found, rest = split_pk(row, schema=schema) + assert found == "PK-001" + assert rest == {"id": "USER-9", "$id": [0.1, 0.2], "title": "t"} + + def test_bare_id_still_ambiguous_when_id_column_exists(self) -> None: + """Even with a ``$id`` vector around, a bare ``id`` cannot be + resolved when the schema declares an ``id`` field.""" + schema = _schema_with_dollar_id_vector(with_id_field=True) + with pytest.raises(InvalidSchemaError) as exc: + split_pk({"id": "USER-9", "$id": [0.1, 0.2], "title": "t"}, schema=schema) + assert exc.value.extra["primaryKeyKey"] == "$$id" + + def test_clashing_export_stays_importable_into_plain_schema(self) -> None: + """A row carrying ``$id`` from an id-column export still resolves + into a plain collection.""" + found, rest = split_pk( + {"$id": "PK-009", "title": "t"}, schema=_schema(with_id_field=False) + ) + assert found == "PK-009" + assert rest == {"title": "t"} + + +class TestRoundTrip: + @pytest.mark.parametrize("with_id_field", [False, True]) + def test_doc_to_row_then_split_pk(self, with_id_field: bool) -> None: + fields: dict[str, object] = {"title": "t"} + if with_id_field: + fields["id"] = "USER-9" + doc = _FakeDoc("PK-1", fields=fields, vectors={"embedding": [0.1, 0.2]}) + schema = _schema(with_id_field=with_id_field) + + row = doc_to_row(doc, schema=schema, include_vector=True) + found, rest = split_pk(row, schema=schema) + + assert found == "PK-1" + assert rest == {**fields, "embedding": [0.1, 0.2]} diff --git a/apps/backend/tests/unit/test_export_formats.py b/apps/backend/tests/unit/test_export_formats.py new file mode 100644 index 0000000..d4ce7f2 --- /dev/null +++ b/apps/backend/tests/unit/test_export_formats.py @@ -0,0 +1,97 @@ +"""Unit tests for the export serialization layer (``formats.JsonlFormat``). + +Design doc §6.6 rules under test: +* one JSON object per line, ``\\n`` terminated (streaming-friendly); +* sparse vector int keys serialised as JSON string keys + (``orjson.OPT_NON_STR_KEYS`` — without it orjson raises); +* non-finite floats are rejected up front with the offending location + (silent ``null`` coercion by orjson would corrupt data); +* nullable-absent semantics pass through untouched. +""" + +from __future__ import annotations + +import pytest + +from zvec_studio.exceptions import ExportNonFiniteError +from zvec_studio.storage.formats import JsonlFormat + + +def _lines(payload: bytes) -> list[str]: + return payload.decode("utf-8").splitlines() + + +class TestJsonlSerialize: + def test_one_object_per_line_newline_terminated(self) -> None: + rows = [{"id": "a", "title": "x"}, {"id": "b", "title": "y"}] + chunks = list(JsonlFormat().serialize(iter(rows))) + + assert len(chunks) == 2 + assert all(c.endswith(b"\n") for c in chunks) + assert _lines(chunks[0] + chunks[1]) == [ + '{"id":"a","title":"x"}', + '{"id":"b","title":"y"}', + ] + + def test_sparse_int_keys_become_string_keys(self) -> None: + """Without OPT_NON_STR_KEYS orjson raises on int dict keys.""" + rows = [{"id": "a", "keywords": {42: 1.0, 314: 0.5}}] + payload = b"".join(JsonlFormat().serialize(iter(rows))) + + assert _lines(payload) == ['{"id":"a","keywords":{"42":1.0,"314":0.5}}'] + + def test_reserved_pk_key_passes_through(self) -> None: + rows = [{"$id": "PK-1", "id": "USER-1", "title": "t"}] + payload = b"".join(JsonlFormat().serialize(iter(rows))) + + assert _lines(payload) == ['{"$id":"PK-1","id":"USER-1","title":"t"}'] + + def test_unicode_is_not_escaped(self) -> None: + rows = [{"id": "a", "title": "héllo 中文"}] + payload = b"".join(JsonlFormat().serialize(iter(rows))) + + assert "héllo 中文".encode() in payload + + def test_nan_scalar_is_rejected_with_location(self) -> None: + rows = [{"id": "doc-1", "score": float("nan")}] + + with pytest.raises(ExportNonFiniteError) as exc: + b"".join(JsonlFormat().serialize(iter(rows))) + + assert exc.value.status_code == 422 + assert exc.value.code == "EXPORT_NON_FINITE_VALUE" + assert exc.value.extra["documentId"] == "doc-1" + assert exc.value.extra["path"] == "score" + + def test_infinity_in_vector_is_rejected_with_path(self) -> None: + rows = [{"id": "doc-1", "embedding": [0.1, float("inf")]}] + + with pytest.raises(ExportNonFiniteError) as exc: + b"".join(JsonlFormat().serialize(iter(rows))) + + assert exc.value.extra["path"] == "embedding[1]" + + def test_non_finite_in_nested_array_field_is_rejected(self) -> None: + rows = [{"id": "doc-1", "matrix": [[1.0], [float("-inf")]]}] + + with pytest.raises(ExportNonFiniteError) as exc: + b"".join(JsonlFormat().serialize(iter(rows))) + + assert exc.value.extra["path"] == "matrix[1][0]" + + def test_finite_rows_before_the_bad_one_are_emitted(self) -> None: + """The generator yields lazily: good rows stream out before the raise.""" + rows = [{"id": "ok"}, {"id": "bad", "score": float("nan")}] + gen = JsonlFormat().serialize(iter(rows)) + + assert next(gen) == b'{"id":"ok"}\n' + with pytest.raises(ExportNonFiniteError): + next(gen) + + def test_document_id_falls_back_to_dollar_id(self) -> None: + rows = [{"$id": "PK-9", "id": "USER-9", "score": float("nan")}] + + with pytest.raises(ExportNonFiniteError) as exc: + b"".join(JsonlFormat().serialize(iter(rows))) + + assert exc.value.extra["documentId"] == "PK-9" diff --git a/apps/backend/tests/unit/test_import_batch_abort.py b/apps/backend/tests/unit/test_import_batch_abort.py new file mode 100644 index 0000000..6ba42b5 --- /dev/null +++ b/apps/backend/tests/unit/test_import_batch_abort.py @@ -0,0 +1,116 @@ +"""Unit tests for the abort-mode batch fallback in the import pipeline. + +When the SDK rejects a whole batch with a ``ValueError`` (batch-level +validation), ``abort`` must behave like the other abort paths: rows that +were already validated inside the same batch still land, and the import +stops at the first failing row — not silently discard everything. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import zvec + +from zvec_studio.schemas.collection import ( + CollectionSchema, + VectorDataType, + VectorSchema, +) +from zvec_studio.storage.import_ import ImportMode, ImportReport, OnErrorMode +from zvec_studio.storage.sdk import CollectionRecord, SdkBackend + + +class _FakeStatus: + def __init__(self, ok: bool) -> None: + self._ok = ok + + def ok(self) -> bool: + return self._ok + + +class _BatchRejectingSdk: + """Rejects multi-doc writes with ValueError; single-doc writes succeed + unless the id is the planted bad row.""" + + def __init__(self, bad_id: str) -> None: + self.bad_id = bad_id + self.written: list[str] = [] + + def insert(self, docs: list) -> list[_FakeStatus]: + if len(docs) > 1: + raise ValueError(f"Invalid doc[{self.bad_id}]: batch-level rejection") + doc = docs[0] + if doc.id == self.bad_id: + raise ValueError(f"Invalid doc[{self.bad_id}]: bad row") + self.written.append(doc.id) + return [_FakeStatus(True)] + + upsert = insert + + +def _record(sdk: _BatchRejectingSdk) -> CollectionRecord: + schema = CollectionSchema.model_construct( + name="target", + vectors=[ + VectorSchema(name="embedding", dataType=VectorDataType.VECTOR_FP32, dimension=2) + ], + fields=[], + ) + return CollectionRecord( + name="target", path=Path("/tmp/target"), schema=schema, sdk_obj=sdk + ) + + +def _docs(*ids: str) -> list: + return [zvec.Doc(id=i, vectors={"embedding": [0.1, 0.2]}) for i in ids] + + +class TestAbortBatchFallback: + def test_abort_keeps_prior_rows_on_batch_value_error(self) -> None: + sdk = _BatchRejectingSdk(bad_id="doc-1") + backend = SdkBackend() + report = ImportReport() + + keep_going = backend._import_write_batch( + _record(sdk), + _docs("doc-0", "doc-1", "doc-2"), + [1, 2, 3], + mode=ImportMode.INSERT, + on_error=OnErrorMode.ABORT, + report=report, + ) + + assert keep_going is False # abort + # The valid row before the failing one still landed. + assert sdk.written == ["doc-0"] + assert report.imported == 1 + assert report.failed == 1 + # Exact line of the offending row (not best-effort). + assert report.errors[0].line == 2 + assert report.errors[0].code == "INVALID_DOCUMENT" + + def test_skip_still_locates_every_offending_row(self) -> None: + """Regression guard: the pre-existing skip behaviour is unchanged.""" + sdk = _BatchRejectingSdk(bad_id="doc-1") + backend = SdkBackend() + report = ImportReport() + + keep_going = backend._import_write_batch( + _record(sdk), + _docs("doc-0", "doc-1", "doc-2"), + [1, 2, 3], + mode=ImportMode.INSERT, + on_error=OnErrorMode.SKIP, + report=report, + ) + + assert keep_going is True + assert sdk.written == ["doc-0", "doc-2"] + assert report.imported == 2 + assert [(e.line, e.code) for e in report.errors] == [(2, "INVALID_DOCUMENT")] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-q"])) diff --git a/apps/backend/tests/unit/test_import_formats.py b/apps/backend/tests/unit/test_import_formats.py new file mode 100644 index 0000000..e6506b2 --- /dev/null +++ b/apps/backend/tests/unit/test_import_formats.py @@ -0,0 +1,98 @@ +"""Unit tests for the import-format layer (``storage/formats.py``). + +Per the design doc (§4.2), formats are the only variable piece of the +import pipeline: ``ImportFormat`` turns a byte stream into ``(line number, +record)`` pairs, and a registry resolves formats by name / file extension. +First version ships JSONL only — the registry exists so adding CSV/Parquet +later is registration, not surgery. +""" + +from __future__ import annotations + +import io + +import pytest + +from zvec_studio.exceptions import InvalidDocumentError +from zvec_studio.storage.formats import ( + EXPORT_FORMATS, + IMPORT_FORMATS, + JsonlFormat, + resolve_import_format, +) + + +def _stream(text: str) -> io.BufferedReader: + return io.BufferedReader(io.BytesIO(text.encode("utf-8"))) + + +class TestJsonlParsing: + def test_parses_valid_rows_with_line_numbers(self) -> None: + stream = _stream('{"id": "a"}\n{"id": "b"}\n') + rows = list(JsonlFormat().parse(stream)) + assert rows == [(1, {"id": "a"}), (2, {"id": "b"})] + + def test_skips_blank_lines_but_keeps_numbering(self) -> None: + stream = _stream('{"id": "a"}\n\n \n{"id": "b"}\n') + rows = list(JsonlFormat().parse(stream)) + assert rows == [(1, {"id": "a"}), (4, {"id": "b"})] + + def test_accepts_missing_trailing_newline(self) -> None: + stream = _stream('{"id": "a"}') + assert list(JsonlFormat().parse(stream)) == [(1, {"id": "a"})] + + def test_invalid_json_line_reports_line_number(self) -> None: + stream = _stream('{"id": "a"}\nnot json\n{"id": "b"}\n') + parsed = JsonlFormat().parse(stream) + assert next(parsed) == (1, {"id": "a"}) + with pytest.raises(InvalidDocumentError) as exc: + next(parsed) + assert exc.value.status_code == 422 + assert exc.value.code == "INVALID_DOCUMENT" + assert exc.value.extra["line"] == 2 + + def test_resumes_after_a_row_error(self) -> None: + """``skip``-mode imports keep going past failing lines, so the + iterator must survive its own raises (a generator would not).""" + stream = _stream('{"id": "a"}\nnot json\n{"id": "b"}\n') + parsed = JsonlFormat().parse(stream) + assert next(parsed) == (1, {"id": "a"}) + with pytest.raises(InvalidDocumentError): + next(parsed) + assert next(parsed) == (3, {"id": "b"}) + with pytest.raises(StopIteration): + next(parsed) + + def test_line_must_be_a_json_object(self) -> None: + for body in ('[1, 2]', '"plain string"', "42", "null"): + stream = _stream(body + "\n") + with pytest.raises(InvalidDocumentError) as exc: + list(JsonlFormat().parse(stream)) + assert exc.value.extra["line"] == 1 + + def test_empty_stream_yields_nothing(self) -> None: + assert list(JsonlFormat().parse(_stream(""))) == [] + assert list(JsonlFormat().parse(_stream("\n\n"))) == [] + + +class TestFormatRegistry: + def test_jsonl_registered_under_name_and_extensions(self) -> None: + assert "jsonl" in IMPORT_FORMATS + assert EXPORT_FORMATS["jsonl"].extension == "jsonl" + assert "ndjson" in JsonlFormat().extensions + + def test_resolve_by_name(self) -> None: + assert resolve_import_format("jsonl") is IMPORT_FORMATS["jsonl"] + + def test_resolve_by_extension(self) -> None: + assert resolve_import_format("ndjson") is IMPORT_FORMATS["jsonl"] + + def test_resolve_is_case_insensitive(self) -> None: + assert resolve_import_format("JSONL") is IMPORT_FORMATS["jsonl"] + + def test_resolve_unknown_lists_supported(self) -> None: + with pytest.raises(InvalidDocumentError) as exc: + resolve_import_format("csv") + assert exc.value.status_code == 400 + assert exc.value.code == "IMPORT_UNSUPPORTED_FORMAT" + assert "jsonl" in exc.value.message diff --git a/apps/backend/tests/unit/test_sdk_query_params.py b/apps/backend/tests/unit/test_sdk_query_params.py index a519586..33d1c70 100644 --- a/apps/backend/tests/unit/test_sdk_query_params.py +++ b/apps/backend/tests/unit/test_sdk_query_params.py @@ -148,6 +148,24 @@ def test_rotated_quantizer_rejects_unsupported_combinations(spec: dict) -> None: _build_index_param(VectorIndexParam.model_validate(spec)) +def test_unknown_vector_dtype_raises_dedicated_error_not_404() -> None: + """A collection with a dtype Studio does not represent (e.g. legacy + VECTOR_FP64) exists — the error must say so, not claim 404 not-found.""" + from zvec_studio.exceptions import UnsupportedVectorDataTypeError + + sdk_schema = zvec.CollectionSchema( + name="demo", + vectors=[zvec.VectorSchema("embedding", zvec.DataType.VECTOR_FP64, 4)], + fields=None, + ) + + with pytest.raises(UnsupportedVectorDataTypeError) as exc: + _from_sdk_schema(sdk_schema, Path("/tmp/demo")) + assert exc.value.status_code == 422 + assert exc.value.code == "UNSUPPORTED_VECTOR_DATA_TYPE" + assert exc.value.extra["vector"] == "embedding" + + def test_sdk_schema_with_pybind_enum_index_param_serializes_for_summary() -> None: sdk_schema = zvec.CollectionSchema( name="demo", diff --git a/apps/backend/tests/unit/test_snapshot.py b/apps/backend/tests/unit/test_snapshot.py new file mode 100644 index 0000000..ac2c7c5 --- /dev/null +++ b/apps/backend/tests/unit/test_snapshot.py @@ -0,0 +1,230 @@ +"""Unit tests for snapshot packaging (``storage/snapshot.py``). + +Design doc §6.2.2: a snapshot is ``manifest.json`` + ``documents.jsonl``. +The manifest carries the schema and export options so an import can +*pre-check compatibility before writing a single row* (fail-fast, 409). +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from zvec_studio.exceptions import ( + ImportManifestInvalidError, + ImportSchemaMismatchError, +) +from zvec_studio.schemas.collection import CollectionSchema +from zvec_studio.storage.snapshot import ( + DATA_FILE_NAME, + MANIFEST_FORMAT, + MANIFEST_NAME, + build_manifest, + check_schema_compatible, + parse_manifest, +) + + +def _schema( + name: str = "demo", + *, + dim: int = 4, + fields: list[dict] | None = None, +) -> CollectionSchema: + return CollectionSchema.model_validate( + { + "name": name, + "vectors": [ + { + "name": "embedding", + "dataType": "VECTOR_FP32", + "dimension": dim, + "indexParam": {"indexType": "FLAT", "metric": "L2"}, + } + ], + "fields": fields if fields is not None else [{"name": "title", "dataType": "STRING"}], + } + ) + + +class TestBuildManifest: + def test_structure(self) -> None: + manifest = build_manifest( + schema=_schema(), include_vector=True, output_fields=None + ) + + assert manifest["format"] == MANIFEST_FORMAT == "zvec-studio.export/1" + assert manifest["zvecVersion"] + assert manifest["studioVersion"] + assert manifest["exportedAt"] + assert manifest["collection"]["name"] == "demo" + assert manifest["options"] == {"includeVector": True, "outputFields": None} + assert manifest["data"]["file"] == DATA_FILE_NAME == "documents.jsonl" + assert manifest["data"]["format"] == "jsonl" + + def test_schema_roundtrips_through_the_model(self) -> None: + manifest = build_manifest( + schema=_schema(fields=[{"name": "score", "dataType": "INT64"}]), + include_vector=False, + output_fields=["score"], + ) + + restored = CollectionSchema.model_validate(manifest["collection"]["schema"]) + assert [f.name for f in restored.fields] == ["score"] + assert restored.vectors[0].dimension == 4 + + def test_no_doc_count_is_recorded(self) -> None: + """Streaming generation cannot know the count up front (design §6.2.2).""" + manifest = build_manifest(schema=_schema(), include_vector=True, output_fields=None) + assert "docCount" not in manifest["data"] + + +class TestParseManifest: + def test_roundtrip(self) -> None: + manifest = build_manifest(schema=_schema(), include_vector=True, output_fields=None) + import orjson + + parsed = parse_manifest(orjson.dumps(manifest)) + assert parsed["format"] == MANIFEST_FORMAT + + def test_invalid_json(self) -> None: + with pytest.raises(ImportManifestInvalidError): + parse_manifest(b"not json") + + def test_not_an_object(self) -> None: + with pytest.raises(ImportManifestInvalidError): + parse_manifest(b"[1, 2]") + + def test_unknown_major_version(self) -> None: + manifest = build_manifest(schema=_schema(), include_vector=True, output_fields=None) + manifest["format"] = "zvec-studio.export/2" + import orjson + + with pytest.raises(ImportManifestInvalidError) as exc: + parse_manifest(orjson.dumps(manifest)) + assert exc.value.status_code == 400 + + def test_missing_schema(self) -> None: + manifest = build_manifest(schema=_schema(), include_vector=True, output_fields=None) + del manifest["collection"]["schema"] + import orjson + + with pytest.raises(ImportManifestInvalidError): + parse_manifest(orjson.dumps(manifest)) + + +class TestSchemaCompatibility: + def _manifest( + self, + schema: CollectionSchema, + *, + include_vector: bool = True, + output_fields: list[str] | None = None, + ) -> dict: + return build_manifest( + schema=schema, include_vector=include_vector, output_fields=output_fields + ) + + def test_identical_schemas_pass(self) -> None: + check_schema_compatible(self._manifest(_schema()), _schema()) + + def test_target_with_extra_fields_passes(self) -> None: + """A pruned export may carry fewer fields than the target has.""" + check_schema_compatible( + self._manifest(_schema(fields=[{"name": "title", "dataType": "STRING"}])), + _schema(fields=[{"name": "title", "dataType": "STRING"}, {"name": "extra", "dataType": "INT64"}]), + ) + + def test_pruned_field_is_not_required_on_target(self) -> None: + """Exported with ``outputFields=["title"]``: the dropped ``body`` + column never appears in the data, so a target without it is fine.""" + source = _schema( + fields=[ + {"name": "title", "dataType": "STRING"}, + {"name": "body", "dataType": "STRING"}, + ] + ) + manifest = self._manifest(source, output_fields=["title"]) + check_schema_compatible(manifest, _schema()) + + def test_kept_field_is_still_checked_when_pruned(self) -> None: + """outputFields pruning exempts only the dropped columns.""" + source = _schema( + fields=[ + {"name": "title", "dataType": "STRING"}, + {"name": "body", "dataType": "STRING"}, + ] + ) + manifest = self._manifest(source, output_fields=["title"]) + target = _schema(fields=[{"name": "title", "dataType": "INT64"}]) + with pytest.raises(ImportSchemaMismatchError) as exc: + check_schema_compatible(manifest, target) + assert any("title" in m for m in exc.value.extra["mismatches"]) + + def test_vector_dimension_mismatch_rejected(self) -> None: + with pytest.raises(ImportSchemaMismatchError) as exc: + check_schema_compatible(self._manifest(_schema(dim=8)), _schema(dim=4)) + assert exc.value.status_code == 409 + assert exc.value.extra["mismatches"] + + def test_missing_vector_in_target_rejected(self) -> None: + target = CollectionSchema.model_validate( + {"name": "tgt", "vectors": [], "fields": [{"name": "title", "dataType": "STRING"}]} + ) + with pytest.raises(ImportSchemaMismatchError): + check_schema_compatible(self._manifest(_schema()), target) + + def test_field_type_mismatch_rejected(self) -> None: + with pytest.raises(ImportSchemaMismatchError): + check_schema_compatible( + self._manifest(_schema(fields=[{"name": "title", "dataType": "STRING"}])), + _schema(fields=[{"name": "title", "dataType": "INT64"}]), + ) + + def test_field_absent_in_target_rejected(self) -> None: + with pytest.raises(ImportSchemaMismatchError): + check_schema_compatible( + self._manifest(_schema(fields=[{"name": "title", "dataType": "STRING"}])), + _schema(fields=[{"name": "other", "dataType": "STRING"}]), + ) + + def test_vectors_not_exported_but_required_rejected(self) -> None: + """includeVector=false cannot repopulate a collection that needs vectors.""" + with pytest.raises(ImportSchemaMismatchError) as exc: + check_schema_compatible( + self._manifest(_schema(), include_vector=False), _schema() + ) + assert any("vector" in m.lower() for m in exc.value.extra["mismatches"]) + + +class TestStreamCancellation: + def test_early_close_unblocks_the_producer_thread(self, tmp_path: Path) -> None: + """A client disconnect (generator closed mid-stream) must not leave + the producer thread parked on the bounded queue forever.""" + import os + import threading + import time + + from zvec_studio.storage.snapshot import ( + SNAPSHOT_WORKER_NAME, + stream_snapshot_package, + ) + + # Incompressible payload so the producer out-runs a one-chunk consumer + # and blocks on the bounded queue. + (tmp_path / DATA_FILE_NAME).write_bytes(os.urandom(4 * 1024 * 1024)) + (tmp_path / MANIFEST_NAME).write_bytes(b"{}") + + stream = stream_snapshot_package( + manifest_path=tmp_path / MANIFEST_NAME, data_path=tmp_path / DATA_FILE_NAME + ) + next(stream) # read a single chunk, then abandon like a disconnect + stream.close() + + deadline = time.monotonic() + 5.0 + while time.monotonic() < deadline: + if not any(t.name == SNAPSHOT_WORKER_NAME for t in threading.enumerate()): + break + time.sleep(0.05) + assert not any(t.name == SNAPSHOT_WORKER_NAME for t in threading.enumerate()) diff --git a/apps/backend/tests/unit/test_write_batches.py b/apps/backend/tests/unit/test_write_batches.py new file mode 100644 index 0000000..fa5d602 --- /dev/null +++ b/apps/backend/tests/unit/test_write_batches.py @@ -0,0 +1,152 @@ +"""Unit tests for the write-path batching helper. + +The Zvec SDK rejects write batches larger than ``_MAX_SDK_WRITE_BATCH`` (1024, +a runtime error, not a documented contract), while Studio's HTTP surface +accepts up to 10,000 documents. ``_write_in_batches`` absorbs that constraint +internally so callers never hit it. +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from zvec_studio.exceptions import ( + DocumentConflictError, + DocumentNotFoundError, + InvalidDocumentError, + ZvecStudioError, +) +from zvec_studio.storage.sdk import ( + _DEFAULT_WRITE_BATCH, + _MAX_SDK_WRITE_BATCH, + _write_in_batches, +) + + +class _FakeStatus: + def __init__(self, ok: bool, message: str = "") -> None: + self._ok = ok + self._message = message + + def ok(self) -> bool: + return self._ok + + def message(self) -> str: + return self._message + + +class TestBatchSizing: + @pytest.mark.parametrize( + ("total", "expected_chunk_sizes"), + [ + (0, []), + (1, [1]), + (511, [511]), + (512, [512]), + (513, [512, 1]), + (1024, [512, 512]), + (1025, [512, 512, 1]), + (2049, [512, 512, 512, 512, 1]), + ], + ) + def test_batches_use_the_default_size( + self, total: int, expected_chunk_sizes: list[int] + ) -> None: + docs = list(range(total)) + seen: list[int] = [] + + def write(chunk: list[int]) -> list[_FakeStatus]: + seen.append(len(chunk)) + assert len(chunk) <= _MAX_SDK_WRITE_BATCH + return [_FakeStatus(True)] * len(chunk) + + _write_in_batches(docs, write, failed_code="INSERT_FAILED") + + assert seen == expected_chunk_sizes + + def test_explicit_batch_size_is_clamped_to_sdk_limit(self) -> None: + """Callers may ask for larger chunks; the SDK limit still caps them.""" + docs = list(range(3000)) + seen: list[int] = [] + + def write(chunk: list[int]) -> list[_FakeStatus]: + seen.append(len(chunk)) + assert len(chunk) <= _MAX_SDK_WRITE_BATCH + return [_FakeStatus(True)] * len(chunk) + + _write_in_batches(docs, write, failed_code="INSERT_FAILED", batch_size=5000) + + assert seen == [1024, 1024, 952] + + def test_all_statuses_are_preserved_in_order(self) -> None: + docs = list(range(_MAX_SDK_WRITE_BATCH + 1)) + + def write(chunk: list[int]) -> list[_FakeStatus]: + return [_FakeStatus(True)] * len(chunk) + + statuses = _write_in_batches(docs, write, failed_code="INSERT_FAILED") + assert len(statuses) == _MAX_SDK_WRITE_BATCH + 1 + + +class TestErrorClassification: + def test_sdk_validation_value_error_maps_to_invalid_document(self) -> None: + def write(chunk: list[Any]) -> list[Any]: + raise ValueError( + "Invalid doc[0]: field[score] is required but not provided" + ) + + with pytest.raises(InvalidDocumentError) as exc: + _write_in_batches([0], write, failed_code="INSERT_FAILED") + + assert exc.value.status_code == 422 + assert "score" in exc.value.message + + def test_batch_limit_value_error_is_not_masked_as_user_input(self) -> None: + """If batching itself regressed, the error must surface as internal.""" + + def write(chunk: list[Any]) -> list[Any]: + raise ValueError("Too many docs: 2000 exceeds max write batch size") + + with pytest.raises(ZvecStudioError) as exc: + _write_in_batches([0], write, failed_code="INSERT_FAILED") + + assert not isinstance(exc.value, InvalidDocumentError) + assert exc.value.status_code == 500 + + def test_duplicate_id_status_maps_to_conflict(self) -> None: + def write(chunk: list[Any]) -> list[_FakeStatus]: + return [ + _FakeStatus(True), + _FakeStatus(False, "insert failed: doc_id[x] already exists in collection"), + ] + + with pytest.raises(DocumentConflictError) as exc: + _write_in_batches([0, 1], write, failed_code="INSERT_FAILED") + + assert exc.value.status_code == 409 + + def test_unknown_id_status_maps_to_not_found(self) -> None: + def write(chunk: list[Any]) -> list[_FakeStatus]: + return [_FakeStatus(False, "Document not found")] + + with pytest.raises(DocumentNotFoundError) as exc: + _write_in_batches([0], write, failed_code="UPDATE_FAILED") + + assert exc.value.status_code == 404 + + def test_other_failed_status_keeps_caller_code(self) -> None: + def write(chunk: list[Any]) -> list[_FakeStatus]: + return [_FakeStatus(False, "some opaque failure")] + + with pytest.raises(ZvecStudioError) as exc: + _write_in_batches([0], write, failed_code="UPSERT_FAILED") + + assert exc.value.code == "UPSERT_FAILED" + assert exc.value.status_code == 500 + + +class TestDefaults: + def test_default_batch_is_below_sdk_limit(self) -> None: + assert 0 < _DEFAULT_WRITE_BATCH <= _MAX_SDK_WRITE_BATCH diff --git a/apps/backend/uv.lock b/apps/backend/uv.lock index fc5015b..c563f9b 100644 --- a/apps/backend/uv.lock +++ b/apps/backend/uv.lock @@ -656,34 +656,34 @@ wheels = [ [package.optional-dependencies] cudart = [ - { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cuda-runtime" }, ] cufft = [ - { name = "nvidia-cufft", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cufft" }, ] cufile = [ - { name = "nvidia-cufile", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cufile" }, ] cupti = [ - { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cuda-cupti" }, ] curand = [ - { name = "nvidia-curand", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-curand" }, ] cusolver = [ - { name = "nvidia-cusolver", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cusolver" }, ] cusparse = [ - { name = "nvidia-cusparse", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cusparse" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-nvjitlink" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cuda-nvrtc" }, ] nvtx = [ - { name = "nvidia-nvtx", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-nvtx" }, ] [[package]] @@ -738,7 +738,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -3186,10 +3186,10 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "joblib", marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "threadpoolctl", marker = "python_full_version < '3.11'" }, + { name = "joblib" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" } }, + { name = "threadpoolctl" }, ] sdist = { url = "https://files.pythonhosted.org/packages/98/c2/a7855e41c9d285dfe86dc50b250978105dce513d6e459ea66a6aeb0e1e0c/scikit_learn-1.7.2.tar.gz", hash = "sha256:20e9e49ecd130598f1ca38a1d85090e1a600147b9c02fa6f15d69cb53d968fda", size = 7193136, upload-time = "2025-09-09T08:21:29.075Z" } wheels = [ @@ -3234,10 +3234,10 @@ resolution-markers = [ "python_full_version >= '3.11' and python_full_version < '3.15'", ] dependencies = [ - { name = "joblib", marker = "python_full_version >= '3.11'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "threadpoolctl", marker = "python_full_version >= '3.11'" }, + { name = "joblib" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" } }, + { name = "threadpoolctl" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0e/d4/40988bf3b8e34feec1d0e6a051446b1f66225f8529b9309becaeef62b6c4/scikit_learn-1.8.0.tar.gz", hash = "sha256:9bccbb3b40e3de10351f8f5068e105d0f4083b1a65fa07b6634fbc401a6287fd", size = 7335585, upload-time = "2025-12-10T07:08:53.618Z" } wheels = [ @@ -3287,7 +3287,7 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } wheels = [ @@ -3347,7 +3347,7 @@ resolution-markers = [ "python_full_version >= '3.11' and python_full_version < '3.15'", ] dependencies = [ - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } wheels = [ @@ -4188,33 +4188,43 @@ wheels = [ [[package]] name = "zvec" -version = "0.6.0" +version = "0.7.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/7f/df/3be46b99573b8fbc232e2ce2c86a6e76155727f65f3a378251a31b930f07/zvec-0.6.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:c9040d309e5e29b5ce442238215b5d5668075afbf6f6128b8b4a8518f0d325d2", size = 12330837, upload-time = "2026-07-20T03:44:17.325Z" }, - { url = "https://files.pythonhosted.org/packages/2f/3a/c160650dc876743e05826cb36471e12b5d522f92952b96e595f119fd1cdc/zvec-0.6.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:5cc4eb12a1665fd0726b9f4e4db9b712ff0fc3e34dc6a745e5143cb2b338096f", size = 17434899, upload-time = "2026-07-20T03:39:45.255Z" }, - { url = "https://files.pythonhosted.org/packages/96/80/edb432349b6da09faba15335539344e650b93809d4e520743e5c97163894/zvec-0.6.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:cd96ea6720ff156756fe13f2398bed3b75577505fabbb25dfdc2b016c44ef4a0", size = 19718007, upload-time = "2026-07-20T03:46:13.115Z" }, - { url = "https://files.pythonhosted.org/packages/07/1f/f843554c669c6029f06538b00d9444ab254592c6e31f353369a2f73d0415/zvec-0.6.0-cp310-cp310-win_amd64.whl", hash = "sha256:72f3080411b6bb75d030bfcbba72dd6b8d00910eb37a49518e66e161e13ec3fd", size = 12948158, upload-time = "2026-07-20T03:30:20.176Z" }, - { url = "https://files.pythonhosted.org/packages/51/a3/1f46239b5b8c1863ae83592d0b62a831f8a8915491d0277c2231cd6a0933/zvec-0.6.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4f0f4d2dbc2552d237c6209b09cc023914b70fc746fae1458511aedd34813383", size = 12331212, upload-time = "2026-07-20T03:44:19.133Z" }, - { url = "https://files.pythonhosted.org/packages/d8/5f/5f3a93461bb18c561913e3cdc4681c650bd05623be2c46522f4ca9429af8/zvec-0.6.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:60e9ba18b97ef2ffaaebbd233e1a91e5bb013ff4a73f54a6ca28dda7f3e4e9e8", size = 17436383, upload-time = "2026-07-20T03:39:48.105Z" }, - { url = "https://files.pythonhosted.org/packages/fa/1b/d7bea9e9cd9f82947a26b7274539922ca66824d844161d69b57225d78e6e/zvec-0.6.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:474de64920fb9d02817e6995f03937a4f3ec3594a0960b4d90624863bf2f3095", size = 19718087, upload-time = "2026-07-20T03:46:16.458Z" }, - { url = "https://files.pythonhosted.org/packages/ae/0b/a91063ea0d6024d81bbfc0f8269bde3158e465baa2603839d2482688a7dd/zvec-0.6.0-cp311-cp311-win_amd64.whl", hash = "sha256:2d168fb8a15f2c6b37a3e9e8f818bab14317f42c5cab33edd628c3ebf8381f34", size = 12950419, upload-time = "2026-07-20T03:30:22.925Z" }, - { url = "https://files.pythonhosted.org/packages/0a/db/67ce3e9cad1ad1d23dbc57880b3148a0cee0c3df66050eda9dd2d29cdc27/zvec-0.6.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:7444b90e0351bbfaefcae601ee5d148926996ce6b67049f31b18f81bc2a27f33", size = 12335004, upload-time = "2026-07-20T03:44:21.548Z" }, - { url = "https://files.pythonhosted.org/packages/6c/74/21124c3b2ee214132118a53d438b95df1ccd9b0f0faf04df71caa6feec5f/zvec-0.6.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:394c568ead1f881ea78b6ee3e8fafdaaa242d34b527449c4a928997c51b78363", size = 17439389, upload-time = "2026-07-20T03:39:50.352Z" }, - { url = "https://files.pythonhosted.org/packages/ff/3e/c31c139867d920fd23af4e374e7fc3dcbed6954ef861a9bd66ac9d040254/zvec-0.6.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:75676b77bf1e280a1822f9911203b5892fcb055b2ff4a0a50fd25ce0334892ce", size = 19727865, upload-time = "2026-07-20T03:46:19.404Z" }, - { url = "https://files.pythonhosted.org/packages/ce/2e/4f4e25f1e647d4d91c5e5a44674d388869c10a90a5de6c40cbd23fd3fabf/zvec-0.6.0-cp312-cp312-win_amd64.whl", hash = "sha256:b846cdbeca48112961f91a7f7971ca6014ff8892b465af8a787700a1b3a0f585", size = 12952987, upload-time = "2026-07-20T03:30:25.741Z" }, - { url = "https://files.pythonhosted.org/packages/ba/2d/ccfa996c829d95f1b4e3465a58fb95bf25b7c6045533a06896573632b9ff/zvec-0.6.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:3d503bbce8275f3efed935c42d67eb771c4b8c277140e934193f56a527597391", size = 12335268, upload-time = "2026-07-20T03:44:23.833Z" }, - { url = "https://files.pythonhosted.org/packages/3c/92/6c3a2a07c09b2cb98ba96f9d14590f2d4c1b5ac93b3a7430d09f95fb997e/zvec-0.6.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5defdf2c6748056af5e7b01c8857c92a3a0d7933316d1c99af83a768e493cdc9", size = 17439157, upload-time = "2026-07-20T03:39:52.564Z" }, - { url = "https://files.pythonhosted.org/packages/b3/6e/ac765097f921e021b109a708d915115cf7ca7097a81dad49b78a0cc174de/zvec-0.6.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:8f3eda20ce95b560e0017619ecfeb8ef9f6d74ef844a3b003958cc0feea419f0", size = 19727944, upload-time = "2026-07-20T03:46:22.677Z" }, - { url = "https://files.pythonhosted.org/packages/cf/90/9e9ad844685d15e5c9018094703ac8502e20d98b20c6902404a32ad6d47b/zvec-0.6.0-cp313-cp313-win_amd64.whl", hash = "sha256:ce301f66f1c10d6df91d234f189239fd5fe07a8157633e143d477a4934fc3992", size = 12953157, upload-time = "2026-07-20T03:30:28.294Z" }, - { url = "https://files.pythonhosted.org/packages/9d/04/d961494ac46fc3041245f10f7f125b6ec5c109b85557a3ceb2a56d00c38e/zvec-0.6.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:44f00c863cdb622e7fb7944abbd81f6d114cedeb6f69a8516504b5358abf3cbe", size = 12336108, upload-time = "2026-07-20T03:44:25.648Z" }, - { url = "https://files.pythonhosted.org/packages/78/0b/d837fafb8bf38e7ce542d396fa213137c1167b643f5f87459b5ef5454875/zvec-0.6.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:e143cb949f9edcbb2c67d5120d43818daf45f1d4abbc8acec7a4e4c5dc5b4130", size = 17441514, upload-time = "2026-07-20T03:39:55.073Z" }, - { url = "https://files.pythonhosted.org/packages/4f/3f/e09523c1e1a586e1b528f7f1fc03cc0db829885075c9da4d9005b8db8b81/zvec-0.6.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:aa91e8bba7f965bc94d60a9fc96aaa2649352cc0c4286890d51a0c1dfbdc87c8", size = 19728221, upload-time = "2026-07-20T03:46:25.954Z" }, - { url = "https://files.pythonhosted.org/packages/31/9e/696c0c26cd56beb13b8cc93dc6f9c3b6fb97c00db0954d8d770e5db71df4/zvec-0.6.0-cp314-cp314-win_amd64.whl", hash = "sha256:22f8a8cf320fdb9692c92b5e5dde6172e2f2ef028b2090876e0023c74d21ff68", size = 13459216, upload-time = "2026-07-20T03:30:31.018Z" }, + { url = "https://files.pythonhosted.org/packages/64/30/40cbb2f598f829a6febacbb92dcfcc23ea5701c7c5dafbcedde4d7f6f37d/zvec-0.7.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8b8ea5d520e34013ac5f20b35d67019ffef408fbb69b14a612201efbbdd4e804", size = 10648442, upload-time = "2026-08-26T04:00:48.433Z" }, + { url = "https://files.pythonhosted.org/packages/10/8f/538fd253e140910e61d030f494e9632e29556d3796fdb81f2ede02776beb/zvec-0.7.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:6c82c84ca3116e7c41d2a5a2c7ec8127e26f58abd3cbd28f33e348bc808c58fc", size = 15145105, upload-time = "2026-08-26T05:04:34.1Z" }, + { url = "https://files.pythonhosted.org/packages/c1/ae/e7331fe950d3a232dd9634fb31d755a4df7561d81a5248784481c7b22168/zvec-0.7.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:1c32348c2f2831e86260e9fade1a1d7442426b03c7cdff40bd99642a3ea61100", size = 16980465, upload-time = "2026-08-26T06:03:44.683Z" }, + { url = "https://files.pythonhosted.org/packages/f2/26/0c487088b02010b0b20cd9d032b94ee7c6742d605838db167d7378732acc/zvec-0.7.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:d958f915b499461db6c48bbbe0e4c10eb2c0756961d22425f58f84d54c0cbc2c", size = 16241407, upload-time = "2026-08-26T05:04:36.595Z" }, + { url = "https://files.pythonhosted.org/packages/9c/fe/af2ffd39d66afa9abcb57075b887c565e6d647062ff3618e267fc0e15ed0/zvec-0.7.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:9a59afcd2185b8026a6e86d5e3a93ee1b5f899d11ca49728cc985a4967b397f3", size = 17806982, upload-time = "2026-08-26T06:03:47.173Z" }, + { url = "https://files.pythonhosted.org/packages/bc/04/a9ae744274fb8ba1542853293d0e96a0124b93b707ad270d06f6947de07e/zvec-0.7.0-cp310-cp310-win_amd64.whl", hash = "sha256:8ac4b1f3e40a56a294c3b43ce79ce6b51c758b0e8313ace94219924c1b863fab", size = 11313033, upload-time = "2026-08-26T03:57:54.678Z" }, + { url = "https://files.pythonhosted.org/packages/57/98/16215d61a136e60d79c670bfb1f0b6c32b69d4f3c239f63b5b311f3b561c/zvec-0.7.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:bdd6b8b033f8ec891b17251f482c50a942d3bdd7f3c7c4f69b7a0ef65570a2d5", size = 10650347, upload-time = "2026-08-26T04:00:50.545Z" }, + { url = "https://files.pythonhosted.org/packages/5b/49/ab9a64c29c32bd21b4348fc63babfabfd53a480d8f59c226b0bdb37feaaa/zvec-0.7.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:1fe6edda5ab478f18db9c68488eb92aaf960928dc2eb7ed025fd589401f010e6", size = 15147570, upload-time = "2026-08-26T05:04:39.027Z" }, + { url = "https://files.pythonhosted.org/packages/2d/87/c97a5ff5a33b7da37475eead1a30ac71524af8b2b91dd45a071dcec259c2/zvec-0.7.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:c1939c0f2bde291b8974275e53aa32f4a50a0d900f61cf31078862cce3aac5e4", size = 16980389, upload-time = "2026-08-26T06:03:49.497Z" }, + { url = "https://files.pythonhosted.org/packages/2e/4f/e6a231b127169635a6349a171e2c44082a9b47b548823cedeec50250853b/zvec-0.7.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:5aaa2058fdea877b3b0b03c544d47b231370406db1d830d0899a5e7b72926d77", size = 16242664, upload-time = "2026-08-26T05:04:41.661Z" }, + { url = "https://files.pythonhosted.org/packages/bc/e4/ffa28e35f6e46e2e66995887da872368d6105b9878a2e917623587ce02ae/zvec-0.7.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f7fd3b908a1685b0cde2f25f1c6e0ab0b3be22142a8221546078902bdd3d391d", size = 17808867, upload-time = "2026-08-26T06:03:52.464Z" }, + { url = "https://files.pythonhosted.org/packages/7a/fc/8e1d15646f349b6b597900e36c6c994c7beee91e5c9d17215b5c13b8e6f9/zvec-0.7.0-cp311-cp311-win_amd64.whl", hash = "sha256:2cae381c9643d57359ff778efb5cbfab4ab2a4d6b36e639e900712ac85362e0e", size = 11314012, upload-time = "2026-08-26T03:57:57.613Z" }, + { url = "https://files.pythonhosted.org/packages/62/93/bcd9048bba4490543306c4442e0edcad17a1121952e996c4d0623d4ec37a/zvec-0.7.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ab0ea8c6ad3a815d366c33c12910ab29330b067fd5d8b298479dbf1a6036cf34", size = 10653019, upload-time = "2026-08-26T04:00:52.778Z" }, + { url = "https://files.pythonhosted.org/packages/eb/e1/1b2f2805268780708d85a6331eb506f0e026b2cd536ee1d3bc30d6c722ed/zvec-0.7.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:b0d956395c3accc191eb6021532b851aa205e3f8b10db857ba5b255fa27c2956", size = 15151647, upload-time = "2026-08-26T05:04:44.346Z" }, + { url = "https://files.pythonhosted.org/packages/ab/91/fc52e0834bebcac5ed0bcf9f07196595074d864d340fbf7ff36a0d3f3332/zvec-0.7.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:89aeb42c9318a7a6a69f705ca92af329a5f9a608d25f947ca46045a506a61982", size = 16987010, upload-time = "2026-08-26T06:03:54.69Z" }, + { url = "https://files.pythonhosted.org/packages/c4/cb/7d685448b6397e8f200416f357a33230ca0985a65ddd20d7dc61aa791b8a/zvec-0.7.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:39b3248583415a5018ebce85a81e3bb2f450988e3c35359fda160ab0ba863808", size = 16247178, upload-time = "2026-08-26T05:04:46.946Z" }, + { url = "https://files.pythonhosted.org/packages/60/dc/396b4f70f393ed858cf74a08157336e1d282e9f4150e8a84079e248d5edf/zvec-0.7.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:b60aac076629320b2d8ea63e3cf44bd3d7d364e63b9006712c04f0dbf8ec4eda", size = 17818035, upload-time = "2026-08-26T06:03:57.293Z" }, + { url = "https://files.pythonhosted.org/packages/41/c0/96c1099659eb4fc00ecbf93d0efa72a35597b886abdea493f2c188e4e810/zvec-0.7.0-cp312-cp312-win_amd64.whl", hash = "sha256:4a29a8064875fc55dd8d5a2066f6524c387f04b168182e2edffe0d175c5d358d", size = 11320749, upload-time = "2026-08-26T03:58:00.295Z" }, + { url = "https://files.pythonhosted.org/packages/c0/44/11e62b9b90908556f3d0b4699b14600637321d4a749e77636d7767803941/zvec-0.7.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:11d7bf36b1db0bbe08270f609990dc007129a96661d9a5ff15c5338f7247c51d", size = 10653127, upload-time = "2026-08-26T04:00:55.055Z" }, + { url = "https://files.pythonhosted.org/packages/74/f4/5aa84ef97f1e4c21106866827c23da4e643a9ec9ca0e859d460f53035091/zvec-0.7.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:e1a06666c16d9caa4012c82223b7d25b02419e8ed73d3596d0f949b6691df22e", size = 15151556, upload-time = "2026-08-26T05:04:49.749Z" }, + { url = "https://files.pythonhosted.org/packages/57/86/e9fdb26a04c0bdcdb357b607b36486b320a86efce9df6b3827d0187a25d3/zvec-0.7.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:1316ef51eb85c11836ad6f309ff2546c0f1237883499d629a2ace7ef92aa8641", size = 16987459, upload-time = "2026-08-26T06:03:59.857Z" }, + { url = "https://files.pythonhosted.org/packages/d5/0e/fd1275c5672dfdc4f98560d55f4671854d058033478946c07a33cd115a86/zvec-0.7.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f0c2106fe60d3ea5e5fd4841e38a190b2e599d929e90aa8a6e6e758f59b7292e", size = 16246636, upload-time = "2026-08-26T05:04:52.363Z" }, + { url = "https://files.pythonhosted.org/packages/bd/dd/959e9f0d31049535a5889249b281406ab0df46d11b078e3907508f47cfa0/zvec-0.7.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:d020728e8b9d019b548a36da68e95e01aaff53ebbfa249bf4253c02896f39a8e", size = 17817652, upload-time = "2026-08-26T06:04:02.166Z" }, + { url = "https://files.pythonhosted.org/packages/4c/30/f1d0334dfda82277ae94c82c215c67a364db00bb003cf048582d5980fbf5/zvec-0.7.0-cp313-cp313-win_amd64.whl", hash = "sha256:d77e8c32f44a21e58c8c8ae66c5fedfd3a378ccd92fe749e251ff6a414257e2f", size = 11320700, upload-time = "2026-08-26T03:58:02.995Z" }, + { url = "https://files.pythonhosted.org/packages/0d/e9/d51c27fea8bd989c23fe3f1d39879abcf2f1123071ea6451e54fc9b6b84f/zvec-0.7.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:7a8baca94d4533ba4a5492ae13626bb44ff347b1c91b086c3731e11dd18aff6f", size = 10654270, upload-time = "2026-08-26T04:00:56.962Z" }, + { url = "https://files.pythonhosted.org/packages/0c/fe/4c2c47d925e58cf5f3622e6fbfd9dde9890edc7376ce4150a8a24afeaa2e/zvec-0.7.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:d3005b6e74527c5ef749c46598b6616701c028d7a7afc324990f5ee3d07871d6", size = 15154171, upload-time = "2026-08-26T05:04:54.925Z" }, + { url = "https://files.pythonhosted.org/packages/12/61/3db280a75e95638f218443d3bfa5438d79bda7b8075d6972d67c61b83097/zvec-0.7.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:41bf392648b9369a5360455fe1a360fa55e3e6cceecc9b5615ba0fdf6545787e", size = 16988594, upload-time = "2026-08-26T06:04:04.357Z" }, + { url = "https://files.pythonhosted.org/packages/ab/f4/a53f15480e7f0f15e8987af7e0fff6299473b1eed993e7f0fec39c6d360e/zvec-0.7.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:0e58182c06ccc55a3b50ca26502c40d7c8e8009693749538acafccc86b79baf0", size = 16252423, upload-time = "2026-08-26T05:04:57.623Z" }, + { url = "https://files.pythonhosted.org/packages/3c/7b/6660218fe288af8fd07075c65a9b896d92cc2aea6150dab0eed68a605bb6/zvec-0.7.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:ac021b603c4c0cbebe8db92a1fee4367c3c74d2cf48cf5874b80576e767648af", size = 17819460, upload-time = "2026-08-26T06:04:06.73Z" }, + { url = "https://files.pythonhosted.org/packages/e1/41/573497f28914f747df84db20067a83aa2c78bdf33e1a3849537d46c95cb1/zvec-0.7.0-cp314-cp314-win_amd64.whl", hash = "sha256:1d1b070e68ffc0670c0833e3f4334ce510e48c9c3a5dd43a624723f3b7b2dbc2", size = 11745529, upload-time = "2026-08-26T03:58:05.774Z" }, ] [[package]] @@ -4282,6 +4292,6 @@ requires-dist = [ { name = "types-setuptools", marker = "extra == 'dev'" }, { name = "typing-extensions", specifier = ">=4.12" }, { name = "uvicorn", extras = ["standard"], specifier = ">=0.30,<0.35" }, - { name = "zvec", specifier = ">=0.6,<0.7" }, + { name = "zvec", specifier = ">=0.7,<0.8" }, ] provides-extras = ["ai", "dev", "packaging"] diff --git a/apps/backend/zvec_studio/exceptions.py b/apps/backend/zvec_studio/exceptions.py index 4219c60..a2847a5 100644 --- a/apps/backend/zvec_studio/exceptions.py +++ b/apps/backend/zvec_studio/exceptions.py @@ -86,6 +86,100 @@ class DocumentNotFoundError(ZvecStudioError): title = "Document Not Found" +class DocumentConflictError(ZvecStudioError): + """The target primary key already exists (``insert`` rejects duplicates).""" + + code = "DOCUMENT_CONFLICT" + status_code = 409 + title = "Document Already Exists" + + +class MaintenanceBlockedError(ZvecStudioError): + """Maintenance was rejected because a snapshot iterator is open. + + Zvec blocks ``flush``/``optimize``/DDL while an export iterator holds the + collection (PR #597 contract). That is a transient conflict — retryable + once the export finishes — not an invalid request. + """ + + code = "MAINTENANCE_BLOCKED" + status_code = 409 + title = "Maintenance Blocked" + + +class UnsupportedVectorDataTypeError(ZvecStudioError): + """The collection carries a vector dtype Studio cannot represent. + + Distinct from ``COLLECTION_NOT_FOUND``: the collection exists and opened + fine at the SDK level — only Studio's dtype surface does not cover it + (e.g. a legacy ``VECTOR_FP64`` collection after the enum was pruned). + """ + + code = "UNSUPPORTED_VECTOR_DATA_TYPE" + status_code = 422 + title = "Unsupported Vector Data Type" + + +class ImportFileNotFoundError(ZvecStudioError): + """The requested import source does not exist or is not a regular file.""" + + code = "IMPORT_FILE_NOT_FOUND" + status_code = 404 + title = "Import File Not Found" + + +class ImportFileNotReadableError(ZvecStudioError): + """The import source exists but cannot be opened for reading.""" + + code = "IMPORT_FILE_NOT_READABLE" + status_code = 403 + title = "Import File Not Readable" + + +class ImportManifestInvalidError(ZvecStudioError): + """A snapshot's manifest.json is unreadable or structurally invalid.""" + + code = "IMPORT_MANIFEST_INVALID" + status_code = 400 + title = "Invalid Import Manifest" + + +class ImportSchemaMismatchError(ZvecStudioError): + """The snapshot's schema is not compatible with the target collection. + + Raised before any row is written (fail-fast); ``extra['mismatches']`` + carries human-readable reasons. + """ + + code = "IMPORT_SCHEMA_MISMATCH" + status_code = 409 + title = "Import Schema Mismatch" + + +class ExportNonFiniteError(ZvecStudioError): + """A document carries NaN/±Inf and cannot be written to JSON safely. + + Raised instead of letting orjson silently coerce non-finite floats to + ``null`` (data loss) or emitting non-standard ``NaN`` literals that no + conforming parser reads back. ``extra`` carries ``documentId`` and the + value ``path`` (e.g. ``embedding[3]``). + """ + + code = "EXPORT_NON_FINITE_VALUE" + status_code = 422 + title = "Non-Finite Value In Document" + + +class InvalidDocumentError(ZvecStudioError): + """A single document fails validation at the SDK level (missing required + field, type mismatch, ...). Distinct from ``InvalidSchemaError``, which + covers the collection schema itself.""" + + code = "INVALID_DOCUMENT" + status_code = 422 + title = "Invalid Document" + + class DimensionMismatchError(ZvecStudioError): code = "DIMENSION_MISMATCH" status_code = 400 diff --git a/apps/backend/zvec_studio/routers/collections.py b/apps/backend/zvec_studio/routers/collections.py index 06ee1d1..5283577 100644 --- a/apps/backend/zvec_studio/routers/collections.py +++ b/apps/backend/zvec_studio/routers/collections.py @@ -2,14 +2,20 @@ from __future__ import annotations +import contextlib +import shutil from datetime import datetime, timezone from pathlib import Path from fastapi import APIRouter, Depends, Query, Request, status from zvec_studio.config_store import ConfigStore +from zvec_studio.exceptions import CollectionAlreadyExistsError +from zvec_studio.routers.documents import import_report_to_response from zvec_studio.schemas import ( CollectionCreateRequest, + CollectionImportRequest, + CollectionImportResponse, CollectionListItem, CollectionListResponse, CollectionOpenRequest, @@ -28,6 +34,12 @@ VectorIndexParam, ) from zvec_studio.storage import CollectionRecord, SdkBackend +from zvec_studio.storage.formats import JsonlFormat +from zvec_studio.storage.import_ import validate_import_source +from zvec_studio.storage.snapshot import ( + check_schema_compatible, + read_snapshot_manifest, +) router = APIRouter(prefix="/collections", tags=["collections"]) @@ -106,6 +118,77 @@ def open_collection( return _summary(record, backend) +@router.post( + ":import", + status_code=status.HTTP_201_CREATED, + response_model=CollectionImportResponse, +) +def import_collection( + body: CollectionImportRequest, + backend: SdkBackend = Depends(_get_backend), + store: ConfigStore = Depends(_get_config_store), +) -> CollectionImportResponse: + """Import a collection from a snapshot package (``.tar.gz``). + + Collection-level lifecycle operation (a sibling of create/open): the + manifest embedded in the snapshot supplies the schema, ``targetPath`` + gives the new collection its home, and the embedded JSONL is loaded in + the same pass. Everything is validated before anything touches the disk + (missing/corrupt package, occupied target, name clash, incompatible + options); if the data load itself fails unexpectedly, the freshly + created collection is rolled back. Row-level failures follow the + standard import report semantics and keep the loaded rows. + """ + source = _normalise_path(body.source.path) + validate_import_source(source) # 404 / 403, before touching anything + manifest = read_snapshot_manifest(source) # 400 on corrupt / malformed + + schema = CollectionSchema.model_validate(manifest["collection"]["schema"]) + if body.name is not None: + # Re-validation applies the collection-name rules to the override. + schema = CollectionSchema.model_validate( + {**schema.model_dump(mode="python"), "name": body.name} + ) + # Fail-fast on snapshot options (e.g. includeVector=false) *before* + # creating anything — otherwise a mismatch would leave a half-imported + # ghost collection behind that permanently occupies targetPath. + check_schema_compatible(manifest, schema) + + target = _normalise_path(body.targetPath) + if target.exists(): + # The Zvec engine refuses to create into *any* existing path (even an + # empty directory — verified empirically), so reject up front with a + # meaningful 409 instead of leaking the engine's path-validation error. + raise CollectionAlreadyExistsError( + f"Target directory '{target}' already exists; choose a new location.", + extra={"targetPath": str(target)}, + ) + + record = backend.create(path=target, schema=schema) + try: + report = backend.import_documents( + record.name, + source_path=str(source), + fmt=JsonlFormat(), + path=str(target), + ) + except Exception: + # The load blew up mid-stream (corrupt member, IO error, ...): roll + # back the freshly created collection so a fixed-up retry is possible. + with contextlib.suppress(Exception): + backend.destroy(record.name, path=str(target)) + shutil.rmtree(target, ignore_errors=True) + raise + + with contextlib.suppress(Exception): + store.touch_recent(record.path, name=record.name) + + return CollectionImportResponse( + collection=_summary(record, backend), + report=import_report_to_response(report), + ) + + # --------------------------------------------------------------------------- # Recent collections (workspace helper). MUST be declared before ``/{name}`` # so the literal ``recent`` segment matches first; ``recent`` is also marked diff --git a/apps/backend/zvec_studio/routers/documents.py b/apps/backend/zvec_studio/routers/documents.py index 1529940..0d19a9b 100644 --- a/apps/backend/zvec_studio/routers/documents.py +++ b/apps/backend/zvec_studio/routers/documents.py @@ -9,10 +9,18 @@ from __future__ import annotations -from typing import Annotated, Any +import itertools +import shutil +import tempfile +from datetime import datetime +from pathlib import Path as _Path +from typing import Annotated, Any, Literal from fastapi import APIRouter, Depends, Path, Query, Request, status +from fastapi.responses import StreamingResponse +from starlette.background import BackgroundTask +from zvec_studio.exceptions import ZvecStudioError from zvec_studio.schemas import ( DocumentBatchDeleteRequest, DocumentBatchDeleteResponse, @@ -20,6 +28,9 @@ DocumentBrowseResponse, DocumentDeleteByFilterRequest, DocumentDeleteByFilterResponse, + DocumentImportErrorEntry, + DocumentImportRequest, + DocumentImportResponse, DocumentInsertRequest, DocumentInsertResponse, DocumentUpdateRequest, @@ -28,6 +39,13 @@ DocumentUpsertResponse, ) from zvec_studio.storage import SdkBackend +from zvec_studio.storage.formats import resolve_export_format, resolve_import_format +from zvec_studio.storage.import_ import ImportMode, ImportReport, OnErrorMode +from zvec_studio.storage.snapshot import ( + build_manifest, + stream_snapshot_package, + write_snapshot_package, +) router = APIRouter(prefix="/collections/{name}/documents", tags=["documents"]) @@ -163,3 +181,177 @@ def delete_by_filter( ) -> DocumentDeleteByFilterResponse: deleted = backend.delete_by_filter(_resolve(backend, name, path), body.filter) return DocumentDeleteByFilterResponse(deleted=deleted) + + +@router.post( + ":import", + response_model=DocumentImportResponse, +) +def import_documents( + name: Annotated[str, Path(min_length=1)], + body: DocumentImportRequest, + backend: Annotated[SdkBackend, Depends(_get_backend)], + path: str | None = Query(default=None), +) -> DocumentImportResponse: + """Bulk-import documents from a JSONL file, streaming. + + The file is read from the backend's local filesystem (Studio is + local-first — the file lives on the same machine), parsed row by row, + and written in SDK-sized batches. Row-level failures are reported in the + 200 body (partial-success semantics); only request-level problems (bad + format, unreadable file, unknown collection) surface as 4xx. + + ``format`` defaults to the file extension; extension-less files are + treated as JSONL. Snapshot packages (``.tar.gz`` / ``.tgz``) always + carry a JSONL data member, so their format resolves to ``jsonl``. + """ + source_path = body.source.path + if source_path.lower().endswith((".tar.gz", ".tgz")): + fmt_name = body.format or "jsonl" + else: + fmt_name = body.format or _Path(source_path).suffix.lstrip(".") or "jsonl" + fmt = resolve_import_format(fmt_name) + + extra: dict[str, Any] = {} + if body.batchSize is not None: + extra["batch_size"] = body.batchSize + report = backend.import_documents( + _resolve(backend, name, path), + source_path=source_path, + fmt=fmt, + mode=ImportMode(body.mode), + on_error=OnErrorMode(body.onError), + **extra, + ) + return import_report_to_response(report) + + +def import_report_to_response(report: ImportReport) -> DocumentImportResponse: + """Map the storage-layer ``ImportReport`` onto its HTTP payload. + + Shared with the collection-level ``:restore`` endpoint, which embeds + the same row-level report in its response. + """ + return DocumentImportResponse( + imported=report.imported, + failed=report.failed, + totalLines=report.total_lines, + aborted=report.aborted, + durationMs=report.duration_ms, + errors=[ + DocumentImportErrorEntry(line=e.line, code=e.code, message=e.message) + for e in report.errors + ], + errorsTruncated=report.errors_truncated, + ) + + +@router.get(":export") +def export_documents( + name: Annotated[str, Path(min_length=1)], + backend: Annotated[SdkBackend, Depends(_get_backend)], + path: str | None = Query(default=None), + includeVector: Annotated[ + bool, + Query(description="Include vector data in each row (default true)."), + ] = True, + outputFields: Annotated[ + str | None, + Query(description="Comma-separated scalar fields to include; omit for all."), + ] = None, + format: Annotated[ + str, + Query(description="Export format name. Supported: jsonl."), + ] = "jsonl", + mode: Annotated[ + Literal["data", "snapshot"], + Query( + description=( + "data: a single JSONL file. snapshot: manifest.json + JSONL " + "bundled in a tar.gz (carries the schema for migration)." + ) + ), + ] = "data", +) -> StreamingResponse: + """Stream every document in the collection as a downloadable file. + + Uses the SDK snapshot iterator (constant memory; writes during the export + are invisible to it), so arbitrarily large collections can be exported. + The response is chunked — the frontend must trigger a *native* download + (````), never ``fetch().blob()``, which would buffer the whole + body in memory. + + In ``data`` mode the first chunk is pulled eagerly so failures that occur + before any byte is sent (blocked iterator, non-finite value in the first + row) still map to a proper 4xx/5xx instead of truncating a stream that + already started. In ``snapshot`` mode the JSONL is staged in a temp dir + first (a tar member must declare its size up front), then the tar.gz is + streamed; the temp dir is removed after the response finishes. + """ + fmt = resolve_export_format(format) + resolved = _resolve(backend, name, path) + fields = ( + [f.strip() for f in outputFields.split(",") if f.strip()] if outputFields else None + ) + stamp = datetime.now().strftime("%Y%m%d-%H%M%S") + + if mode == "snapshot": + tmp_dir = _Path(tempfile.mkdtemp(prefix="zvec-studio-export-")) + try: + record = backend.get(resolved) + manifest = build_manifest( + schema=record.schema, + include_vector=includeVector, + output_fields=fields, + ) + rows = backend.iter_documents( + resolved, include_vector=includeVector, output_fields=fields + ) + manifest_path, data_path = write_snapshot_package( + rows=rows, serialize=fmt.serialize, manifest=manifest, tmp_dir=tmp_dir + ) + except RuntimeError as exc: + shutil.rmtree(tmp_dir, ignore_errors=True) + raise ZvecStudioError( + f"Export cannot start: {exc}", + code="EXPORT_BLOCKED", + status_code=409, + ) from exc + except Exception: + shutil.rmtree(tmp_dir, ignore_errors=True) + raise + return StreamingResponse( + stream_snapshot_package(manifest_path=manifest_path, data_path=data_path), + media_type="application/gzip", + headers={ + "Content-Disposition": f'attachment; filename="{resolved}-{stamp}.tar.gz"' + }, + background=BackgroundTask(shutil.rmtree, str(tmp_dir), True), + ) + + rows = backend.iter_documents( + resolved, include_vector=includeVector, output_fields=fields + ) + chunks = fmt.serialize(rows) + + try: + first = next(chunks) + except StopIteration: + body: Any = iter(()) + except RuntimeError as exc: + # The SDK raises while an iterator cannot be opened (e.g. a + # maintenance op is running) — surface as a retryable 409. + raise ZvecStudioError( + f"Export cannot start: {exc}", + code="EXPORT_BLOCKED", + status_code=409, + ) from exc + else: + body = itertools.chain((first,), chunks) + + filename = f"{resolved}-{stamp}.{fmt.extension}" + return StreamingResponse( + body, + media_type=fmt.content_type, + headers={"Content-Disposition": f'attachment; filename="{filename}"'}, + ) diff --git a/apps/backend/zvec_studio/routers/fs.py b/apps/backend/zvec_studio/routers/fs.py index 27af854..d4b3075 100644 --- a/apps/backend/zvec_studio/routers/fs.py +++ b/apps/backend/zvec_studio/routers/fs.py @@ -7,10 +7,12 @@ model), which is useless when the user wants to specify an absolute storage path. -This router exposes a single read-only endpoint that lists *directory* -entries on the host filesystem. The frontend renders a custom modal picker -on top of it, navigating up/down the tree and returning a real absolute -path. The endpoint never reads file contents. +This router exposes a read-only listing endpoint that walks the host +filesystem for the picker UIs, navigating up/down the tree and returning +real absolute paths, plus a reveal helper. By default only directories are +listed (the collection-open picker); with ``includeFiles=true`` files appear +as well (the import file picker), optionally filtered by extension. The +endpoints never read file contents. """ from __future__ import annotations @@ -18,7 +20,7 @@ import platform import subprocess from pathlib import Path -from typing import Annotated +from typing import Annotated, Literal from fastapi import APIRouter, HTTPException, Query, status from pydantic import BaseModel, Field @@ -27,10 +29,18 @@ class FsEntry(BaseModel): - """A single subdirectory entry.""" + """A single directory (or file) entry.""" - name: str = Field(..., description="Leaf name of the directory.") - path: str = Field(..., description="Absolute path of the directory.") + name: str = Field(..., description="Leaf name of the entry.") + path: str = Field(..., description="Absolute path of the entry.") + kind: Literal["dir", "file"] = Field( + default="dir", + description="Entry kind. 'file' entries only appear when includeFiles is set.", + ) + size: int | None = Field( + default=None, + description="Size in bytes for files; null for directories.", + ) class FsListing(BaseModel): @@ -44,7 +54,10 @@ class FsListing(BaseModel): home: str = Field(..., description="The current user's home directory (absolute).") entries: list[FsEntry] = Field( default_factory=list, - description="Sorted subdirectories (excludes hidden entries by default).", + description=( + "Entries sorted by name (case-insensitive). Files only appear " + "when includeFiles=true." + ), ) @@ -107,6 +120,28 @@ def reveal_in_file_manager(body: RevealRequest) -> None: ) from exc +def _parse_extensions(raw: str | None) -> tuple[str, ...] | None: + """Normalise a comma-separated extension filter (``.jsonl,.tar.gz``). + + Returns lowered suffixes (including multi-part ones like ``.tar.gz``), + or None when no filter was requested. + """ + if not raw: + return None + parts = [p.strip().lower() for p in raw.split(",") if p.strip()] + if not parts: + return None + return tuple(p if p.startswith(".") else f".{p}" for p in parts) + + +def _file_size(path: Path) -> int | None: + try: + return path.stat().st_size + except OSError: + # Unstatable entries still show up, just without a size. + return None + + @router.get("/list", response_model=FsListing) def list_directory( path: Annotated[ @@ -117,25 +152,65 @@ def list_directory( bool, Query(description="Include dotfile-prefixed entries when true."), ] = False, + include_files: Annotated[ + bool, + Query( + alias="includeFiles", + description=( + "Also list files (not only directories). Needed by the " + "import file picker; default keeps the legacy directory-only " + "behaviour." + ), + ), + ] = False, + extensions: Annotated[ + str | None, + Query( + description=( + "Comma-separated file-extension filter applied when " + "includeFiles is set, e.g. '.jsonl,.tar.gz'. Directories " + "are never filtered out (they are needed for navigation). " + "Case-insensitive." + ), + ), + ] = None, ) -> FsListing: - """List subdirectories of ``path`` (or ``$HOME`` when omitted). + """List entries of ``path`` (or ``$HOME`` when omitted). - Returns the resolved absolute path, parent path, and sorted entries. - Files are intentionally excluded — the picker selects a directory only. + By default only subdirectories are returned (the original directory + picker contract). With ``includeFiles=true`` files are listed as well, + optionally narrowed by ``extensions``; each entry carries ``kind`` and, + for files, ``size``. """ target = _resolve(path) + ext_filter = _parse_extensions(extensions) if include_files else None entries: list[FsEntry] = [] try: for child in sorted(target.iterdir(), key=lambda p: p.name.lower()): if not show_hidden and child.name.startswith("."): continue try: - if not child.is_dir(): - continue + is_dir = child.is_dir() except OSError: # Broken symlinks / permission errors -> skip silently. continue - entries.append(FsEntry(name=child.name, path=str(child))) + if is_dir: + entries.append(FsEntry(name=child.name, path=str(child), kind="dir")) + continue + if not include_files: + continue + if ext_filter is not None and not any( + child.name.lower().endswith(ext) for ext in ext_filter + ): + continue + entries.append( + FsEntry( + name=child.name, + path=str(child), + kind="file", + size=_file_size(child), + ) + ) except PermissionError as exc: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, diff --git a/apps/backend/zvec_studio/schemas/__init__.py b/apps/backend/zvec_studio/schemas/__init__.py index 7b8f960..6a5406c 100644 --- a/apps/backend/zvec_studio/schemas/__init__.py +++ b/apps/backend/zvec_studio/schemas/__init__.py @@ -38,6 +38,8 @@ ) from zvec_studio.schemas.collection import ( CollectionCreateRequest, + CollectionImportRequest, + CollectionImportResponse, CollectionListItem, CollectionListResponse, CollectionOpenRequest, @@ -68,12 +70,16 @@ DocumentBrowseResponse, DocumentDeleteByFilterRequest, DocumentDeleteByFilterResponse, + DocumentImportErrorEntry, + DocumentImportRequest, + DocumentImportResponse, DocumentInsertRequest, DocumentInsertResponse, DocumentUpdateRequest, DocumentUpdateResponse, DocumentUpsertRequest, DocumentUpsertResponse, + ImportSourceSpec, ) from zvec_studio.schemas.search import ( DiskAnnQueryParamSpec, @@ -96,6 +102,8 @@ "BM25Config", "BM25Language", "CollectionCreateRequest", + "CollectionImportRequest", + "CollectionImportResponse", "CollectionListItem", "CollectionListResponse", "CollectionOpenRequest", @@ -112,6 +120,9 @@ "DocumentBrowseResponse", "DocumentDeleteByFilterRequest", "DocumentDeleteByFilterResponse", + "DocumentImportErrorEntry", + "DocumentImportRequest", + "DocumentImportResponse", "DocumentInsertRequest", "DocumentInsertResponse", "DocumentUpdateRequest", @@ -134,6 +145,7 @@ "FtsSpec", "HnswQueryParamSpec", "HnswRabitqQueryParamSpec", + "ImportSourceSpec", "IndexCreateRequest", "IndexType", "IvfQueryParamSpec", diff --git a/apps/backend/zvec_studio/schemas/collection.py b/apps/backend/zvec_studio/schemas/collection.py index 447e30a..c011ef2 100644 --- a/apps/backend/zvec_studio/schemas/collection.py +++ b/apps/backend/zvec_studio/schemas/collection.py @@ -22,9 +22,13 @@ # least 3 characters long (verified empirically against zvec 0.4.0). _COLLECTION_NAME_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_]{2,63}$") -# Field names colliding with Zvec's built-in ``Doc.id``. The SDK injects ``id`` -# automatically; user-defined columns must not shadow it. -_RESERVED_FIELD_NAMES: frozenset[str] = frozenset({"id", "_id"}) +# ``id`` and ``_id`` used to be rejected here on the assumption that Zvec +# "injects" ``id`` and a same-named column would shadow it. That assumption is +# wrong: ``Doc.id`` lives beside the columns, and Zvec accepts a schema that +# declares a scalar field or vector named ``id`` (or ``_id``) — the two coexist +# independently. Rejecting them made Studio unable to *open* collections +# created through the SDK, so the restriction is gone; the primary key gets its +# own row key (``$id``) whenever a column takes ``id``, see storage/doc_repr.py. _FTS_TOKENIZERS: frozenset[str] = frozenset({"standard", "whitespace", "jieba"}) _FTS_FILTERS: frozenset[str] = frozenset({"lowercase", "ascii_folding", "stemmer"}) @@ -36,11 +40,15 @@ class VectorDataType(str, Enum): - """Vector dtypes supported by the Zvec SDK.""" + """Vector dtypes supported by the Zvec SDK. + + ``VECTOR_FP64`` is intentionally absent: zvec's schema validation rejects + it ("dense_vector's data type only support FP32"), so exposing it would + only let users build collections that fail at create time. + """ VECTOR_FP32 = "VECTOR_FP32" VECTOR_FP16 = "VECTOR_FP16" - VECTOR_FP64 = "VECTOR_FP64" VECTOR_INT8 = "VECTOR_INT8" SPARSE_VECTOR_FP32 = "SPARSE_VECTOR_FP32" SPARSE_VECTOR_FP16 = "SPARSE_VECTOR_FP16" @@ -106,8 +114,6 @@ class FieldSchema(BaseModel): def _validate_name(cls, v: str) -> str: if not _NAME_RE.match(v): raise ValueError("field name must match ^[A-Za-z_][A-Za-z0-9_]{0,63}$") - if v in _RESERVED_FIELD_NAMES: - raise ValueError(f"field name '{v}' is reserved by the SDK and cannot be used") return v @model_validator(mode="after") @@ -179,8 +185,6 @@ class VectorSchema(BaseModel): def _validate_name(cls, v: str) -> str: if not _NAME_RE.match(v): raise ValueError("vector name must match ^[A-Za-z_][A-Za-z0-9_]{0,63}$") - if v in _RESERVED_FIELD_NAMES: - raise ValueError(f"vector name '{v}' is reserved by the SDK and cannot be used") return v @@ -350,8 +354,6 @@ class FieldRenameRequest(BaseModel): def _validate_name(cls, v: str) -> str: if not _NAME_RE.match(v): raise ValueError("new field name must match ^[A-Za-z_][A-Za-z0-9_]{0,63}$") - if v in _RESERVED_FIELD_NAMES: - raise ValueError(f"field name '{v}' is reserved by the SDK and cannot be used") return v @@ -401,3 +403,62 @@ class MaintenanceResponse(BaseModel): operation: str = Field(..., description="Operation performed (``flush``, ``optimize``).") timestamp: str = Field(..., description="ISO-8601 timestamp when the operation finished.") + + +# --------------------------------------------------------------------------- +# Snapshot restore (collection-level lifecycle operation). +# --------------------------------------------------------------------------- + + +class CollectionImportRequest(BaseModel): + """Import a collection from a snapshot package (``.tar.gz``). + + A sibling of create/open, not a variant of document import: the manifest + embedded in the snapshot supplies the schema, ``targetPath`` gives the new + collection a home on disk, and the data rows are loaded in the same pass. + """ + + model_config = ConfigDict(extra="forbid") + + source: ImportSourceSpec = Field( + ..., description="The snapshot package to import from (local path)." + ) + targetPath: str = Field( + ..., + min_length=1, + description=( + "Directory for the imported collection. Must not exist (409 if it " + "does, empty directories included); the collection is created there." + ), + ) + name: str | None = Field( + default=None, + description=( + "Optional new collection name. Defaults to the name recorded in " + "the snapshot manifest; override it to avoid clashing with an " + "open collection." + ), + ) + + @field_validator("name") + @classmethod + def _validate_name(cls, v: str | None) -> str | None: + if v is not None and not _COLLECTION_NAME_RE.match(v): + raise ValueError("collection name must match ^[A-Za-z][A-Za-z0-9_]{2,63}$") + return v + + +class CollectionImportResponse(BaseModel): + """The imported collection plus the row-level load report.""" + + collection: CollectionSummary + report: DocumentImportResponse + + +from zvec_studio.schemas.document import ( # noqa: E402 (avoid import cycle) + DocumentImportResponse, + ImportSourceSpec, +) + +CollectionImportRequest.model_rebuild() +CollectionImportResponse.model_rebuild() diff --git a/apps/backend/zvec_studio/schemas/document.py b/apps/backend/zvec_studio/schemas/document.py index 3543817..1e98f14 100644 --- a/apps/backend/zvec_studio/schemas/document.py +++ b/apps/backend/zvec_studio/schemas/document.py @@ -7,7 +7,7 @@ from __future__ import annotations -from typing import Any +from typing import Any, Literal from pydantic import BaseModel, ConfigDict, Field @@ -110,3 +110,82 @@ class DocumentDeleteByFilterRequest(BaseModel): class DocumentDeleteByFilterResponse(BaseModel): deleted: int + + +# --------------------------------------------------------------------------- +# Bulk import from a file (JSONL). The row-level report mirrors +# ``storage/import_.py::ImportReport`` verbatim; row failures are carried in +# the 200 body (partial-success batch semantics), only request-level problems +# map to 4xx. +# --------------------------------------------------------------------------- + + +class ImportSourceSpec(BaseModel): + """Where the import reads its data from. + + Only ``localPath`` exists today — Studio is local-first, the backend runs + on the user's machine, and reading a path avoids shovelling gigabytes + through an HTTP upload. ``upload`` is reserved for the remote deployment + enhancement. + """ + + model_config = ConfigDict(extra="forbid") + + kind: Literal["localPath"] = "localPath" + path: str = Field(..., min_length=1, description="Absolute path of the file to import.") + + +class DocumentImportRequest(BaseModel): + """Body for ``POST /collections/{name}/documents:import``.""" + + model_config = ConfigDict(extra="forbid") + + source: ImportSourceSpec + mode: Literal["insert", "replace"] = Field( + default="replace", + description=( + "insert: existing primary keys fail the row. replace: the row's " + "content becomes the whole document for that key (idempotent " + "re-import)." + ), + ) + onError: Literal["abort", "skip"] = Field( + default="abort", + description="abort: stop at the first failing row; skip: record it and continue.", + ) + format: str | None = Field( + default=None, + description=( + "Import format name ('jsonl'). When omitted, inferred from the " + "file extension; extension-less files default to jsonl." + ), + ) + batchSize: int | None = Field( + default=None, + ge=1, + le=1024, + description=( + "Internal write batch size. Capped at the SDK's 1024-document " + "limit; defaults to the benchmarked sweet spot (512)." + ), + ) + + +class DocumentImportErrorEntry(BaseModel): + """One failing row.""" + + line: int = Field(..., description="1-based physical line in the file.") + code: str = Field(..., description="Studio error code for this row.") + message: str + + +class DocumentImportResponse(BaseModel): + imported: int + failed: int + totalLines: int = Field(..., description="Data rows read (blank lines excluded).") + aborted: bool + durationMs: float + errors: list[DocumentImportErrorEntry] = Field( + ..., description="First 100 failing rows; see errorsTruncated." + ) + errorsTruncated: bool diff --git a/apps/backend/zvec_studio/storage/doc_repr.py b/apps/backend/zvec_studio/storage/doc_repr.py new file mode 100644 index 0000000..d925add --- /dev/null +++ b/apps/backend/zvec_studio/storage/doc_repr.py @@ -0,0 +1,176 @@ +"""Document wire representation: ``zvec.Doc`` <-> flat row dict. + +Every place that turns a document into JSON (HTTP responses, JSONL export) +or reads one back (HTTP writes, JSONL import) goes through here, so the two +directions can never drift apart. + +Primary-key representation +-------------------------- +A Zvec document has a primary key (``Doc.id``) that lives *beside* its +columns, and Zvec accepts a column named ``id`` (field or vector — they are +independent of the primary key), which a naive flat mapping would let shadow +the key. The row key carrying the primary key is therefore chosen from a +reserved chain: + +* no column named ``id`` -> primary key sits under ``id`` +* a column named ``id`` -> primary key sits under ``$id`` +* a column named ``$id`` -> primary key under ``$$id`` +* …and so on, prepending one more ``$`` per clash + +Today the chain stops at ``$id`` in practice: the Zvec engine rejects ``$`` +in column names at create time (fields via its schema regex, vectors +empirically too), and Studio enforces the same regex — so ``$id`` can never +be a user column. The deeper links are defense in depth: they keep the +mapping total for hand-crafted manifests and for any future loosening of the +engine's naming rules, so the two directions can never corrupt a key. +""" + +from __future__ import annotations + +import re +from collections.abc import Sequence +from typing import Any, Protocol + +from zvec_studio.exceptions import InvalidSchemaError + +#: Row key carrying the primary key for an ordinary schema. +PK_KEY = "id" +#: First reserved fallback when a column occupies ``id``. +RESERVED_PK_KEY = "$id" +#: Keys of the reserved chain: one or more ``$`` prefixes + ``id``. +_RESERVED_KEY_RE = re.compile(r"^\$+id$") + + +class _NamedColumn(Protocol): + name: str + + +class _SchemaLike(Protocol): + """The slice of ``CollectionSchema`` this module needs. + + Declared read-only (properties rather than attributes) so structural + matching stays covariant: the concrete SDK/Studio types carry narrower + value types than the protocol needs to know about. + """ + + @property + def fields(self) -> Sequence[Any]: ... + + @property + def vectors(self) -> Sequence[Any]: ... + + +class _DocLike(Protocol): + """The slice of ``zvec.Doc`` this module needs (read-only, see above).""" + + @property + def id(self) -> str: ... + + @property + def fields(self) -> Any: ... + + @property + def vectors(self) -> Any: ... + + +def _column_names(schema: _SchemaLike) -> set[str]: + # Field and vector names share one namespace in Zvec. + return {c.name for c in (*schema.fields, *schema.vectors)} + + +def declares_id_column(schema: _SchemaLike) -> bool: + """Return True when *schema* has a column literally named ``id``.""" + return PK_KEY in _column_names(schema) + + +def pk_key(schema: _SchemaLike) -> str: + """Return the row key that carries the primary key for *schema*. + + Walks the reserved chain ``id`` -> ``$id`` -> ``$$id`` -> ... until a key + no user column can occupy is found (see module docstring). + """ + names = _column_names(schema) + key = PK_KEY + while key in names: + key = "$" + key + return key + + +def doc_to_row( + doc: _DocLike, *, schema: _SchemaLike, include_vector: bool +) -> dict[str, Any]: + """Flatten a document into a row dict. + + The primary key is written first under :func:`pk_key`; because that key + is never occupied by a user column, nothing can shadow it. + """ + out: dict[str, Any] = {pk_key(schema): doc.id} + if doc.fields: + out.update(doc.fields) + if include_vector and doc.vectors: + out.update(doc.vectors) + return out + + +def split_pk( + row: dict[str, Any], *, schema: _SchemaLike +) -> tuple[str | None, dict[str, Any]]: + """Split *row* into ``(primary key or None, remaining columns)``. + + ``None`` means the row carries no primary key and the caller should mint + one (the SDK requires a str id; callers generate a ULID). An explicit + ``null`` value behaves exactly like an omitted key (the historical write + contract), so ``{"id": null}`` still auto-generates instead of failing. + + Reserved keys take precedence over a bare ``id`` and are accepted even + when the schema does not need them, so exports stay importable across + schema shapes. + + Raises: + InvalidSchemaError: if the primary key is not a string, or if the row + is ambiguous: a column occupies ``id`` (so the bare ``id`` could + be that column) and no reserved key is present. Guessing either + way risks silently misplacing the key. + """ + rest = dict(row) + names = _column_names(schema) + key = pk_key(schema) + + # Accept any reserved-chain key the row carries (the deepest first), but + # never one a column occupies — that value belongs to the column, not to + # the primary key. This is independent of the target schema's own chain + # depth, so a file exported from a clashing schema stays importable into + # any other schema. + reserved = sorted( + (k for k in rest if _RESERVED_KEY_RE.match(k)), key=len, reverse=True + ) + for candidate in reserved: + if candidate in names: + continue + raw = rest.pop(candidate) + if raw is None: + return None, rest + return _coerce_pk(raw, key=candidate), rest + + if PK_KEY in rest: + if declares_id_column(schema): + raise InvalidSchemaError( + "Ambiguous 'id': a column occupies 'id' on this collection, so " + f"the primary key must be given as '{key}'.", + extra={"column": PK_KEY, "primaryKeyKey": key}, + ) + raw = rest.pop(PK_KEY) + if raw is None: + return None, rest + return _coerce_pk(raw, key=PK_KEY), rest + + return None, rest + + +def _coerce_pk(raw: Any, *, key: str) -> str: + if not isinstance(raw, str): + raise InvalidSchemaError( + f"Document '{key}' must be a string (Zvec requires str ids).", + extra={key: raw}, + ) + return raw diff --git a/apps/backend/zvec_studio/storage/formats.py b/apps/backend/zvec_studio/storage/formats.py new file mode 100644 index 0000000..007085f --- /dev/null +++ b/apps/backend/zvec_studio/storage/formats.py @@ -0,0 +1,212 @@ +"""Import/export format protocols. + +The format is the only variable piece of the import/export pipeline (design +doc §4.2). An :class:`ImportFormat` turns a byte stream into ``(line number, +record)`` pairs; formats are looked up by name or file extension through a +registry. Shipping JSONL today keeps the dependency surface empty (orjson is +already a core dependency); CSV/Parquet can join later by implementing the +protocol and registering — the pipeline itself stays untouched. +""" + +from __future__ import annotations + +import math +from collections.abc import Iterator +from typing import Any, BinaryIO, Protocol, runtime_checkable + +import orjson + +from zvec_studio.exceptions import ExportNonFiniteError, InvalidDocumentError +from zvec_studio.storage.doc_repr import PK_KEY, RESERVED_PK_KEY + + +@runtime_checkable +class ImportFormat(Protocol): + """Parses a file byte stream into ``(line number, record)`` pairs. + + ``line number`` is the 1-based physical line the record came from, so + downstream errors can point the user at the exact spot in their editor. + Implementations must stream: holding more than the current record is + forbidden (files may be gigabytes). + """ + + name: str + extensions: tuple[str, ...] + + def parse(self, stream: BinaryIO) -> Iterator[tuple[int, dict[str, Any]]]: ... + + +@runtime_checkable +class ExportFormat(Protocol): + """Serializes document rows into file bytes (used by the export API).""" + + name: str + content_type: str + extension: str + + def serialize(self, rows: Iterator[dict[str, Any]]) -> Iterator[bytes]: ... + + +def _row_document_id(row: dict[str, Any]) -> str: + """Best-effort primary key for error reporting (``$id`` wins over ``id``). + + Key order mirrors :mod:`zvec_studio.storage.doc_repr` — the reserved key + wins whenever the schema declares its own ``id`` column. + """ + for key in (RESERVED_PK_KEY, PK_KEY): + value = row.get(key) + if value is not None: + return str(value) + return "" + + +def _check_finite(value: Any, *, path: str, document_id: str) -> None: + """Raise :class:`ExportNonFiniteError` at the first NaN/±Inf found. + + Walks scalars, lists and dicts so vectors (``list[float]``), sparse + vectors (``{index: weight}``) and array fields are all covered. + """ + if isinstance(value, bool): + return + if isinstance(value, float) and not math.isfinite(value): + raise ExportNonFiniteError( + f"Document {document_id} contains a non-finite value at '{path}'.", + extra={"documentId": document_id, "path": path}, + ) + if isinstance(value, list): + for i, item in enumerate(value): + _check_finite(item, path=f"{path}[{i}]", document_id=document_id) + elif isinstance(value, dict): + for k, v in value.items(): + _check_finite(v, path=f"{path}.{k}", document_id=document_id) + + +class _JsonlLines: + """Stateful line iterator for :class:`JsonlFormat`. + + Deliberately a class (not a generator): after raising on a bad line the + iterator must keep going — ``skip``-mode imports continue past row + errors, and a generator would be finished by the first raise. + """ + + def __init__(self, stream: BinaryIO) -> None: + self._stream = stream + self._line_number = 0 + + def __iter__(self) -> _JsonlLines: + return self + + def __next__(self) -> tuple[int, dict[str, Any]]: + while True: + raw = self._stream.readline() + if not raw: + raise StopIteration + self._line_number += 1 + text = raw.strip() + if not text: + continue + try: + record = orjson.loads(text) + except orjson.JSONDecodeError as exc: + raise InvalidDocumentError( + f"Line {self._line_number} is not valid JSON: {exc}", + code="INVALID_DOCUMENT", + extra={"line": self._line_number}, + ) from exc + if not isinstance(record, dict): + raise InvalidDocumentError( + f"Line {self._line_number} must be a JSON object, got " + f"{type(record).__name__}.", + code="INVALID_DOCUMENT", + extra={"line": self._line_number}, + ) + return self._line_number, record + + +class JsonlFormat: + """JSON Lines: one JSON object per line, UTF-8. + + The row structure matches the document API representation (see + ``storage/doc_repr.py``) so exports can be re-imported verbatim. + """ + + name = "jsonl" + content_type = "application/x-ndjson" + extension = "jsonl" + # Annotated so mypy widens the literal to the protocol's + # ``tuple[str, ...]`` (a bare literal infers as fixed-length). + extensions: tuple[str, ...] = ("jsonl", "ndjson") + + def parse(self, stream: BinaryIO) -> Iterator[tuple[int, dict[str, Any]]]: + return _JsonlLines(stream) + + def serialize(self, rows: Iterator[dict[str, Any]]) -> Iterator[bytes]: + """Yield one ``\\n``-terminated JSON line per row, lazily. + + Each row is validated for non-finite floats *before* serialization: + orjson would otherwise coerce NaN/±Inf to ``null`` (silent data loss). + ``OPT_NON_STR_KEYS`` lets sparse vectors keep their ``int`` indexes + (orjson rejects non-str dict keys otherwise). + """ + for row in rows: + document_id = _row_document_id(row) + for key, value in row.items(): + _check_finite(value, path=key, document_id=document_id) + yield orjson.dumps(row, option=orjson.OPT_NON_STR_KEYS) + b"\n" + + +#: Import formats keyed by canonical name; extensions also resolve here. +IMPORT_FORMATS: dict[str, ImportFormat] = { + "jsonl": JsonlFormat(), +} + +#: Export formats keyed by canonical name (see ``storage/export.py``). +EXPORT_FORMATS: dict[str, ExportFormat] = { + "jsonl": JsonlFormat(), +} + +_EXTENSION_TO_FORMAT: dict[str, ImportFormat] = { + ext: fmt for fmt in IMPORT_FORMATS.values() for ext in fmt.extensions +} + + +def resolve_import_format(name_or_extension: str) -> ImportFormat: + """Look up an import format by name or file extension (case-insensitive). + + Raises: + InvalidDocumentError: with code ``IMPORT_UNSUPPORTED_FORMAT`` (400) + when the value matches nothing; the message lists what is + supported so the caller can fix it. + """ + key = name_or_extension.strip().lower().lstrip(".") + fmt = IMPORT_FORMATS.get(key) or _EXTENSION_TO_FORMAT.get(key) + if fmt is None: + supported = ", ".join(sorted(IMPORT_FORMATS)) + raise InvalidDocumentError( + f"Unsupported import format '{name_or_extension}'. Supported: {supported}.", + code="IMPORT_UNSUPPORTED_FORMAT", + status_code=400, + extra={"format": name_or_extension, "supported": sorted(IMPORT_FORMATS)}, + ) + return fmt + + +def resolve_export_format(name: str) -> ExportFormat: + """Look up an export format by name (case-insensitive). + + Raises: + InvalidDocumentError: with code ``EXPORT_UNSUPPORTED_FORMAT`` (400) + when the name matches nothing; the message lists what is + supported so the caller can fix it. + """ + key = name.strip().lower() + fmt = EXPORT_FORMATS.get(key) + if fmt is None: + supported = ", ".join(sorted(EXPORT_FORMATS)) + raise InvalidDocumentError( + f"Unsupported export format '{name}'. Supported: {supported}.", + code="EXPORT_UNSUPPORTED_FORMAT", + status_code=400, + extra={"format": name, "supported": sorted(EXPORT_FORMATS)}, + ) + return fmt diff --git a/apps/backend/zvec_studio/storage/import_.py b/apps/backend/zvec_studio/storage/import_.py new file mode 100644 index 0000000..d14c47a --- /dev/null +++ b/apps/backend/zvec_studio/storage/import_.py @@ -0,0 +1,104 @@ +"""Import contracts and source handling. + +The pipeline itself (parse -> build -> batched write) is orchestrated by +:meth:`SdkBackend.import_documents`; this module holds the pieces that do not +need a live collection: + +* the import vocabulary (modes, error policies) and report dataclasses that + the HTTP layer serializes verbatim; +* validation of the source path, mapping filesystem problems onto the Studio + error taxonomy (404 / 403) instead of leaking ``OSError``s as 500s. + +See design doc §6.4 for the semantics of each mode and policy. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path + +from zvec_studio.exceptions import ( + ImportFileNotFoundError, + ImportFileNotReadableError, +) + +#: Maximum number of per-row errors carried in a report; the import keeps +#: counting beyond this but only returns the first slice (the pattern is +#: enough for the user to fix their file). +MAX_REPORTED_ERRORS = 100 + + +class ImportMode(str, Enum): + """How incoming documents are written. + + * ``insert`` - strict: existing primary keys fail the row. + * ``replace`` - the file's content becomes the whole document for that + primary key (SDK-native ``upsert``). Chosen over field-level merge so + re-importing the same file is idempotent; see design doc §6.4. + """ + + INSERT = "insert" + REPLACE = "replace" + + +class OnErrorMode(str, Enum): + """What happens when a row fails. + + * ``abort`` - stop at the first failing row (rows written so far stay). + * ``skip`` - record the row error and continue to the end. + """ + + ABORT = "abort" + SKIP = "skip" + + +@dataclass +class ImportFailure: + """A single failing row.""" + + line: int + code: str + message: str + + +@dataclass +class ImportReport: + """Outcome of an import run (serialised as the API response).""" + + imported: int = 0 + failed: int = 0 + total_lines: int = 0 + aborted: bool = False + duration_ms: float = 0.0 + errors: list[ImportFailure] = field(default_factory=list) + errors_truncated: bool = False + + def add_failure(self, failure: ImportFailure) -> None: + self.failed += 1 + if len(self.errors) < MAX_REPORTED_ERRORS: + self.errors.append(failure) + else: + self.errors_truncated = True + + +def validate_import_source(raw_path: str | Path) -> Path: + """Resolve *raw_path* to a readable regular file. + + Raises: + ImportFileNotFoundError: missing path, or not a regular file. + ImportFileNotReadableError: exists but cannot be read. + """ + candidate = Path(raw_path).expanduser() + if not candidate.exists() or not candidate.is_file(): + raise ImportFileNotFoundError( + f"Import source is not a readable file: {candidate}", + extra={"path": str(candidate)}, + ) + if not os.access(candidate, os.R_OK): + raise ImportFileNotReadableError( + f"Import source is not readable: {candidate}", + extra={"path": str(candidate)}, + ) + return candidate diff --git a/apps/backend/zvec_studio/storage/sdk.py b/apps/backend/zvec_studio/storage/sdk.py index 999b56b..1fae6b1 100644 --- a/apps/backend/zvec_studio/storage/sdk.py +++ b/apps/backend/zvec_studio/storage/sdk.py @@ -18,14 +18,18 @@ from __future__ import annotations +import contextlib import gc import json import math import re +import tarfile +import time +from collections.abc import Callable, Iterator from dataclasses import dataclass, field from pathlib import Path from threading import RLock -from typing import Any +from typing import Any, BinaryIO, cast import zvec from ulid import ULID @@ -70,9 +74,14 @@ CollectionAlreadyExistsError, CollectionNotFoundError, DimensionMismatchError, + DocumentConflictError, DocumentNotFoundError, + ImportManifestInvalidError, + InvalidDocumentError, InvalidFilterExpressionError, InvalidSchemaError, + MaintenanceBlockedError, + UnsupportedVectorDataTypeError, ZvecStudioError, ) from zvec_studio.schemas import ( @@ -94,6 +103,30 @@ VectorIndexParam, VectorQuerySpec, ) +from zvec_studio.storage.doc_repr import doc_to_row, pk_key, split_pk +from zvec_studio.storage.formats import ImportFormat +from zvec_studio.storage.import_ import ( + ImportFailure, + ImportMode, + ImportReport, + OnErrorMode, + validate_import_source, +) +from zvec_studio.storage.snapshot import ( + DATA_FILE_NAME, + MANIFEST_NAME, + check_schema_compatible, + parse_manifest, +) + +# Zvec's C++ core rejects write batches larger than this (observed error: +# "Too many docs: N exceeds max write batch size of 1024"). It is a runtime +# guard, not a documented contract -- keep the name honest and adjust it +# (plus any test asserting the value) if the SDK ever publishes otherwise. +_MAX_SDK_WRITE_BATCH = 1024 +# Benchmarked sweet spot: writing in 512-doc chunks measured faster than +# 1024-doc chunks on the same workload (~68k vs ~54k docs/s, M-series SSD). +_DEFAULT_WRITE_BATCH = 512 _ZVEC_EXPORTS = vars(zvec) FtsIndexParam: type[Any] = _ZVEC_EXPORTS["FtsIndexParam"] @@ -116,6 +149,21 @@ def _exc_msg(exc: BaseException) -> str: return repr(exc) +def _raise_if_maintenance_blocked(exc: BaseException) -> None: + """Map Zvec's ``iterators are open`` rejection onto a retryable 409. + + While a snapshot iterator is open (an export is running), Zvec refuses + maintenance operations with ``RuntimeError('... while iterators are + open')``. That is a transient conflict, not an invalid request — surface + it as ``MAINTENANCE_BLOCKED`` so clients can retry after the export + finishes instead of seeing an opaque 4xx/5xx. + """ + if isinstance(exc, RuntimeError) and "iterator" in _exc_msg(exc).lower(): + raise MaintenanceBlockedError( + _exc_msg(exc) or "Maintenance is blocked while an export is running." + ) from exc + + @dataclass class CollectionRecord: """Live representation of an opened Collection. @@ -154,7 +202,7 @@ class CollectionRecord: _VECTOR_TO_SDK: dict[VectorDataType, SdkDataType] = { VectorDataType.VECTOR_FP32: SdkDataType.VECTOR_FP32, VectorDataType.VECTOR_FP16: SdkDataType.VECTOR_FP16, - VectorDataType.VECTOR_FP64: SdkDataType.VECTOR_FP64, + # VECTOR_FP64 deliberately omitted: zvec rejects it at schema validation. VectorDataType.VECTOR_INT8: SdkDataType.VECTOR_INT8, VectorDataType.SPARSE_VECTOR_FP32: SdkDataType.SPARSE_VECTOR_FP32, VectorDataType.SPARSE_VECTOR_FP16: SdkDataType.SPARSE_VECTOR_FP16, @@ -530,9 +578,19 @@ def _from_sdk_schema(sdk_schema: Any, path: Path) -> CollectionSchema: for v in (sdk_schema.vectors or []): vdt = _SDK_TO_VECTOR.get(v.data_type) if vdt is None: - raise CollectionNotFoundError( - f"Unknown vector data type {v.data_type} at {path}", - extra={"path": str(path)}, + # The collection exists and opened fine — this is a dtype Studio + # cannot represent (e.g. a legacy VECTOR_FP64 collection), not a + # missing collection. + raise UnsupportedVectorDataTypeError( + f"Vector '{v.name}' uses data type {v.data_type}, which this " + f"version of Studio does not support " + f"(supported: {', '.join(t.value for t in VectorDataType)}).", + extra={ + "path": str(path), + "vector": v.name, + "dataType": str(v.data_type), + "supported": [t.value for t in VectorDataType], + }, ) vectors.append( StudioVectorSchema( @@ -558,13 +616,16 @@ def _from_sdk_schema(sdk_schema: Any, path: Path) -> CollectionSchema: return CollectionSchema(name=sdk_schema.name, vectors=vectors, fields=fields) -def _doc_to_dict(doc: SdkDoc, *, include_vector: bool) -> dict[str, Any]: - out: dict[str, Any] = {"id": doc.id} - if doc.fields: - out.update(doc.fields) - if include_vector and doc.vectors: - out.update(doc.vectors) - return out +def _doc_to_dict( + doc: SdkDoc, *, schema: CollectionSchema, include_vector: bool +) -> dict[str, Any]: + """Flatten an SDK document into the wire row representation. + + Delegates to :mod:`zvec_studio.storage.doc_repr` so the primary key lands + on ``$id`` when the schema declares its own ``id`` column instead of being + silently overwritten by it. + """ + return doc_to_row(doc, schema=schema, include_vector=include_vector) def _dir_size(path: Path) -> int: @@ -611,17 +672,6 @@ def _coerce_stats(raw: Any, collection_path: str | Path) -> CollectionStats: ) -def _ensure_id(doc: dict[str, Any]) -> str: - raw = doc.get("id") - if raw is None: - return str(ULID()) - if not isinstance(raw, str): - raise InvalidSchemaError( - "Document 'id' must be a string (Zvec requires str ids).", - extra={"id": raw}, - ) - return raw - # ── Field type validation ── @@ -694,16 +744,15 @@ def _validate_field_value(name: str, value: Any, field: FieldSchema) -> None: def _build_doc(doc: dict[str, Any], schema: CollectionSchema) -> SdkDoc: - doc_id = _ensure_id(doc) + found_id, columns = split_pk(doc, schema=schema) + doc_id = found_id if found_id is not None else str(ULID()) vec_defs = {v.name: v for v in schema.vectors} field_map = {f.name: f for f in schema.fields} # ``Any`` mirrors the SDK's invariant value type (list[float]/list[int]/ # ndarray/sparse dict) so mypy does not complain about narrowing. vectors: dict[str, Any] = {} fields: dict[str, Any] = {} - for k, v in doc.items(): - if k == "id": - continue + for k, v in columns.items(): if k in vec_defs: vec_def = vec_defs[k] if _is_sparse_vector_type(vec_def.dataType): @@ -742,22 +791,20 @@ def _build_doc_partial(doc: dict[str, Any], schema: CollectionSchema) -> SdkDoc: Unlike :func:`_build_doc`, this does **not** require every vector field to be present; only the columns the caller is changing are validated. - The ``id`` is mandatory (the SDK rejects updates without an id). + The primary key is mandatory (the SDK rejects updates without an id). """ - raw_id = doc.get("id") - if not isinstance(raw_id, str) or not raw_id: + found_id, columns = split_pk(doc, schema=schema) + if not found_id: raise InvalidSchemaError( - "Update payload requires an explicit string 'id'.", - extra={"id": raw_id}, + f"Update payload requires an explicit string '{pk_key(schema)}'.", + extra={"primaryKeyKey": pk_key(schema)}, ) vec_defs = {v.name: v for v in schema.vectors} field_map = {f.name: f for f in schema.fields} # See note on ``_build_doc`` -- ``Any`` matches the SDK's invariant value type. vectors: dict[str, Any] = {} fields: dict[str, Any] = {} - for k, v in doc.items(): - if k == "id": - continue + for k, v in columns.items(): if k in vec_defs: vec_def = vec_defs[k] if _is_sparse_vector_type(vec_def.dataType): @@ -782,7 +829,7 @@ def _build_doc_partial(doc: dict[str, Any], schema: CollectionSchema) -> SdkDoc: f"Unknown column '{k}' in document.", extra={"column": k}, ) - return SdkDoc(id=raw_id, fields=fields or None, vectors=vectors) + return SdkDoc(id=found_id, fields=fields or None, vectors=vectors) def _status_ok(status: Any) -> bool: @@ -811,6 +858,100 @@ def _status_ok(status: Any) -> bool: return True +def _status_msg(status: Any) -> str: + """Best-effort human-readable message of a Zvec ``Status``. + + ``Status.message()`` is a method on the pybind11 binding; some fallback + shapes expose it differently (or not at all). + """ + getter = getattr(status, "message", None) + if callable(getter): + try: + return str(getter()) + except Exception: + pass + try: + return str(status) + except Exception: + return "" + + +def _write_in_batches( + sdk_docs: list[SdkDoc], + write_fn: Callable[[list[SdkDoc]], Any], + *, + failed_code: str, + batch_size: int = _DEFAULT_WRITE_BATCH, +) -> list[Any]: + """Run ``write_fn`` in SDK-sized chunks and classify failures. + + Absorbs the SDK's hard batch limit (:data:`_MAX_SDK_WRITE_BATCH`) so + callers may pass any number of documents. Per-document failures are + reported through the SDK ``Status`` list; this helper maps them onto the + Studio error taxonomy: + + * duplicate primary key -> :class:`DocumentConflictError` (409) + * unknown id (update) -> :class:`DocumentNotFoundError` (404) + * anything else -> :class:`ZvecStudioError` (*failed_code*) + + A ``ValueError`` raised *by the SDK itself* (document-level validation) + becomes :class:`InvalidDocumentError` (422); if it still mentions the + batch limit, our own batching has regressed and it must stay a 500. + """ + statuses: list[Any] = [] + chunk = max(1, min(batch_size, _MAX_SDK_WRITE_BATCH)) + try: + for start in range(0, len(sdk_docs), chunk): + result = write_fn(sdk_docs[start : start + chunk]) + if not isinstance(result, list): + result = [result] + statuses.extend(result) + except ValueError as exc: + msg = _exc_msg(exc) + if "too many docs" in msg.lower(): + raise ZvecStudioError( + f"Internal write-batch regression: {msg}", + code="INTERNAL_ERROR", + ) from exc + raise InvalidDocumentError(msg, sdk_exception="ValueError") from exc + _check_write_statuses(statuses, failed_code=failed_code) + return statuses + + +def _check_write_statuses(statuses: list[Any], *, failed_code: str) -> None: + """Raise the first non-OK ``Status``, mapped to the closest HTTP error. + + Classification is shared with the import row reports via + :func:`_classify_status_code`; only the opaque-failure code differs + (callers pass their own operation code, e.g. ``INSERT_FAILED``). + """ + for s in statuses: + if _status_ok(s): + continue + msg = _status_msg(s) + code = _classify_status_code(msg) + if code == "DOCUMENT_CONFLICT": + raise DocumentConflictError( + msg or "Document already exists.", sdk_exception="Status" + ) + if code == "DOCUMENT_NOT_FOUND": + raise DocumentNotFoundError(msg or "Document not found.", sdk_exception="Status") + raise ZvecStudioError( + f"Zvec write returned non-zero status: {msg or s}", code=failed_code + ) + + +def _classify_status_code(message: str) -> str: + """Map a non-OK ``Status`` message onto a Studio error code (for import + row reports, which must not raise).""" + lowered = message.lower() + if "already exists" in lowered: + return "DOCUMENT_CONFLICT" + if "not found" in lowered: + return "DOCUMENT_NOT_FOUND" + return "WRITE_FAILED" + + class SdkBackend: """Zvec-backed Collection backend (only backend in v0.2.0). @@ -856,11 +997,19 @@ def stats(self, name: str, *, path: str | None = None) -> CollectionStats: def flush(self, name: str, *, path: str | None = None) -> None: record = self.get(name, path=path) - record.sdk_obj.flush() + try: + record.sdk_obj.flush() + except RuntimeError as exc: + _raise_if_maintenance_blocked(exc) + raise def optimize(self, name: str, *, path: str | None = None) -> None: record = self.get(name, path=path) - record.sdk_obj.optimize(zvec.OptimizeOption()) + try: + record.sdk_obj.optimize(zvec.OptimizeOption()) + except RuntimeError as exc: + _raise_if_maintenance_blocked(exc) + raise def destroy(self, name: str, *, path: str | None = None) -> None: with self._lock: @@ -907,6 +1056,7 @@ def add_field( sdk_field, expression, zvec.AddColumnOption() ) except (ValueError, RuntimeError) as exc: + _raise_if_maintenance_blocked(exc) raise InvalidSchemaError( f"Zvec add_column failed: {_exc_msg(exc)}", extra={"name": name, "column": field.name}, @@ -924,6 +1074,7 @@ def drop_field(self, name: str, field_name: str, *, path: str | None = None) -> try: record.sdk_obj.drop_column(field_name) except (ValueError, RuntimeError) as exc: + _raise_if_maintenance_blocked(exc) raise InvalidSchemaError( f"Zvec drop_column failed: {_exc_msg(exc)}", extra={"name": name, "column": field_name}, @@ -953,6 +1104,7 @@ def rename_field( old_name, new_name=new_name, option=zvec.AlterColumnOption() ) except (ValueError, RuntimeError) as exc: + _raise_if_maintenance_blocked(exc) raise InvalidSchemaError( f"Zvec alter_column failed: {_exc_msg(exc)}", extra={"name": name, "oldName": old_name, "newName": new_name}, @@ -983,6 +1135,7 @@ def create_index( vector_field, sdk_param, zvec.IndexOption() ) except (ValueError, RuntimeError) as exc: + _raise_if_maintenance_blocked(exc) raise InvalidSchemaError( f"Zvec create_index failed: {_exc_msg(exc)}", extra={"name": name, "vectorField": vector_field}, @@ -1001,6 +1154,7 @@ def drop_index(self, name: str, vector_field: str, *, path: str | None = None) - try: record.sdk_obj.drop_index(vector_field) except (ValueError, RuntimeError) as exc: + _raise_if_maintenance_blocked(exc) raise InvalidSchemaError( f"Zvec drop_index failed: {_exc_msg(exc)}", extra={"name": name, "vectorField": vector_field}, @@ -1038,6 +1192,7 @@ def create_scalar_index( param = _build_scalar_index_param(index_param) record.sdk_obj.create_index(field_name, param, zvec.IndexOption()) except (ValueError, RuntimeError) as exc: + _raise_if_maintenance_blocked(exc) raise InvalidSchemaError( f"Zvec create_index (scalar) failed: {_exc_msg(exc)}", extra={"name": name, "field": field_name}, @@ -1055,6 +1210,7 @@ def drop_scalar_index(self, name: str, field_name: str, *, path: str | None = No try: record.sdk_obj.drop_index(field_name) except (ValueError, RuntimeError) as exc: + _raise_if_maintenance_blocked(exc) raise InvalidSchemaError( f"Zvec drop_index (scalar) failed: {_exc_msg(exc)}", extra={"name": name, "field": field_name}, @@ -1064,6 +1220,23 @@ def drop_scalar_index(self, name: str, field_name: str, *, path: str | None = No # ---- documents ---- + def _ensure_unique_name(self, name: str, path: Path) -> None: + """Enforce the registry invariant: open collection names are unique. + + Nearly every document API resolves a collection by name alone (the + optional ``path`` hint is a fast-path), so two same-named collections + can never be open at once — requests would silently resolve to + whichever registered first. Callers must close the blocking one or + pick another name. + """ + for record in self._by_path.values(): + if record.name == name and record.path != path: + raise CollectionAlreadyExistsError( + f"Collection '{name}' is already open from '{record.path}'. " + f"Close it first, or choose another name.", + extra={"name": name, "conflictingPath": str(record.path)}, + ) + def create(self, *, path: Path, schema: CollectionSchema) -> CollectionRecord: path = Path(path) with self._lock: @@ -1084,6 +1257,7 @@ def create(self, *, path: Path, schema: CollectionSchema) -> CollectionRecord: f"Parent directory {path.parent} does not exist.", extra={"path": str(path)}, ) + self._ensure_unique_name(schema.name, path) sdk_schema = _to_sdk_schema(schema) try: sdk_obj = zvec.create_and_open(str(path), sdk_schema) @@ -1129,6 +1303,7 @@ def open(self, path: Path) -> CollectionRecord: extra={"path": str(path)}, ) from exc schema = _from_sdk_schema(sdk_obj.schema, path) + self._ensure_unique_name(schema.name, path) record = CollectionRecord( name=schema.name, path=path, @@ -1159,15 +1334,9 @@ def close(self, name: str, *, path: str | None = None) -> CollectionRecord: def insert_documents(self, name: str, docs: list[dict[str, Any]]) -> list[str]: record = self.get(name) sdk_docs = [_build_doc(d, record.schema) for d in docs] - statuses = record.sdk_obj.insert(sdk_docs) - if not isinstance(statuses, list): - statuses = [statuses] - for s in statuses: - if not _status_ok(s): - raise ZvecStudioError( - f"Zvec insert returned non-zero status: {s}", - code="INSERT_FAILED", - ) + # The SDK rejects write batches above _MAX_SDK_WRITE_BATCH, so split + # internally (the HTTP contract accepts up to 10,000 documents). + _write_in_batches(sdk_docs, record.sdk_obj.insert, failed_code="INSERT_FAILED") record.sdk_obj.flush() return [d.id for d in sdk_docs] @@ -1179,34 +1348,38 @@ def upsert_documents(self, name: str, docs: list[dict[str, Any]]) -> list[str]: For new ids, the document is inserted as-is. Implementation: uses ``update()`` for existing docs (partial merge) - and falls back to ``insert()`` for new docs. + and falls back to ``insert()`` for new docs. The primary key is + resolved through :mod:`zvec_studio.storage.doc_repr`, so a schema that + declares its own ``id`` column (primary key carried on ``$id``) still + merges by the real key, and ambiguous bare-``id`` rows are rejected. """ record = self.get(name) - # Split into docs with explicit id vs without - with_id: list[dict[str, Any]] = [] + # Split into docs with an explicit primary key vs without (split_pk + # also validates ambiguous rows before any write happens). + keyed: list[tuple[str, dict[str, Any]]] = [] without_id: list[dict[str, Any]] = [] for d in docs: - if d.get("id"): - with_id.append(d) + found_id, _ = split_pk(d, schema=record.schema) + if found_id: + keyed.append((found_id, d)) else: without_id.append(d) result_ids: list[str] = [] # For docs with id: try update (partial merge), fallback to insert - if with_id: + if keyed: # Check which ids already exist - ids_to_check = [str(d["id"]) for d in with_id] - fetched = record.sdk_obj.fetch(ids_to_check) + fetched = record.sdk_obj.fetch([doc_id for doc_id, _ in keyed]) existing_ids: set[str] = set() if isinstance(fetched, dict): existing_ids = {k for k, v in fetched.items() if v is not None} to_update: list[dict[str, Any]] = [] to_insert: list[dict[str, Any]] = [] - for d in with_id: - if str(d["id"]) in existing_ids: + for doc_id, d in keyed: + if doc_id in existing_ids: to_update.append(d) else: to_insert.append(d) @@ -1214,43 +1387,25 @@ def upsert_documents(self, name: str, docs: list[dict[str, Any]]) -> list[str]: # Update existing docs (partial — only provided fields change) if to_update: sdk_update_docs = [_build_doc_partial(d, record.schema) for d in to_update] - statuses = record.sdk_obj.update(sdk_update_docs) - if not isinstance(statuses, list): - statuses = [statuses] - for s in statuses: - if not _status_ok(s): - raise ZvecStudioError( - f"Zvec update returned non-zero status: {s}", - code="UPSERT_FAILED", - ) + _write_in_batches( + sdk_update_docs, record.sdk_obj.update, failed_code="UPSERT_FAILED" + ) result_ids.extend(d.id for d in sdk_update_docs) # Insert new docs with explicit id (full doc required) if to_insert: sdk_insert_docs = [_build_doc(d, record.schema) for d in to_insert] - statuses = record.sdk_obj.insert(sdk_insert_docs) - if not isinstance(statuses, list): - statuses = [statuses] - for s in statuses: - if not _status_ok(s): - raise ZvecStudioError( - f"Zvec insert returned non-zero status: {s}", - code="UPSERT_FAILED", - ) + _write_in_batches( + sdk_insert_docs, record.sdk_obj.insert, failed_code="UPSERT_FAILED" + ) result_ids.extend(d.id for d in sdk_insert_docs) # Insert docs without id (auto-generate ULID) if without_id: sdk_new_docs = [_build_doc(d, record.schema) for d in without_id] - statuses = record.sdk_obj.insert(sdk_new_docs) - if not isinstance(statuses, list): - statuses = [statuses] - for s in statuses: - if not _status_ok(s): - raise ZvecStudioError( - f"Zvec insert returned non-zero status: {s}", - code="UPSERT_FAILED", - ) + _write_in_batches( + sdk_new_docs, record.sdk_obj.insert, failed_code="UPSERT_FAILED" + ) result_ids.extend(d.id for d in sdk_new_docs) record.sdk_obj.flush() @@ -1259,33 +1414,325 @@ def upsert_documents(self, name: str, docs: list[dict[str, Any]]) -> list[str]: def update_documents(self, name: str, docs: list[dict[str, Any]]) -> list[str]: record = self.get(name) sdk_docs = [_build_doc_partial(d, record.schema) for d in docs] - statuses = record.sdk_obj.update(sdk_docs) - if not isinstance(statuses, list): - statuses = [statuses] - for s, d in zip(statuses, sdk_docs, strict=False): - if _status_ok(s): + # ``update`` reports unknown ids through a per-doc ``Status`` whose + # message contains "not found"; _check_write_statuses maps that to + # DocumentNotFoundError (404) instead of an opaque 5xx. + _write_in_batches(sdk_docs, record.sdk_obj.update, failed_code="UPDATE_FAILED") + record.sdk_obj.flush() + return [d.id for d in sdk_docs] + + def import_documents( + self, + name: str, + *, + source_path: str, + fmt: ImportFormat, + mode: ImportMode = ImportMode.REPLACE, + on_error: OnErrorMode = OnErrorMode.ABORT, + batch_size: int = _DEFAULT_WRITE_BATCH, + path: str | None = None, + ) -> ImportReport: + """Import documents from a file, streaming. + + Pipeline: parse row -> build ``SdkDoc`` (schema-validated) -> write in + SDK-sized batches (``upsert`` for ``replace``, ``insert`` for + ``insert``) -> aggregate a per-row report. Memory stays bounded by one + batch regardless of file size; the collection is flushed once at the + end (including aborts, so partially imported rows are durable). + + Sources ending in ``.tar.gz`` / ``.tgz`` are treated as snapshot + packages: the embedded ``manifest.json`` is parsed and schema-checked + *before any row is written* (fail-fast 409), then the embedded JSONL + member is consumed through the same pipeline. + + ``path`` fast-paths collection lookup (the registry guarantees open + names are unique, so it is never needed for disambiguation — handy + all the same when a caller already holds the absolute path). + """ + record = self.get(name, path=path) + source = validate_import_source(source_path) + chunk = max(1, min(batch_size, _MAX_SDK_WRITE_BATCH)) + report = ImportReport() + started = time.perf_counter() + + with contextlib.ExitStack() as stack: + parser = self._open_import_stream(stack, source, record.schema, fmt) + aborted = self._consume_rows( + record, parser, mode=mode, on_error=on_error, chunk=chunk, report=report + ) + + # One flush for the whole run (see design doc §6.4: per-batch flushing + # would pay needless segment-seal costs). + record.sdk_obj.flush() + report.aborted = aborted + report.duration_ms = (time.perf_counter() - started) * 1000 + return report + + def _open_import_stream( + self, + stack: contextlib.ExitStack, + source: Path, + schema: CollectionSchema, + fmt: ImportFormat, + ) -> Iterator[tuple[int, dict[str, Any]]]: + """Open the row stream for plain JSONL files or snapshot packages. + + A ``*.tar.gz`` source that is not actually a readable gzip/tar stream + (or breaks mid-read) is a user input error, mapped to 400 + ``IMPORT_MANIFEST_INVALID`` instead of leaking tarfile/gzip internals + as a 500. + """ + if source.name.lower().endswith((".tar.gz", ".tgz")): + try: + tar = stack.enter_context(tarfile.open(source, mode="r:gz")) # noqa: SIM115 + members = {m.name: m for m in tar if m.isfile()} + except (tarfile.TarError, EOFError, OSError) as exc: + raise ImportManifestInvalidError( + f"'{source.name}' is not a readable snapshot package: {exc}", + extra={"path": str(source)}, + ) from exc + if MANIFEST_NAME not in members or DATA_FILE_NAME not in members: + raise ImportManifestInvalidError( + f"Snapshot must contain '{MANIFEST_NAME}' and '{DATA_FILE_NAME}'.", + extra={"members": sorted(members)}, + ) + manifest_file = tar.extractfile(members[MANIFEST_NAME]) + if manifest_file is None: + raise ImportManifestInvalidError( + f"'{MANIFEST_NAME}' is not a readable file inside the snapshot." + ) + try: + manifest = parse_manifest(manifest_file.read()) + except (tarfile.TarError, EOFError, OSError) as exc: + raise ImportManifestInvalidError( + f"'{MANIFEST_NAME}' inside '{source.name}' is corrupted: {exc}", + extra={"path": str(source)}, + ) from exc + check_schema_compatible(manifest, schema) + data_file = tar.extractfile(members[DATA_FILE_NAME]) + if data_file is None: + raise ImportManifestInvalidError( + f"'{DATA_FILE_NAME}' is not a readable file inside the snapshot." + ) + return self._guarded_snapshot_rows(fmt.parse(cast(BinaryIO, data_file)), source) + stream = stack.enter_context(source.open("rb")) + return fmt.parse(stream) + + @staticmethod + def _guarded_snapshot_rows( + rows: Iterator[tuple[int, dict[str, Any]]], source: Path + ) -> Iterator[tuple[int, dict[str, Any]]]: + """Re-raise mid-stream gzip/tar corruption as a 400 manifest error. + + Row-level parse failures (``InvalidDocumentError``) pass through + untouched — only archive-level breakage is reclassified. + """ + try: + yield from rows + except (tarfile.TarError, EOFError, OSError) as exc: + raise ImportManifestInvalidError( + f"'{source.name}' is corrupted and cannot be read to the end: {exc}", + extra={"path": str(source)}, + ) from exc + + def _consume_rows( + self, + record: CollectionRecord, + parser: Iterator[tuple[int, dict[str, Any]]], + *, + mode: ImportMode, + on_error: OnErrorMode, + chunk: int, + report: ImportReport, + ) -> bool: + """Build rows into batches, write them, and report; return aborted.""" + batch: list[SdkDoc] = [] + batch_lines: list[int] = [] + aborted = False + + def flush_batch() -> bool: + """Write the pending batch; return True to abort the whole import.""" + nonlocal batch, batch_lines + if not batch: + return False + abort_now = not self._import_write_batch( + record, + batch, + batch_lines, + mode=mode, + on_error=on_error, + report=report, + ) + batch, batch_lines = [], [] + return abort_now + + while True: + try: + line_number, row = next(parser) + except StopIteration: + break + except ZvecStudioError as exc: + # Format-level row error (invalid JSON, wrong shape, ...). + raw_line = exc.extra.get("line", 0) + report.add_failure( + ImportFailure( + line=int(raw_line) if isinstance(raw_line, int) else 0, + code=exc.code, + message=exc.message, + ) + ) + if on_error is OnErrorMode.ABORT: + aborted = True + break continue - # Zvec returns ``code=1, message='Document not found'`` for unknown - # ids in update(). Surface that as 404 so the UI can react sensibly - # rather than treating it as an opaque 5xx. - msg = "" - getter = getattr(s, "message", None) - if callable(getter): - try: - msg = str(getter()) - except Exception: - msg = "" - if "not found" in msg.lower(): - raise DocumentNotFoundError( - f"Document '{d.id}' not found in '{name}'.", - extra={"collection": name, "id": d.id}, + + report.total_lines += 1 + try: + batch.append(_build_doc(row, record.schema)) + except ZvecStudioError as exc: + report.add_failure( + ImportFailure(line=line_number, code=exc.code, message=exc.message) ) - raise ZvecStudioError( - f"Zvec update returned non-zero status: {s}", - code="UPDATE_FAILED", + if on_error is OnErrorMode.ABORT: + aborted = True + break + continue + batch_lines.append(line_number) + + if len(batch) >= chunk and flush_batch(): + aborted = True + break + + # Write whatever valid rows are still buffered. After an abort the + # buffer holds only rows validated *before* the failing one, and + # "stop at the first error" must not silently discard them. + if flush_batch(): + aborted = True + return aborted + + def _import_write_batch( + self, + record: CollectionRecord, + sdk_docs: list[SdkDoc], + lines: list[int], + *, + mode: ImportMode, + on_error: OnErrorMode, + report: ImportReport, + ) -> bool: + """Write one batch; return True when the import should continue. + + A ``ValueError`` from the SDK means the whole batch was rejected at + validation time; either way we retry row by row so the valid rows in + the batch still land — under ``skip`` to locate every offending line, + under ``abort`` to keep the rows before the first failing one + (consistent with the other abort paths, which never discard rows that + were already validated). + """ + write_fn = ( + record.sdk_obj.upsert if mode is ImportMode.REPLACE else record.sdk_obj.insert + ) + try: + statuses = write_fn(sdk_docs) + if not isinstance(statuses, list): + statuses = [statuses] + except ValueError as exc: + msg = _exc_msg(exc) + if "too many docs" in msg.lower(): + # Batching regression — never disguise as user input. + raise ZvecStudioError( + f"Internal write-batch regression: {msg}", code="INTERNAL_ERROR" + ) from exc + return self._import_write_one_by_one( + record, + sdk_docs, + lines, + mode=mode, + report=report, + stop_at_first_failure=on_error is OnErrorMode.ABORT, ) - record.sdk_obj.flush() - return [d.id for d in sdk_docs] + + failed_in_batch = False + first_failure_recorded = False + for i, s in enumerate(statuses): + if _status_ok(s): + report.imported += 1 + continue + failed_in_batch = True + # ``abort`` reports only the first failing row (the import stops); + # ``skip`` records all of them. + if on_error is OnErrorMode.SKIP or not first_failure_recorded: + msg = _status_msg(s) + report.add_failure( + ImportFailure( + line=lines[i], code=_classify_status_code(msg), message=msg + ) + ) + first_failure_recorded = True + return not (failed_in_batch and on_error is OnErrorMode.ABORT) + + def _import_write_one_by_one( + self, + record: CollectionRecord, + sdk_docs: list[SdkDoc], + lines: list[int], + *, + mode: ImportMode, + report: ImportReport, + stop_at_first_failure: bool = False, + ) -> bool: + """Row-by-row fallback after a batch-level ValueError. + + ``skip`` uses it to locate every offending line (and keep good rows); + ``abort`` uses it with ``stop_at_first_failure`` so the rows before + the first failing one still land. Returns False to abort the import. + """ + write_fn = ( + record.sdk_obj.upsert if mode is ImportMode.REPLACE else record.sdk_obj.insert + ) + for doc, line in zip(sdk_docs, lines, strict=True): + try: + statuses = write_fn([doc]) + status = statuses[0] if isinstance(statuses, list) else statuses + except ValueError as exc: + report.add_failure( + ImportFailure(line=line, code="INVALID_DOCUMENT", message=_exc_msg(exc)) + ) + if stop_at_first_failure: + return False + continue + if _status_ok(status): + report.imported += 1 + else: + msg = _status_msg(status) + report.add_failure( + ImportFailure(line=line, code=_classify_status_code(msg), message=msg) + ) + if stop_at_first_failure: + return False + return True + + def iter_documents( + self, + name: str, + *, + include_vector: bool, + output_fields: list[str] | None = None, + path: str | None = None, + ) -> Iterator[dict[str, Any]]: + """Stream every document as a wire row (design doc §6.3). + + Wraps the SDK snapshot iterator (zvec 0.7 ``iter_docs``) in ``with`` so + the iterator — which blocks maintenance operations while open — is + released on EVERY exit path: normal exhaustion, early ``close()`` + (client disconnect), and exceptions mid-stream. + """ + record = self.get(name, path=path) + with record.sdk_obj.iter_docs( + output_fields=output_fields, include_vector=include_vector + ) as iterator: + for doc in iterator: + yield doc_to_row(doc, schema=record.schema, include_vector=include_vector) def get_document(self, name: str, doc_id: str) -> dict[str, Any]: record = self.get(name) @@ -1297,7 +1744,7 @@ def get_document(self, name: str, doc_id: str) -> dict[str, Any]: f"Document '{sid}' not found in '{name}'.", extra={"collection": name, "id": sid}, ) - return _doc_to_dict(doc, include_vector=True) + return _doc_to_dict(doc, schema=record.schema, include_vector=True) def delete_document(self, name: str, doc_id: str) -> None: record = self.get(name) @@ -1319,6 +1766,8 @@ def delete_documents(self, name: str, ids: list[str]) -> int: present = record.sdk_obj.fetch(sids) deleted_ids = list(present.keys()) if isinstance(present, dict) else [] if deleted_ids: + # No batching needed here: the write-batch limit is insert/update + # specific (a 2048-id delete was verified to succeed in one call). record.sdk_obj.delete(deleted_ids) record.sdk_obj.flush() return len(deleted_ids) @@ -1379,7 +1828,10 @@ def browse( str(exc), extra={"filter": filter_expr}, ) from exc - return [_doc_to_dict(d, include_vector=include_vector) for d in docs] + return [ + _doc_to_dict(d, schema=record.schema, include_vector=include_vector) + for d in docs + ] # ---- search ---- @@ -1434,7 +1886,7 @@ def search( ( d.id, float(d.score) if d.score is not None else 0.0, - _doc_to_dict(d, include_vector=include_vector), + _doc_to_dict(d, schema=record.schema, include_vector=include_vector), ) for d in docs ] @@ -1521,7 +1973,7 @@ def group_by_search( ( doc.id, float(doc.score) if doc.score is not None else 0.0, - _doc_to_dict(doc, include_vector=include_vector), + _doc_to_dict(doc, schema=record.schema, include_vector=include_vector), str(group.group_by_value), ) for group in groups diff --git a/apps/backend/zvec_studio/storage/snapshot.py b/apps/backend/zvec_studio/storage/snapshot.py new file mode 100644 index 0000000..389a53b --- /dev/null +++ b/apps/backend/zvec_studio/storage/snapshot.py @@ -0,0 +1,325 @@ +"""Snapshot packaging: ``manifest.json`` + ``documents.jsonl``. + +A snapshot export bundles the collection schema and export options in a +``manifest.json`` next to the JSONL data (design doc §6.2.2). On import the +manifest is read *before any row is written* so schema incompatibility fails +fast (409) instead of surfacing mid-file. + +Naming / versioning rules: + +* ``MANIFEST_FORMAT`` is ``zvec-studio.export/``. A major bump means + incompatible manifest changes; parsers reject unknown majors. +* No ``docCount`` is recorded: streaming generation cannot know it up front, + and the importer counts actual rows anyway (design doc §6.2.2). +""" + +from __future__ import annotations + +import queue +import tarfile +import threading +from collections.abc import Iterator +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, BinaryIO, cast + +import orjson +import zvec + +from zvec_studio.__about__ import __version__ as STUDIO_VERSION +from zvec_studio.exceptions import ( + ImportManifestInvalidError, + ImportSchemaMismatchError, +) +from zvec_studio.schemas.collection import CollectionSchema + +#: Member file names inside the snapshot package. +MANIFEST_NAME = "manifest.json" +DATA_FILE_NAME = "documents.jsonl" + +#: Format identifier written into the manifest (major version = contract). +MANIFEST_FORMAT = "zvec-studio.export/1" +_MANIFEST_PREFIX = "zvec-studio.export/" +_SUPPORTED_MAJOR = "1" + + +def build_manifest( + *, + schema: CollectionSchema, + include_vector: bool, + output_fields: list[str] | None, +) -> dict[str, Any]: + """Build the manifest document for a snapshot export.""" + return { + "format": MANIFEST_FORMAT, + "exportedAt": datetime.now(timezone.utc).isoformat(timespec="seconds"), + "zvecVersion": getattr(zvec, "__version__", "unknown"), + "studioVersion": STUDIO_VERSION, + "collection": { + "name": schema.name, + "schema": schema.model_dump(mode="json"), + }, + "options": { + "includeVector": include_vector, + "outputFields": output_fields, + }, + "data": { + "file": DATA_FILE_NAME, + "format": "jsonl", + "encoding": "utf-8", + }, + } + + +def parse_manifest(raw: bytes) -> dict[str, Any]: + """Parse and structurally validate a manifest. + + Raises: + ImportManifestInvalidError: unreadable JSON, wrong shape, unknown + format major, or missing required members. + """ + try: + manifest = orjson.loads(raw) + except orjson.JSONDecodeError as exc: + raise ImportManifestInvalidError( + f"manifest.json is not valid JSON: {exc}" + ) from exc + if not isinstance(manifest, dict): + raise ImportManifestInvalidError("manifest.json must be a JSON object.") + + fmt = manifest.get("format") + if not isinstance(fmt, str) or not fmt.startswith(_MANIFEST_PREFIX): + raise ImportManifestInvalidError( + f"manifest.json carries an unrecognized format: {fmt!r}." + ) + major = fmt.removeprefix(_MANIFEST_PREFIX) + if major != _SUPPORTED_MAJOR: + raise ImportManifestInvalidError( + f"manifest.json format major '{major}' is not supported " + f"(supported: {_SUPPORTED_MAJOR}). Upgrade Studio." + ) + + collection = manifest.get("collection") + if not isinstance(collection, dict) or not isinstance(collection.get("schema"), dict): + raise ImportManifestInvalidError( + "manifest.json is missing 'collection.schema'." + ) + try: + CollectionSchema.model_validate(collection["schema"]) + except ValueError as exc: + raise ImportManifestInvalidError( + f"manifest.json schema is invalid: {exc}" + ) from exc + return manifest + + +def read_snapshot_manifest(path: Path) -> dict[str, Any]: + """Open a snapshot package and return its parsed manifest. + + Used by collection restore (which needs the schema before creating + anything). Raises the same 400-class errors the import path reports for + unreadable / malformed packages. + """ + try: + with tarfile.open(path, mode="r:gz") as tar: + members = {m.name: m for m in tar if m.isfile()} + if MANIFEST_NAME not in members: + raise ImportManifestInvalidError( + f"Snapshot must contain '{MANIFEST_NAME}'.", + extra={"members": sorted(members)}, + ) + manifest_file = tar.extractfile(members[MANIFEST_NAME]) + if manifest_file is None: + raise ImportManifestInvalidError( + f"'{MANIFEST_NAME}' is not a readable file inside the snapshot." + ) + return parse_manifest(manifest_file.read()) + except (tarfile.TarError, EOFError, OSError) as exc: + raise ImportManifestInvalidError( + f"'{path.name}' is not a readable snapshot package: {exc}", + extra={"path": str(path)}, + ) from exc + + +def check_schema_compatible(manifest: dict[str, Any], target: CollectionSchema) -> None: + """Verify the snapshot can be imported into *target*; raise otherwise. + + Rules (all collected, reported together): + * every vector in the snapshot must exist in the target with the same + dtype and dimension (indexes are rebuilt, not compared); + * every scalar field in the snapshot must exist in the target with the + same dtype (the target may have extra fields — a pruned export is fine); + fields dropped by ``options.outputFields`` never appear in the data, so + they are exempt from this check; + * if the target declares vectors but the snapshot was exported without + them, repopulating is impossible. + + Raises: + ImportSchemaMismatchError: with ``extra['mismatches']`` listing all + incompatibilities. + """ + source = CollectionSchema.model_validate(manifest["collection"]["schema"]) + options = manifest.get("options") or {} + mismatches: list[str] = [] + + target_vectors = {v.name: v for v in target.vectors} + for v in source.vectors: + tv = target_vectors.get(v.name) + if tv is None: + mismatches.append(f"Vector '{v.name}' is not defined on the target collection.") + elif tv.dataType != v.dataType: + mismatches.append( + f"Vector '{v.name}' dtype differs: snapshot {v.dataType.value}, " + f"target {tv.dataType.value}." + ) + elif v.dimension != tv.dimension: + mismatches.append( + f"Vector '{v.name}' dimension differs: snapshot {v.dimension}, " + f"target {tv.dimension}." + ) + + # ``outputFields`` prunes columns out of the data rows; only the kept + # columns need to line up with the target schema. + raw_output_fields = options.get("outputFields") + pruned: set[str] | None = None + if isinstance(raw_output_fields, list) and raw_output_fields: + pruned = {str(name) for name in raw_output_fields} + + target_fields = {f.name: f for f in target.fields} + for f in source.fields: + if pruned is not None and f.name not in pruned: + continue + tf = target_fields.get(f.name) + if tf is None: + mismatches.append(f"Field '{f.name}' is not defined on the target collection.") + elif tf.dataType != f.dataType: + mismatches.append( + f"Field '{f.name}' dtype differs: snapshot {f.dataType.value}, " + f"target {tf.dataType.value}." + ) + + if target.vectors and not options.get("includeVector", True): + mismatches.append( + "The snapshot was exported without vectors, but the target " + "collection requires them." + ) + + if mismatches: + raise ImportSchemaMismatchError( + "Snapshot schema is not compatible with the target collection: " + + " ".join(mismatches), + extra={"mismatches": mismatches}, + ) + + +# --------------------------------------------------------------------------- +# Package assembly / streaming +# --------------------------------------------------------------------------- + + +def write_snapshot_package( + *, + rows: Iterator[dict[str, Any]], + serialize: Any, + manifest: dict[str, Any], + tmp_dir: Path, +) -> tuple[Path, Path]: + """Materialize ``documents.jsonl`` + ``manifest.json`` into *tmp_dir*. + + Streaming caveat (see design doc §5.7): a tar member's header must carry + its size up front, and the JSONL size is unknown until fully serialized. + The data is therefore staged on disk first (temp space ~= data size), + then tar-streamed without holding it in memory. Returns + ``(manifest_path, data_path)``. + """ + data_path = tmp_dir / DATA_FILE_NAME + with data_path.open("wb") as out: + for chunk in serialize(rows): + out.write(chunk) + manifest_path = tmp_dir / MANIFEST_NAME + manifest_path.write_bytes(orjson.dumps(manifest, option=orjson.OPT_INDENT_2)) + return manifest_path, data_path + + +class _QueueWriter: + """File-like adapter feeding a producer thread's bytes into a queue. + + ``cancelled`` lets the consumer (an aborted HTTP download) wake a + producer that is blocked pushing into the bounded queue — without it the + thread would sit on ``q.put`` until process exit. + """ + + def __init__(self, q: queue.Queue[bytes | None], cancelled: threading.Event) -> None: + self._q = q + self._cancelled = cancelled + + def write(self, data: bytes) -> int: + if data: + payload = bytes(data) + while True: + if self._cancelled.is_set(): + raise RuntimeError("snapshot export cancelled") + try: + # A plain ``put`` could block forever on a full queue once the + # consumer is gone; poll so cancellation wins promptly. + self._q.put(payload, timeout=0.1) + break + except queue.Full: + continue + return len(data) + + def close(self) -> None: # tarfile calls close(); the queue owns lifetime + pass + + +#: Worker thread name — tests use it to assert no producer leaks on abort. +SNAPSHOT_WORKER_NAME = "zvec-studio-snapshot-export" + + +def stream_snapshot_package( + *, manifest_path: Path, data_path: Path +) -> Iterator[bytes]: + """Yield a gzip tar (manifest first, then documents) without buffering it. + + ``tarfile`` writes synchronously, so compression runs on a worker thread + pushing chunks through a bounded queue; the generator consumes them, + keeping peak memory at a handful of chunks regardless of package size. + On early close (client disconnect) a cancellation event unblocks the + producer so the thread exits promptly instead of leaking. + """ + q: queue.Queue[bytes | None] = queue.Queue(maxsize=32) + errors: list[BaseException] = [] + cancelled = threading.Event() + + def produce() -> None: + try: + writer = cast(BinaryIO, _QueueWriter(q, cancelled)) + with tarfile.open(fileobj=writer, mode="w|gz") as tar: + tar.add(str(manifest_path), arcname=MANIFEST_NAME) + tar.add(str(data_path), arcname=DATA_FILE_NAME) + except BaseException as exc: # re-raised in consumer + if not cancelled.is_set(): + errors.append(exc) + finally: + # The sentinel is only needed by a live consumer; once cancelled, + # nobody drains the queue, so don't block on a full one. + while not cancelled.is_set(): + try: + q.put(None, timeout=0.1) + break + except queue.Full: + continue + + worker = threading.Thread(target=produce, name=SNAPSHOT_WORKER_NAME, daemon=True) + worker.start() + try: + while True: + chunk = q.get() + if chunk is None: + break + yield chunk + finally: + cancelled.set() + worker.join(timeout=5) + if errors: + raise errors[0] From 41fcd74f1034d7e62bba01415d6711b459ca7890 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=90=AA=E9=9C=84?= Date: Thu, 27 Aug 2026 17:48:38 +0800 Subject: [PATCH 2/7] feat(frontend): import/export dialogs, sidebar import entry, $id-aware tables MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Browse tab gains Import/Export dialogs: local JSONL import (write mode and error policy with row-level report), native streamed export download in data/snapshot modes (never fetch().blob()). - Sidebar "+" menu gains "Import Collection": rebuilds a whole collection from a snapshot package with an auto-prefilled target directory and optional rename; the in-collection import dialog stays data-file only (.jsonl/.ndjson). - Document tables across Browse/Query/Write follow the backend primary-key rules ($id when a column occupies id), including detail fetch and copies. - File picker lists files for imports; partial imports surface as warnings; full zh-CN coverage; api-client types regenerated. Signed-off-by: 琪霄 --- .../components/ui/DirectoryPickerDialog.css | 7 + .../ui/DirectoryPickerDialog.test.tsx | 4 +- .../src/components/ui/FilePickerDialog.tsx | 222 ++ apps/frontend/src/components/ui/index.ts | 2 + apps/frontend/src/features/collections/api.ts | 72 +- .../src/features/collections/hooks.ts | 69 +- .../src/features/collections/index.ts | 3 + .../src/features/documents/api.test.ts | 52 + apps/frontend/src/features/documents/api.ts | 48 +- apps/frontend/src/features/documents/hooks.ts | 19 + apps/frontend/src/features/documents/index.ts | 6 +- apps/frontend/src/i18n/en.json | 84 +- apps/frontend/src/i18n/zh.json | 82 +- apps/frontend/src/layouts/AppShell.css | 38 + apps/frontend/src/layouts/AppShell.test.tsx | 44 +- apps/frontend/src/layouts/AppShell.tsx | 262 +- apps/frontend/src/lib/fs-api.test.ts | 21 + apps/frontend/src/lib/fs-api.ts | 17 +- .../src/pages/collections/BrowseTab.test.tsx | 67 +- .../src/pages/collections/BrowseTab.tsx | 103 +- .../collections/CreateCollectionDialog.tsx | 10 +- .../src/pages/collections/DocumentsPanel.tsx | 71 +- .../ExportDocumentsDialog.test.tsx | 99 + .../collections/ExportDocumentsDialog.tsx | 161 + .../ImportCollectionDialog.test.tsx | 159 + .../collections/ImportCollectionDialog.tsx | 215 ++ .../collections/ImportDocumentsDialog.css | 62 + .../ImportDocumentsDialog.test.tsx | 129 + .../collections/ImportDocumentsDialog.tsx | 251 ++ .../src/pages/collections/MutateTab.tsx | 628 +++- .../src/pages/collections/QueryTab.tsx | 554 ++- .../pages/collections/SchemaDdlSection.tsx | 15 +- .../src/pages/collections/doc-repr.test.ts | 62 + .../src/pages/collections/doc-repr.ts | 42 + .../src/pages/collections/import-utils.ts | 21 + packages/api-client/src/schema.ts | 365 +- pnpm-lock.yaml | 3294 ++++++++++++++++- 37 files changed, 6724 insertions(+), 636 deletions(-) create mode 100644 apps/frontend/src/components/ui/FilePickerDialog.tsx create mode 100644 apps/frontend/src/features/documents/api.test.ts create mode 100644 apps/frontend/src/pages/collections/ExportDocumentsDialog.test.tsx create mode 100644 apps/frontend/src/pages/collections/ExportDocumentsDialog.tsx create mode 100644 apps/frontend/src/pages/collections/ImportCollectionDialog.test.tsx create mode 100644 apps/frontend/src/pages/collections/ImportCollectionDialog.tsx create mode 100644 apps/frontend/src/pages/collections/ImportDocumentsDialog.css create mode 100644 apps/frontend/src/pages/collections/ImportDocumentsDialog.test.tsx create mode 100644 apps/frontend/src/pages/collections/ImportDocumentsDialog.tsx create mode 100644 apps/frontend/src/pages/collections/doc-repr.test.ts create mode 100644 apps/frontend/src/pages/collections/doc-repr.ts create mode 100644 apps/frontend/src/pages/collections/import-utils.ts diff --git a/apps/frontend/src/components/ui/DirectoryPickerDialog.css b/apps/frontend/src/components/ui/DirectoryPickerDialog.css index d21f90a..a93e8b4 100644 --- a/apps/frontend/src/components/ui/DirectoryPickerDialog.css +++ b/apps/frontend/src/components/ui/DirectoryPickerDialog.css @@ -73,3 +73,10 @@ color: var(--zv-color-text-muted, #6b7280); font-size: 0.85rem; } + +.zv-dirpicker__entry-size { + margin-left: auto; + color: var(--zv-color-text-muted, #6b7280); + font-size: 0.75rem; + font-variant-numeric: tabular-nums; +} diff --git a/apps/frontend/src/components/ui/DirectoryPickerDialog.test.tsx b/apps/frontend/src/components/ui/DirectoryPickerDialog.test.tsx index f7027e3..c98e71f 100644 --- a/apps/frontend/src/components/ui/DirectoryPickerDialog.test.tsx +++ b/apps/frontend/src/components/ui/DirectoryPickerDialog.test.tsx @@ -41,13 +41,13 @@ function fixtures(): Record { path: '/', parent: null, home: '/home/bob', - entries: [{ name: 'tmp', path: '/tmp' }], + entries: [{ name: 'tmp', path: '/tmp', kind: 'dir', size: null }], }, '/tmp': { path: '/tmp', parent: '/', home: '/home/bob', - entries: [{ name: 'projects', path: '/tmp/projects' }], + entries: [{ name: 'projects', path: '/tmp/projects', kind: 'dir', size: null }], }, '/tmp/projects': { path: '/tmp/projects', diff --git a/apps/frontend/src/components/ui/FilePickerDialog.tsx b/apps/frontend/src/components/ui/FilePickerDialog.tsx new file mode 100644 index 0000000..e95d497 --- /dev/null +++ b/apps/frontend/src/components/ui/FilePickerDialog.tsx @@ -0,0 +1,222 @@ +/** + * FilePickerDialog: a modal browser for picking a *file* on the host + * filesystem (the import source). + * + * Built on the same ``GET /fs/list`` backend endpoint as the directory + * picker, but requests ``includeFiles=true`` so files appear alongside + * directories. Directories navigate on click; files select on click. + * An optional comma-separated ``extensions`` filter narrows the listed + * files (directories are never filtered, they are needed for navigation). + */ +import { useCallback, useEffect, useRef, useState } from 'react'; +import { useTranslation } from 'react-i18next'; + +import { Button } from './Button'; +import { Dialog } from './Dialog'; +import { Spinner } from './Spinner'; +import { useApiClient } from '@/lib/api-client-context'; +import { ApiError } from '@/lib/api-client'; +import { createFsApi, type FsListing } from '@/lib/fs-api'; + +import './DirectoryPickerDialog.css'; + +export interface FilePickerDialogProps { + readonly open: boolean; + readonly onClose: () => void; + readonly onSelect: (absolutePath: string) => void; + /** Optional initial directory. Defaults to the user's home. */ + readonly initialPath?: string; + /** Comma-separated extension filter, e.g. ``'.jsonl,.tar.gz'``. */ + readonly extensions?: string; + readonly title?: string; +} + +function formatSize(bytes: number | null): string { + if (bytes === null) return ''; + if (bytes < 1024) return `${bytes} B`; + if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)} KB`; + return `${(bytes / (1024 * 1024)).toFixed(1)} MB`; +} + +export function FilePickerDialog({ + open, + onClose, + onSelect, + initialPath, + extensions, + title, +}: FilePickerDialogProps): JSX.Element { + const { t } = useTranslation(); + const apiClient = useApiClient(); + const fsApi = createFsApi(apiClient); + + const [listing, setListing] = useState(null); + const [draft, setDraft] = useState(''); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(null); + const abortRef = useRef(null); + + const load = useCallback( + async (path: string | undefined): Promise => { + abortRef.current?.abort(); + const controller = new AbortController(); + abortRef.current = controller; + setLoading(true); + setError(null); + try { + const next = await fsApi.list({ + path, + includeFiles: true, + extensions, + signal: controller.signal, + }); + if (controller.signal.aborted) return; + setListing(next); + setDraft(next.path); + } catch (err) { + if (controller.signal.aborted) return; + if (err instanceof ApiError) { + setError(t(err.error.messageKey, { defaultValue: err.error.message })); + } else if (err instanceof Error) { + setError(err.message); + } else { + setError(String(err)); + } + } finally { + if (!controller.signal.aborted) setLoading(false); + } + }, + [fsApi, t, extensions], + ); + + useEffect(() => { + if (open) { + void load(initialPath); + } + return () => { + abortRef.current?.abort(); + }; + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [open]); + + function handleGo(): void { + if (!draft || draft === listing?.path) return; + void load(draft); + } + + return ( + + {t('components.filePicker.cancel')} + + } + > +
+
+ setDraft(event.target.value)} + onKeyDown={(event) => { + if (event.key === 'Enter') { + event.preventDefault(); + handleGo(); + } + }} + aria-label={t('components.filePicker.pathLabel')} + data-testid="zv-filepicker-path" + /> + +
+ +
+ {listing?.home && ( + + )} + {listing?.parent && ( + + )} +
+ + {loading && ( +
+ +
+ )} + + {error && !loading && ( +

+ {error} +

+ )} + + {!loading && !error && listing && ( +
    + {listing.entries.length === 0 && ( +
  • {t('components.filePicker.empty')}
    • + )} + {listing.entries.map((entry) => + entry.kind === 'dir' ? ( +
  • + +
    • + ) : ( +
  • + +
    • + ), + )} +
+ )} +
+ + ); +} diff --git a/apps/frontend/src/components/ui/index.ts b/apps/frontend/src/components/ui/index.ts index 1798b0f..2ab1fd2 100644 --- a/apps/frontend/src/components/ui/index.ts +++ b/apps/frontend/src/components/ui/index.ts @@ -26,6 +26,8 @@ export type { DirectoryInputProps } from './DirectoryInput'; export { DirectoryPickerDialog } from './DirectoryPickerDialog'; export type { DirectoryPickerDialogProps } from './DirectoryPickerDialog'; +export { FilePickerDialog } from './FilePickerDialog'; +export type { FilePickerDialogProps } from './FilePickerDialog'; export { Select } from './Select'; export type { SelectOption, SelectProps } from './Select'; diff --git a/apps/frontend/src/features/collections/api.ts b/apps/frontend/src/features/collections/api.ts index fd5be35..2595519 100644 --- a/apps/frontend/src/features/collections/api.ts +++ b/apps/frontend/src/features/collections/api.ts @@ -14,6 +14,8 @@ export type CollectionListResponse = components['schemas']['CollectionListRespon export type CollectionListItem = components['schemas']['CollectionListItem']; export type CollectionCreateRequest = components['schemas']['CollectionCreateRequest']; export type CollectionOpenRequest = components['schemas']['CollectionOpenRequest']; +export type CollectionImportRequest = components['schemas']['CollectionImportRequest']; +export type CollectionImportResponse = components['schemas']['CollectionImportResponse']; export type CollectionSchemaPayload = components['schemas']['CollectionSchema']; export type FieldAddRequest = components['schemas']['FieldAddRequest']; export type FieldRenameRequest = components['schemas']['FieldRenameRequest']; @@ -21,8 +23,7 @@ export type IndexCreateRequest = components['schemas']['IndexCreateRequest']; export type ScalarIndexCreateRequest = components['schemas']['ScalarIndexCreateRequest']; export type MaintenanceResponse = components['schemas']['MaintenanceResponse']; export type RecentCollectionItem = components['schemas']['RecentCollectionItem']; -export type RecentCollectionListResponse = - components['schemas']['RecentCollectionListResponse']; +export type RecentCollectionListResponse = components['schemas']['RecentCollectionListResponse']; export type RecentForgetRequest = components['schemas']['RecentForgetRequest']; export interface CollectionsApi { @@ -34,16 +35,8 @@ export interface CollectionsApi { flush(name: string, signal?: AbortSignal): Promise; optimize(name: string, signal?: AbortSignal): Promise; destroy(name: string, signal?: AbortSignal): Promise; - addField( - name: string, - body: FieldAddRequest, - signal?: AbortSignal, - ): Promise; - dropField( - name: string, - field: string, - signal?: AbortSignal, - ): Promise; + addField(name: string, body: FieldAddRequest, signal?: AbortSignal): Promise; + dropField(name: string, field: string, signal?: AbortSignal): Promise; renameField( name: string, field: string, @@ -55,22 +48,14 @@ export interface CollectionsApi { body: IndexCreateRequest, signal?: AbortSignal, ): Promise; - dropIndex( - name: string, - vectorField: string, - signal?: AbortSignal, - ): Promise; + dropIndex(name: string, vectorField: string, signal?: AbortSignal): Promise; createScalarIndex( name: string, field: string, body?: ScalarIndexCreateRequest, signal?: AbortSignal, ): Promise; - dropScalarIndex( - name: string, - field: string, - signal?: AbortSignal, - ): Promise; + dropScalarIndex(name: string, field: string, signal?: AbortSignal): Promise; /** * List the recently-opened collections (persisted to ``config.json``, * survives process restarts). Distinct from ``list`` which only reflects @@ -81,6 +66,9 @@ export interface CollectionsApi { forgetRecent(body: RecentForgetRequest, signal?: AbortSignal): Promise; /** Drop every entry from the recent list. Idempotent. */ clearRecent(signal?: AbortSignal): Promise; + /** Restore a collection from a snapshot package (collection-level op). */ + /** Import a collection from a snapshot package (collection-level op). */ + importCollection(body: CollectionImportRequest, signal?: AbortSignal): Promise; } /** Build a CollectionsApi bound to the given transport. */ @@ -102,25 +90,26 @@ export function createCollectionsApi(client: ApiClient): CollectionsApi { return client.request(url, { method: 'DELETE', signal }); }, flush: (name, signal) => - client.request( - `/collections/${encodeURIComponent(name)}:flush`, - { method: 'POST', signal }, - ), + client.request(`/collections/${encodeURIComponent(name)}:flush`, { + method: 'POST', + signal, + }), optimize: (name, signal) => - client.request( - `/collections/${encodeURIComponent(name)}:optimize`, - { method: 'POST', signal }, - ), + client.request(`/collections/${encodeURIComponent(name)}:optimize`, { + method: 'POST', + signal, + }), destroy: (name, signal) => client.request(`/collections/${encodeURIComponent(name)}:destroy`, { method: 'POST', signal, }), addField: (name, body, signal) => - client.request( - `/collections/${encodeURIComponent(name)}/fields`, - { method: 'POST', body, signal }, - ), + client.request(`/collections/${encodeURIComponent(name)}/fields`, { + method: 'POST', + body, + signal, + }), dropField: (name, field, signal) => client.request( `/collections/${encodeURIComponent(name)}/fields/${encodeURIComponent(field)}`, @@ -132,10 +121,11 @@ export function createCollectionsApi(client: ApiClient): CollectionsApi { { method: 'PATCH', body, signal }, ), createIndex: (name, body, signal) => - client.request( - `/collections/${encodeURIComponent(name)}/indexes`, - { method: 'POST', body, signal }, - ), + client.request(`/collections/${encodeURIComponent(name)}/indexes`, { + method: 'POST', + body, + signal, + }), dropIndex: (name, vectorField, signal) => client.request( `/collections/${encodeURIComponent(name)}/indexes/${encodeURIComponent(vectorField)}`, @@ -161,5 +151,11 @@ export function createCollectionsApi(client: ApiClient): CollectionsApi { }), clearRecent: (signal) => client.request('/collections/recent', { method: 'DELETE', signal }), + importCollection: (body, signal) => + client.request('/collections:import', { + method: 'POST', + body, + signal, + }), }; } diff --git a/apps/frontend/src/features/collections/hooks.ts b/apps/frontend/src/features/collections/hooks.ts index c7aa3c5..22fa161 100644 --- a/apps/frontend/src/features/collections/hooks.ts +++ b/apps/frontend/src/features/collections/hooks.ts @@ -20,6 +20,8 @@ import { type CollectionCreateRequest, type CollectionListResponse, type CollectionOpenRequest, + type CollectionImportRequest, + type CollectionImportResponse, type CollectionSummary, type FieldAddRequest, type FieldRenameRequest, @@ -35,7 +37,9 @@ export const collectionsQueryKey = ['collections', 'list'] as const; /** Query key factory for a single Collection detail. */ export const collectionDetailQueryKey = (name: string, path?: string) => - path ? (['collections', 'detail', name, path] as const) : (['collections', 'detail', name] as const); + path + ? (['collections', 'detail', name, path] as const) + : (['collections', 'detail', name] as const); /** Canonical query key for the recently-opened collections list. */ export const recentCollectionsQueryKey = ['collections', 'recent', 'list'] as const; @@ -96,6 +100,28 @@ export function useOpenCollection(): UseMutationResult< }); } +/** + * Import a collection from a snapshot package — a collection-level + * lifecycle operation (create the collection from the manifest schema and + * load its data in one pass), not a variant of document import. + */ +export function useImportCollection(): UseMutationResult< + CollectionImportResponse, + unknown, + CollectionImportRequest +> { + const client = useApiClient(); + const api = createCollectionsApi(client); + const qc = useQueryClient(); + return useMutation({ + mutationFn: (body: CollectionImportRequest) => api.importCollection(body), + onSuccess: () => { + void qc.invalidateQueries({ queryKey: collectionsQueryKey }); + void qc.invalidateQueries({ queryKey: recentCollectionsQueryKey }); + }, + }); +} + export interface CloseCollectionVariables { readonly name: string; readonly path?: string; @@ -125,11 +151,7 @@ function invalidateCollection( ]); } -export function useFlushCollection(): UseMutationResult< - MaintenanceResponse, - unknown, - string -> { +export function useFlushCollection(): UseMutationResult { const client = useApiClient(); const api = createCollectionsApi(client); const qc = useQueryClient(); @@ -139,11 +161,7 @@ export function useFlushCollection(): UseMutationResult< }); } -export function useOptimizeCollection(): UseMutationResult< - MaintenanceResponse, - unknown, - string -> { +export function useOptimizeCollection(): UseMutationResult { const client = useApiClient(); const api = createCollectionsApi(client); const qc = useQueryClient(); @@ -176,11 +194,7 @@ export interface AddFieldVariables { readonly body: FieldAddRequest; } -export function useAddField(): UseMutationResult< - CollectionSummary, - unknown, - AddFieldVariables -> { +export function useAddField(): UseMutationResult { const client = useApiClient(); const api = createCollectionsApi(client); const qc = useQueryClient(); @@ -195,11 +209,7 @@ export interface DropFieldVariables { readonly field: string; } -export function useDropField(): UseMutationResult< - CollectionSummary, - unknown, - DropFieldVariables -> { +export function useDropField(): UseMutationResult { const client = useApiClient(); const api = createCollectionsApi(client); const qc = useQueryClient(); @@ -253,11 +263,7 @@ export interface DropIndexVariables { readonly vectorField: string; } -export function useDropIndex(): UseMutationResult< - CollectionSummary, - unknown, - DropIndexVariables -> { +export function useDropIndex(): UseMutationResult { const client = useApiClient(); const api = createCollectionsApi(client); const qc = useQueryClient(); @@ -310,10 +316,7 @@ export function useDropScalarIndex(): UseMutationResult< // ── Recently opened collections ───────────────────────────────────────────── -export function useListRecent(): UseQueryResult< - RecentCollectionListResponse, - unknown -> { +export function useListRecent(): UseQueryResult { const client = useApiClient(); const api = createCollectionsApi(client); return useQuery({ @@ -322,11 +325,7 @@ export function useListRecent(): UseQueryResult< }); } -export function useForgetRecent(): UseMutationResult< - void, - unknown, - RecentForgetRequest -> { +export function useForgetRecent(): UseMutationResult { const client = useApiClient(); const api = createCollectionsApi(client); const qc = useQueryClient(); diff --git a/apps/frontend/src/features/collections/index.ts b/apps/frontend/src/features/collections/index.ts index 321f8ca..48ba9ee 100644 --- a/apps/frontend/src/features/collections/index.ts +++ b/apps/frontend/src/features/collections/index.ts @@ -7,6 +7,8 @@ export { type CollectionListItem, type CollectionCreateRequest, type CollectionOpenRequest, + type CollectionImportRequest, + type CollectionImportResponse, type CollectionSchemaPayload, type FieldAddRequest, type FieldRenameRequest, @@ -25,6 +27,7 @@ export { useCollection, useCreateCollection, useOpenCollection, + useImportCollection, useCloseCollection, useFlushCollection, useOptimizeCollection, diff --git a/apps/frontend/src/features/documents/api.test.ts b/apps/frontend/src/features/documents/api.test.ts new file mode 100644 index 0000000..d659017 --- /dev/null +++ b/apps/frontend/src/features/documents/api.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, it } from 'vitest'; + +import { buildExportUrl } from './api'; + +describe('buildExportUrl', () => { + const base = '/api/v1'; + + it('defaults to jsonl with vectors included', () => { + expect(buildExportUrl(base, 'demo', { includeVector: true })).toBe( + '/api/v1/collections/demo/documents:export?includeVector=true&format=jsonl', + ); + }); + + it('omits vectors when disabled', () => { + expect(buildExportUrl(base, 'demo', { includeVector: false })).toContain('includeVector=false'); + }); + + it('joins output fields with commas', () => { + const url = buildExportUrl(base, 'demo', { + includeVector: false, + outputFields: ['title', 'score'], + }); + expect(url).toContain('outputFields=title%2Cscore'); + }); + + it('skips an empty output field list', () => { + const url = buildExportUrl(base, 'demo', { includeVector: true, outputFields: [] }); + expect(url).not.toContain('outputFields'); + }); + + it('encodes the collection name', () => { + const url = buildExportUrl(base, 'my collection', { includeVector: true }); + expect(url).toContain('/collections/my%20collection/documents:export'); + }); + + it('honours an absolute base url (Tauri sidecar)', () => { + const url = buildExportUrl('http://127.0.0.1:7861/api/v1', 'demo', { + includeVector: true, + }); + expect(url.startsWith('http://127.0.0.1:7861/api/v1/collections/demo')).toBe(true); + }); + + it('adds mode=snapshot for snapshot exports', () => { + const url = buildExportUrl(base, 'demo', { includeVector: true, mode: 'snapshot' }); + expect(url).toContain('mode=snapshot'); + }); + + it('omits the mode param for data exports', () => { + const url = buildExportUrl(base, 'demo', { includeVector: true, mode: 'data' }); + expect(url).not.toContain('mode='); + }); +}); diff --git a/apps/frontend/src/features/documents/api.ts b/apps/frontend/src/features/documents/api.ts index 832b7b3..eb46577 100644 --- a/apps/frontend/src/features/documents/api.ts +++ b/apps/frontend/src/features/documents/api.ts @@ -21,10 +21,11 @@ export type DocumentUpsertRequest = components['schemas']['DocumentUpsertRequest export type DocumentUpsertResponse = components['schemas']['DocumentUpsertResponse']; export type DocumentUpdateRequest = components['schemas']['DocumentUpdateRequest']; export type DocumentUpdateResponse = components['schemas']['DocumentUpdateResponse']; -export type DocumentDeleteByFilterRequest = - components['schemas']['DocumentDeleteByFilterRequest']; +export type DocumentDeleteByFilterRequest = components['schemas']['DocumentDeleteByFilterRequest']; export type DocumentDeleteByFilterResponse = components['schemas']['DocumentDeleteByFilterResponse']; +export type DocumentImportRequest = components['schemas']['DocumentImportRequest']; +export type DocumentImportResponse = components['schemas']['DocumentImportResponse']; /** A single document row. The backend returns an open-ended dict of fields. */ export type DocumentRecord = Record; @@ -61,6 +62,11 @@ export interface DocumentsApi { body: DocumentDeleteByFilterRequest, signal?: AbortSignal, ): Promise; + importDocuments( + collection: string, + body: DocumentImportRequest, + signal?: AbortSignal, + ): Promise; } /** Build a DocumentsApi bound to the given transport. */ @@ -106,5 +112,43 @@ export function createDocumentsApi(client: ApiClient): DocumentsApi { `/collections/${encodeURIComponent(collection)}/documents:deleteByFilter`, { method: 'POST', body, signal }, ), + importDocuments: (collection, body, signal) => + client.request( + `/collections/${encodeURIComponent(collection)}/documents:import`, + { method: 'POST', body, signal }, + ), }; } + +export interface ExportOptions { + readonly includeVector: boolean; + readonly outputFields?: ReadonlyArray; + readonly format?: string; + /** ``data`` (single JSONL file, default) or ``snapshot`` (tar.gz bundle). */ + readonly mode?: 'data' | 'snapshot'; +} + +/** + * Build the download URL for ``GET /collections/{name}/documents:export``. + * + * Exported through a *native* browser download (````), never + * ``fetch``: the response streams a potentially gigabyte-sized file and + * buffering it in JS memory would crash the tab (design doc §6.3). + */ +export function buildExportUrl( + baseUrl: string, + collection: string, + options: ExportOptions, +): string { + const params = new URLSearchParams(); + params.set('includeVector', String(options.includeVector)); + if (options.outputFields && options.outputFields.length > 0) { + params.set('outputFields', options.outputFields.join(',')); + } + params.set('format', options.format ?? 'jsonl'); + if (options.mode === 'snapshot') { + params.set('mode', 'snapshot'); + } + const qs = params.toString(); + return `${baseUrl}/collections/${encodeURIComponent(collection)}/documents:export?${qs}`; +} diff --git a/apps/frontend/src/features/documents/hooks.ts b/apps/frontend/src/features/documents/hooks.ts index 089df5d..46d7679 100644 --- a/apps/frontend/src/features/documents/hooks.ts +++ b/apps/frontend/src/features/documents/hooks.ts @@ -25,6 +25,8 @@ import { type DocumentBrowseResponse, type DocumentDeleteByFilterRequest, type DocumentDeleteByFilterResponse, + type DocumentImportRequest, + type DocumentImportResponse, type DocumentInsertRequest, type DocumentInsertResponse, type DocumentRecord, @@ -205,3 +207,20 @@ export function useDeleteDocumentsByFilter( onSuccess: () => invalidateAfterWrite(queryClient, collection), }); } + +/** + * Bulk import from a local JSONL file. The response carries a per-row report + * (imported / failed / errors); the HTTP status stays 200 for partial + * success, so callers inspect ``data`` rather than treating it as failure. + */ +export function useImportDocuments( + collection: string, +): UseMutationResult { + const client = useApiClient(); + const api = createDocumentsApi(client); + const queryClient = useQueryClient(); + return useMutation({ + mutationFn: (body) => api.importDocuments(collection, body), + onSuccess: () => invalidateAfterWrite(queryClient, collection), + }); +} diff --git a/apps/frontend/src/features/documents/index.ts b/apps/frontend/src/features/documents/index.ts index 14d69d6..495101b 100644 --- a/apps/frontend/src/features/documents/index.ts +++ b/apps/frontend/src/features/documents/index.ts @@ -1,7 +1,8 @@ /** * Public surface of the documents feature module. */ -export { createDocumentsApi } from './api'; +export { buildExportUrl, createDocumentsApi } from './api'; +export type { ExportOptions } from './api'; export type { DocumentBatchDeleteRequest, DocumentBatchDeleteResponse, @@ -9,6 +10,8 @@ export type { DocumentBrowseResponse, DocumentDeleteByFilterRequest, DocumentDeleteByFilterResponse, + DocumentImportRequest, + DocumentImportResponse, DocumentInsertRequest, DocumentInsertResponse, DocumentRecord, @@ -28,6 +31,7 @@ export { useDeleteDocument, useDeleteDocumentsBatch, useDeleteDocumentsByFilter, + useImportDocuments, useUpdateDocuments, useUpsertDocuments, } from './hooks'; diff --git a/apps/frontend/src/i18n/en.json b/apps/frontend/src/i18n/en.json index 58a1f67..ab3c226 100644 --- a/apps/frontend/src/i18n/en.json +++ b/apps/frontend/src/i18n/en.json @@ -17,6 +17,15 @@ "select": "Select", "cancel": "Cancel", "empty": "No subdirectories." + }, + "filePicker": { + "title": "Choose a file", + "pathLabel": "Path", + "go": "Go", + "home": "Home", + "parent": "Up", + "cancel": "Cancel", + "empty": "No entries." } }, "nav": { @@ -166,9 +175,11 @@ "dimRange": "Dimension must be an integer between 1 and 32768.", "fieldsRequired": "Add at least one vector or scalar field.", "fieldNameRequired": "Each field needs a valid name.", - "reservedName": "Field name \"{{name}}\" is reserved by Zvec — pick another name.", "duplicateField": "Field name \"{{name}}\" is used more than once — each field must have a unique name.", "invalidIndexParams": "Index params must be a JSON object." + }, + "hints": { + "idColumn": "Naming a column \"id\" is allowed. This collection's primary key will then be shown and exchanged under \"$id\" to avoid the clash." } }, "open": { @@ -259,8 +270,7 @@ "successDescription": "\"{{name}}\" added.", "failureTitle": "Add field failed", "errors": { - "nameRequired": "Field name is required.", - "reservedName": "Field name \"{{name}}\" is reserved." + "nameRequired": "Field name is required." } }, "renameFieldDialog": { @@ -440,6 +450,48 @@ "successTitle": "Documents deleted", "successDescription": "{{count}} document(s) removed.", "failureTitle": "Delete by filter failed" + }, + "import": { + "openLabel": "Import", + "title": "Import documents from file", + "fileLabel": "JSONL file", + "fileHelp": "A JSONL file — one JSON document per line. The file is read from this machine. To import a whole collection from a snapshot (.tar.gz), use Collections → Import Collection.", + "browse": "Browse…", + "pickerTitle": "Choose a JSONL file", + "modeLabel": "Write mode", + "modeReplace": "Replace (overwrite documents with the same id)", + "modeReplaceHelp": "The file's content becomes the whole document for that id — fields absent from the file are cleared. Importing the same file twice gives the same result.", + "modeInsert": "Insert only (fail rows whose id already exists)", + "onErrorLabel": "On row error", + "onErrorAbort": "Abort at the first failing row", + "onErrorSkip": "Skip failing rows and continue", + "submit": "Import", + "submitting": "Importing…", + "cancel": "Cancel", + "done": "Done", + "resultTitle": "Imported {{imported}}, failed {{failed}}", + "reportSummary": "{{imported}} of {{total}} row(s) imported; {{failed}} failed.", + "duration": "took {{seconds}}s", + "aborted": "Import stopped at the first error (rows written before it are kept).", + "errorsCaption": "Failing rows:", + "errorsTruncated": "(showing the first 100)" + }, + "export": { + "openLabel": "Export", + "title": "Export documents", + "modeLabel": "Export mode", + "modeData": "Data file (.jsonl)", + "modeSnapshot": "Collection snapshot (.tar.gz)", + "modeHelp": "A snapshot bundles the schema with the data so the collection can be recreated elsewhere.", + "includeVector": "Include vector data", + "includeVectorHelp": "Vectors are the main source of file size; uncheck to export fields only.", + "fieldsLabel": "Field filter (optional)", + "fieldsHelp": "Select specific scalar fields, or leave all unchecked to export every field.", + "submit": "Start export", + "cancel": "Cancel", + "startedTitle": "Export started", + "startedBody": "Your download has begun — check the browser's download area.", + "maintenanceHint": "Reads and writes stay available during the export; maintenance (index creation, optimize, …) is briefly rejected until it finishes." } }, "searchPanel": { @@ -660,6 +712,28 @@ "modeJson": "JSON", "idRequired": "Document ID" } + }, + "importCollection": { + "menuItem": "Import Collection", + "title": "Import Collection", + "fileLabel": "Snapshot file (.tar.gz)", + "fileHelp": "A snapshot produced by Export → Collection snapshot: schema and data bundled together.", + "browse": "Browse…", + "pickerTitle": "Choose a snapshot file", + "targetLabel": "Target directory", + "targetHelp": "The imported collection is created here. Defaults to a directory next to the snapshot file — editable. Must be a path that does not exist yet.", + "nameLabel": "Collection name (optional)", + "nameHelp": "Defaults to the name recorded in the snapshot; rename to avoid a clash.", + "namePlaceholder": "Leave empty to keep the snapshot's name", + "submit": "Import", + "submitting": "Importing…", + "successTitle": "Collection imported", + "successBody": "\"{{name}}\" imported — {{count}} document(s) loaded.", + "errors": { + "required": "Both the snapshot file and the target directory are required." + }, + "partialTitle": "Import finished with problems", + "partialBody": "\"{{name}}\" was imported with {{failed}} failing row(s) — {{imported}} of {{count}} loaded. See the import report." } }, "documents": { @@ -822,7 +896,7 @@ "title": "Recently opened", "aria": "Recently opened collections", "empty": "No recently opened collections yet.", - "loading": "Loading recently opened collections\u2026", + "loading": "Loading recently opened collections…", "clearAll": "Clear all", "forgetAria": "Forget {{path}}" } @@ -848,7 +922,7 @@ "validation": "Invalid input. Please review the highlighted fields.", "code": { "COLLECTION_NOT_FOUND": "Collection was not found.", - "COLLECTION_ALREADY_EXISTS": "A Collection with this name or path already exists.", + "COLLECTION_ALREADY_EXISTS": "Conflicts with an already-open Collection's name or path. Close it first or pick another name / path.", "INVALID_SCHEMA": "The Collection schema is invalid.", "INVALID_FILTER_EXPRESSION": "The filter expression is invalid.", "DOCUMENT_NOT_FOUND": "Document was not found.", diff --git a/apps/frontend/src/i18n/zh.json b/apps/frontend/src/i18n/zh.json index 780a98e..a74d1e3 100644 --- a/apps/frontend/src/i18n/zh.json +++ b/apps/frontend/src/i18n/zh.json @@ -17,6 +17,15 @@ "select": "选择", "cancel": "取消", "empty": "没有子目录。" + }, + "filePicker": { + "title": "选择文件", + "pathLabel": "路径", + "go": "前往", + "home": "主目录", + "parent": "上级", + "cancel": "取消", + "empty": "没有条目。" } }, "nav": { @@ -166,9 +175,11 @@ "dimRange": "维度必须是 1 到 32768 之间的整数。", "fieldsRequired": "至少添加一个向量或标量字段。", "fieldNameRequired": "每个字段需要一个有效名称。", - "reservedName": "字段名 \"{{name}}\" 是 Zvec 保留字段 — 请换一个名称。", "duplicateField": "字段名 \"{{name}}\" 重复了 — 每个字段名必须唯一。", "invalidIndexParams": "索引参数必须是 JSON 对象。" + }, + "hints": { + "idColumn": "允许使用 \"id\" 作为列名。此时该集合的主键将显示为 \"$id\"，以避免混淆。" } }, "open": { @@ -259,8 +270,7 @@ "successDescription": "\"{{name}}\" 已添加。", "failureTitle": "添加字段失败", "errors": { - "nameRequired": "字段名为必填项。", - "reservedName": "字段名 \"{{name}}\" 是保留字段。" + "nameRequired": "字段名为必填项。" } }, "renameFieldDialog": { @@ -440,6 +450,48 @@ "successTitle": "文档已删除", "successDescription": "已移除 {{count}} 条文档。", "failureTitle": "按过滤器删除失败" + }, + "import": { + "openLabel": "导入", + "title": "从文件导入文档", + "fileLabel": "JSONL 文件", + "fileHelp": "JSONL 文件（每行一个 JSON 文档）。文件将从本机读取。若要从快照（.tar.gz）导入整个集合，请使用侧边栏「导入集合」。", + "browse": "浏览…", + "pickerTitle": "选择 JSONL 文件", + "modeLabel": "写入模式", + "modeReplace": "覆盖（同 id 的文档整条覆盖）", + "modeReplaceHelp": "文件内容将成为该 id 的完整文档 — 文件中未出现的字段会被清空。重复导入同一文件结果一致。", + "modeInsert": "仅新增（id 已存在的行将失败）", + "onErrorLabel": "遇到行错误时", + "onErrorAbort": "在第一个失败行中止", + "onErrorSkip": "跳过失败行并继续", + "submit": "导入", + "submitting": "导入中…", + "cancel": "取消", + "done": "完成", + "resultTitle": "已导入 {{imported}} 条，失败 {{failed}} 条", + "reportSummary": "共 {{total}} 行，已导入 {{imported}} 条，失败 {{failed}} 条。", + "duration": "耗时 {{seconds}} 秒", + "aborted": "导入在第一个错误处停止（此前已写入的行保留）。", + "errorsCaption": "失败的行：", + "errorsTruncated": "（仅显示前 100 条）" + }, + "export": { + "openLabel": "导出", + "title": "导出文档", + "modeLabel": "导出模式", + "modeData": "数据文件（.jsonl）", + "modeSnapshot": "集合快照（.tar.gz）", + "modeHelp": "快照将 schema 与数据一起打包，便于在别处重建集合。", + "includeVector": "包含向量数据", + "includeVectorHelp": "向量是文件体积的主要来源；取消勾选则仅导出字段。", + "fieldsLabel": "字段筛选（可选）", + "fieldsHelp": "勾选特定标量字段，或全部不勾以导出所有字段。", + "submit": "开始导出", + "cancel": "取消", + "startedTitle": "导出已开始", + "startedBody": "下载已开始 — 请查看浏览器的下载区域。", + "maintenanceHint": "导出期间读写仍可用；建索引、优化等维护操作会在导出完成前被暂时拒绝。" } }, "searchPanel": { @@ -660,6 +712,28 @@ "modeJson": "JSON", "idRequired": "文档 ID" } + }, + "importCollection": { + "menuItem": "导入集合", + "title": "导入集合", + "fileLabel": "快照文件（.tar.gz）", + "fileHelp": "由「导出 → 集合快照」产生的快照包，内含 schema 与全部数据。", + "browse": "浏览…", + "pickerTitle": "选择快照文件", + "targetLabel": "目标目录", + "targetHelp": "导入的集合将创建在该目录。默认填充为快照文件旁的同名目录，可修改；必须是尚不存在的路径。", + "nameLabel": "集合名（可选）", + "nameHelp": "默认使用快照中记录的名称；改名可避免与已有集合重名。", + "namePlaceholder": "留空则沿用快照中的名称", + "submit": "导入", + "submitting": "导入中…", + "successTitle": "集合已导入", + "successBody": "\"{{name}}\" 导入完成，共载入 {{count}} 条文档。", + "errors": { + "required": "请同时填写快照文件与目标目录。" + }, + "partialTitle": "导入完成，但有部分问题", + "partialBody": "「{{name}}」已导入，{{count}} 条中有 {{failed}} 条失败、{{imported}} 条成功。请查看导入报告。" } }, "documents": { @@ -848,7 +922,7 @@ "validation": "输入无效。请检查高亮的字段。", "code": { "COLLECTION_NOT_FOUND": "未找到集合。", - "COLLECTION_ALREADY_EXISTS": "已存在同名或同路径的集合。", + "COLLECTION_ALREADY_EXISTS": "与已打开集合的名称或路径冲突。请先关闭对应集合，或另选名称 / 路径。", "INVALID_SCHEMA": "集合 schema 无效。", "INVALID_FILTER_EXPRESSION": "过滤表达式无效。", "DOCUMENT_NOT_FOUND": "未找到文档。", diff --git a/apps/frontend/src/layouts/AppShell.css b/apps/frontend/src/layouts/AppShell.css index 6b7949e..84ebb8b 100644 --- a/apps/frontend/src/layouts/AppShell.css +++ b/apps/frontend/src/layouts/AppShell.css @@ -81,6 +81,44 @@ display: flex; gap: 2px; } +.zv-sidebar__menu-anchor { + position: relative; +} +.zv-sidebar__menu-overlay { + position: fixed; + inset: 0; + z-index: 40; +} +.zv-sidebar__menu { + position: absolute; + top: calc(100% + 4px); + right: 0; + z-index: 41; + min-width: 180px; + padding: 4px; + border-radius: 8px; + border: 1px solid var(--zv-color-border, #e5e7eb); + background: var(--zv-color-surface, #ffffff); + box-shadow: 0 4px 16px rgba(0, 0, 0, 0.12); + display: flex; + flex-direction: column; + gap: 2px; +} +.zv-sidebar__menu button { + appearance: none; + border: none; + background: transparent; + text-align: left; + padding: 6px 10px; + border-radius: 6px; + font-size: 0.85rem; + color: var(--zv-color-text, #111827); + cursor: pointer; + white-space: nowrap; +} +.zv-sidebar__menu button:hover { + background: var(--zv-color-surface-hover, rgba(0, 0, 0, 0.05)); +} .zv-sidebar__action-btn { width: 20px; height: 20px; diff --git a/apps/frontend/src/layouts/AppShell.test.tsx b/apps/frontend/src/layouts/AppShell.test.tsx index 6af985f..a5f3fb6 100644 --- a/apps/frontend/src/layouts/AppShell.test.tsx +++ b/apps/frontend/src/layouts/AppShell.test.tsx @@ -31,11 +31,13 @@ function freshState(): FakeShellState { collections: [{ name: 'open-demo', path: '/tmp/open-demo' }], recent: [{ name: 'closed-demo', path: '/tmp/closed-demo' }], embeddings: [ - { name: 'local-dense', description: null, config: { type: 'default_local_dense', dimension: 4 } }, - ], - rerankers: [ - { name: 'rrf', description: null, config: { type: 'rrf', rankConstant: 60 } }, + { + name: 'local-dense', + description: null, + config: { type: 'default_local_dense', dimension: 4 }, + }, ], + rerankers: [{ name: 'rrf', description: null, config: { type: 'rrf', rankConstant: 60 } }], openedPaths: [], closed: [], forgotten: [], @@ -129,10 +131,22 @@ function renderShell( }> root} /> - collections} /> - collection detail} /> - embedding detail} /> - reranker detail} /> + collections} + /> + collection detail} + /> + embedding detail} + /> + reranker detail} + /> @@ -219,7 +233,9 @@ describe('AppShell', () => { const user = userEvent.setup(); renderShell('/collections', freshState()); - await user.click(await screen.findByTitle('Create Collection')); + // The "+" button toggles a menu holding the collection-level actions. + await user.click(await screen.findByTestId('zv-collections-add')); + await user.click(screen.getByTestId('zv-collections-menu-create')); expect(await screen.findByRole('dialog', { name: /create collection/i })).toBeInTheDocument(); await user.click(screen.getByRole('button', { name: /cancel/i })); @@ -227,6 +243,16 @@ describe('AppShell', () => { expect(await screen.findByRole('dialog', { name: /open collection/i })).toBeInTheDocument(); }); + it('opens the import-collection dialog from the "+" menu', async () => { + const user = userEvent.setup(); + renderShell('/collections', freshState()); + + await user.click(await screen.findByTestId('zv-collections-add')); + await user.click(screen.getByTestId('zv-collections-menu-import')); + + expect(await screen.findByRole('dialog', { name: /import collection/i })).toBeInTheDocument(); + }); + it('toggles sidebar visibility, theme, language, and the guide tour', async () => { const user = userEvent.setup(); renderShell('/collections', freshState()); diff --git a/apps/frontend/src/layouts/AppShell.tsx b/apps/frontend/src/layouts/AppShell.tsx index da9bad4..cb5f68f 100644 --- a/apps/frontend/src/layouts/AppShell.tsx +++ b/apps/frontend/src/layouts/AppShell.tsx @@ -9,14 +9,26 @@ import { useToast } from '@/components/ui/toast-context'; import { ApiError } from '@/lib/api-client'; import { useTheme } from '@/styles/theme'; import { setLanguage, type Language } from '@/i18n'; -import { useCollectionsList, useListRecent, useCloseCollection, useOpenCollection, useForgetRecent } from '@/features/collections/hooks'; -import { useListEmbeddings, useListRerankers, useDeleteEmbedding, useDeleteReranker } from '@/features/ai/hooks'; +import { + useCollectionsList, + useListRecent, + useCloseCollection, + useOpenCollection, + useForgetRecent, +} from '@/features/collections/hooks'; +import { + useListEmbeddings, + useListRerankers, + useDeleteEmbedding, + useDeleteReranker, +} from '@/features/ai/hooks'; import type { CollectionListItem, RecentCollectionItem } from '@/features/collections/api'; import type { EmbeddingFunctionRecord, RerankerFunctionRecord } from '@/features/ai/api'; import logoLight from '@/assets/logo-light.svg'; import logoDark from '@/assets/logo-dark.svg'; import { getEmbeddingTag } from '@/features/ai/utils'; import { CreateCollectionDialog } from '@/pages/collections/CreateCollectionDialog'; +import { ImportCollectionDialog } from '@/pages/collections/ImportCollectionDialog'; import { OpenCollectionDialog } from '@/pages/collections/OpenCollectionDialog'; import { CreateEmbeddingDialog } from '@/pages/embeddings/CreateEmbeddingDialog'; import { CreateRerankerDialog } from '@/pages/rerankers/CreateRerankerDialog'; @@ -89,7 +101,7 @@ function DemoTabs(): JSX.Element { const mainRect = main?.getBoundingClientRect(); const style: React.CSSProperties = { position: 'fixed', - top: topbarRect ? topbarRect.bottom : (mainRect ? mainRect.top + 52 : 52), + top: topbarRect ? topbarRect.bottom : mainRect ? mainRect.top + 52 : 52, left: mainRect ? mainRect.left : 240, width: mainRect ? mainRect.width : 'calc(100vw - 240px)', zIndex: 9997, @@ -102,7 +114,9 @@ function DemoTabs(): JSX.Element { className="zv-detail-tabs" style={{ margin: 0, background: 'var(--zv-color-surface)' }} > - + @@ -163,6 +177,8 @@ export function AppShell(): JSX.Element { const [sidebarCollapsed, setSidebarCollapsed] = useState(false); const [showCreateCollection, setShowCreateCollection] = useState(false); const [showOpenCollection, setShowOpenCollection] = useState(false); + const [showImportCollection, setShowRestoreSnapshot] = useState(false); + const [collectionMenuOpen, setCollectionMenuOpen] = useState(false); const [showCreateEmbedding, setShowCreateEmbedding] = useState(false); const [showCreateReranker, setShowCreateReranker] = useState(false); const [deleteEmbeddingTarget, setDeleteEmbeddingTarget] = useState(null); @@ -218,26 +234,33 @@ export function AppShell(): JSX.Element { function handleCloseCollection(e: React.MouseEvent, name: string, path?: string) { e.stopPropagation(); e.preventDefault(); - closeCollection.mutate({ name, path }, { - onSuccess: () => { - if (currentPath === `/collections/${encodeURIComponent(name)}`) { - navigate('/collections'); - } + closeCollection.mutate( + { name, path }, + { + onSuccess: () => { + if (currentPath === `/collections/${encodeURIComponent(name)}`) { + navigate('/collections'); + } + }, }, - }); + ); } function handleForgetRecent(e: React.MouseEvent, path: string) { e.stopPropagation(); e.preventDefault(); - forgetRecent.mutate({ path }, { - onError: (err) => { - const title = err instanceof ApiError - ? t(err.error.messageKey, { defaultValue: err.error.message }) - : t('errors.unknown'); - toast.push({ title, severity: 'error' }); + forgetRecent.mutate( + { path }, + { + onError: (err) => { + const title = + err instanceof ApiError + ? t(err.error.messageKey, { defaultValue: err.error.message }) + : t('errors.unknown'); + toast.push({ title, severity: 'error' }); + }, }, - }); + ); } function handleDeleteEmbeddingClick(e: React.MouseEvent, name: string) { @@ -256,9 +279,10 @@ export function AppShell(): JSX.Element { if (currentPath === `/embeddings/${encodeURIComponent(name)}`) navigate('/collections'); }, onError: (err) => { - const title = err instanceof ApiError - ? t(err.error.messageKey, { defaultValue: err.error.message }) - : t('errors.unknown'); + const title = + err instanceof ApiError + ? t(err.error.messageKey, { defaultValue: err.error.message }) + : t('errors.unknown'); toast.push({ title, severity: 'error' }); }, }); @@ -280,9 +304,10 @@ export function AppShell(): JSX.Element { if (currentPath === `/rerankers/${encodeURIComponent(name)}`) navigate('/collections'); }, onError: (err) => { - const title = err instanceof ApiError - ? t(err.error.messageKey, { defaultValue: err.error.message }) - : t('errors.unknown'); + const title = + err instanceof ApiError + ? t(err.error.messageKey, { defaultValue: err.error.message }) + : t('errors.unknown'); toast.push({ title, severity: 'error' }); }, }); @@ -291,8 +316,14 @@ export function AppShell(): JSX.Element { return (
{/* ── Sidebar ── */} -
{view === 'insert' && ( - + )} {view === 'upsert' && ( - + )} {view === 'update' && ( - + )} {view === 'delete' && } @@ -107,14 +126,22 @@ interface DocSlot { function makeDocSlot( fields: ReadonlyArray<{ name: string }>, vectors: ReadonlyArray, + pkKey: string, ): DocSlot { const fv: Record = {}; - for (const f of fields) if (f.name !== 'id') fv[f.name] = ''; + // The primary-key column is entered separately (``docId``), whatever its + // row key is (``id`` or ``$id``). + for (const f of fields) if (f.name !== pkKey) fv[f.name] = ''; const vi: Record = {}; for (const v of vectors) { vi[v.name] = { mode: 'raw', rawText: '', embedding: '', embedText: '' }; } - return { id: `doc-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 6)}`, docId: '', fieldValues: fv, vectorInputs: vi }; + return { + id: `doc-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 6)}`, + docId: '', + fieldValues: fv, + vectorInputs: vi, + }; } function embedResultMatchesVector(result: EmbedResponse, vector: { dataType: string }): boolean { @@ -151,8 +178,12 @@ function VectorFieldInput({ return (
- {vector.name} - {vectorDimensionLabel(vector)} {vector.dataType} + + {vector.name} + + + {vectorDimensionLabel(vector)} {vector.dataType} +
@@ -177,7 +208,10 @@ function VectorFieldInput({ {input.mode === 'embed' ? ( <> {dimMismatch && ( - + {t('pages.collections.detail.mutate.dimMismatch', { embDim: dimMismatch.embDim, fieldDim: dimMismatch.fieldDim, @@ -186,9 +220,11 @@ function VectorFieldInput({ )}