diff --git a/Cargo.lock b/Cargo.lock index e0847c9d..d1f4c8e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2625,6 +2625,7 @@ dependencies = [ "serde", "serde_json", "tabled", + "tempfile", "tokio", "toml", ] diff --git a/docs/ja/src/laurus-cli/commands.md b/docs/ja/src/laurus-cli/commands.md index 72f97a39..633c7342 100644 --- a/docs/ja/src/laurus-cli/commands.md +++ b/docs/ja/src/laurus-cli/commands.md @@ -237,25 +237,45 @@ laurus add doc --id --data | `--id ` | はい | 外部ドキュメント ID（文字列） | | `--data ` | はい | JSON 文字列としてのドキュメントフィールド | -JSON フォーマットはフィールド名と値を対応付けたフラットなオブジェクトです: +JSON はドキュメントの serde 形式です: `fields` オブジェクトの各フィールド名に、外部タグ付きの値（`Text`・`Int64`・`Float64`・`Bool`・`VectorValue` など）を対応付けます。 ```json { - "title": "Introduction to Rust", - "body": "Rust is a systems programming language.", - "category": "programming" + "fields": { + "title": {"Text": "Introduction to Rust"}, + "body": {"Text": "Rust is a systems programming language."}, + "year": {"Int64": 2024} + } } ``` **例:** ```bash -laurus add doc --id doc1 --data '{"title":"Hello World","body":"This is a test document."}' +laurus add doc --id doc1 --data '{"fields":{"title":{"Text":"Hello World"},"body":{"Text":"This is a test document."}}}' # Document 'doc1' added. Run 'commit' to persist changes. ``` > **ヒント:** 複数のドキュメントが同じ外部 ID を共有できます（チャンキングパターン）。各チャンクに対して `add doc` を使用してください。 +### `add docs` + +JSONL ファイルからドキュメントチャンクをバルク追加します — 1 行に 1 エントリの `{"id": "...", "document": {"fields": {...}}}` 形式で、`document` は `add doc --data` と同じ JSON 形式です。エントリはエンジンのバッチ API（バッチごとに WAL fsync 1 回）で適用され、`add doc` と異なり**自動的にコミット**します（`--commit-every` 件ごと + 最後に 1 回）。 + +```bash +laurus add docs --file [--batch-size 1000] [--commit-every 0] +``` + +**引数:** + +| フラグ | 必須 | 説明 | +| :--- | :--- | :--- | +| `--file ` | はい | 取り込む JSONL ファイルのパス | +| `--batch-size ` | いいえ | エンジンのバッチ呼び出しあたりのドキュメント数（既定 `1000`） | +| `--commit-every ` | いいえ | N 件適用ごとにコミット。`0` = 最後の 1 回のみ（既定） | + +繰り返した ID はチャンクとして蓄積されます。途中で失敗した場合、エラーは該当行を示し、適用済みの prefix はコミットされるため、残りの行から再実行してインジェストを継続できます。 + --- ## `put` — リソースの上書き（Upsert） @@ -278,12 +298,33 @@ laurus put doc --id --data **例:** ```bash -laurus put doc --id doc1 --data '{"title":"Updated Title","body":"This replaces the existing document."}' +laurus put doc --id doc1 --data '{"fields":{"title":{"Text":"Updated Title"},"body":{"Text":"This replaces the existing document."}}}' # Document 'doc1' put (upserted). Run 'commit' to persist changes. ``` > **注意:** `add doc` とは異なり、`put doc` は指定 ID の既存チャンクをすべて置き換えます。チャンクを追記したい場合は `add doc` を、ドキュメント全体を置き換えたい場合は `put doc` を使用してください。 +### `put docs` + +JSONL ファイルからドキュメントをバルク Upsert します — 1 行に 1 エントリの `{"id": "...", "document": {"fields": {...}}}` 形式で、エンジンのバッチ API（バッチごとに WAL fsync 1 回）で適用されます。重複した ID は順にデデュープされます（最後の出現が勝ち）。`add docs` と同じく**自動的にコミット**します。 + +```bash +laurus put docs --file [--batch-size 1000] [--commit-every 0] +``` + +引数は `add docs` と同じです。途中で失敗した場合、エラーは該当行を示し、適用済みの prefix はコミットされます。put は冪等なので、ファイル全体（または残りの suffix）の再実行は安全です。 + +**例:** + +```bash +cat > docs.jsonl <<'JSONL' +{"id": "doc1", "document": {"fields": {"title": {"Text": "Hello"}}}} +{"id": "doc2", "document": {"fields": {"title": {"Text": "World"}}}} +JSONL +laurus put docs --file docs.jsonl +# 2 documents put (upserted) and committed. +``` + --- ### `add field` diff --git a/docs/ja/src/laurus-mcp/tools.md b/docs/ja/src/laurus-mcp/tools.md index 2ffbe4ed..970ae258 100644 --- a/docs/ja/src/laurus-mcp/tools.md +++ b/docs/ja/src/laurus-mcp/tools.md @@ -183,6 +183,54 @@ document: {"title": "Hello World - Part 2", "body": "これは続きです。"} --- +## put_documents + +1 回の呼び出しで多数のドキュメントを Put（Upsert）します。エントリは入力順に逐次適用され、バッチ全体で WAL fsync は 1 回です — ドキュメントごとに `put_document` を呼ぶよりはるかに高速です。1 バッチ内で重複した ID はデデュープされます（最後の出現が勝ち）。実行後に `commit` を呼び出してください。 + +### パラメータ + +| 名前 | 型 | 必須 | 説明 | +| :--- | :--- | :--- | :--- | +| `documents` | array | はい | `{"id": "...", "document": {...}}` エントリの配列。各 `document` は `put_document` と同じ形式 | + +### 例 + +```text +Tool: put_documents +documents: [ + {"id": "doc-1", "document": {"title": "Hello"}}, + {"id": "doc-2", "document": {"title": "World"}} +] +``` + +結果: `2 documents put (upserted). Call commit to persist changes.` + +失敗した場合は該当エントリで中断し、適用済みの prefix はロールバックされないため、バッチ（またはその suffix）の再試行は冪等です。 + +--- + +## add_documents + +1 回の呼び出しで多数のドキュメントを新しいチャンクとして追加します。`put_documents` と異なり既存ドキュメントは削除されないため、ID を繰り返すと同一論理ドキュメントの複数チャンクになります。バッチ全体で WAL fsync は 1 回です。実行後に `commit` を呼び出してください。 + +### パラメータ + +`put_documents` と同じです。 + +### 例 + +```text +Tool: add_documents +documents: [ + {"id": "doc-1", "document": {"title": "Part 1"}}, + {"id": "doc-1", "document": {"title": "Part 2"}} +] +``` + +結果: `2 documents added as chunks. Call commit to persist changes.` + +--- + ## get_documents 外部 ID で全ドキュメント（チャンクを含む）を取得します。 diff --git a/docs/ja/src/laurus-server/grpc_api.md b/docs/ja/src/laurus-server/grpc_api.md index f83165bc..47ce667e 100644 --- a/docs/ja/src/laurus-server/grpc_api.md +++ b/docs/ja/src/laurus-server/grpc_api.md @@ -8,7 +8,7 @@ | :--- | :--- | :--- | | `HealthService` | `Check` | ヘルスチェック | | `IndexService` | `CreateIndex`, `GetIndex`, `GetSchema`, `AddField`, `DeleteField` | インデックスのライフサイクルとスキーマ | -| `DocumentService` | `PutDocument`, `AddDocument`, `GetDocuments`, `DeleteDocuments`, `Commit`, `FlushWal` | ドキュメント CRUD・コミット・WAL flush | +| `DocumentService` | `PutDocument`, `AddDocument`, `PutDocuments`, `AddDocuments`, `GetDocuments`, `DeleteDocuments`, `Commit`, `FlushWal` | ドキュメント CRUD・バルクインジェスト・コミット・WAL flush | | `SearchService` | `Search`, `SearchStream` | 単発検索とストリーミング検索 | --- @@ -282,6 +282,39 @@ rpc AddDocument(AddDocumentRequest) returns (AddDocumentResponse); リクエストフィールドは `PutDocument` と同じです。 +### `PutDocuments` + +バッチ Upsert。エントリは入力順に逐次適用され、バッチ全体で WAL fsync は 1 回です — ドキュメントごとに `PutDocument` を呼ぶよりはるかに高速です。1 バッチ内で重複した ID は、同じ put を 1 件ずつ発行した場合とまったく同じようにデデュープされます（最後の出現が勝ち）。 + +```protobuf +rpc PutDocuments(PutDocumentsRequest) returns (PutDocumentsResponse); + +message DocumentEntry { + string id = 1; + Document document = 2; +} + +message PutDocumentsRequest { + repeated DocumentEntry documents = 1; +} + +message PutDocumentsResponse { + uint32 applied = 1; // 成功時はリクエストサイズと一致 +} +``` + +適用できない最初のエントリで fail-fast します。適用済みエントリはロールバック**されず**（次のコミットで永続化）、エラーステータスのメッセージに失敗位置・その ID・適用済み件数が含まれるため、バッチ（またはその suffix）の再試行は冪等です。呼び出し側の誤り（スキーマ違反など）で失敗したバッチは `INVALID_ARGUMENT`、ストレージ障害は `INTERNAL` を返します。 + +### `AddDocuments` + +バッチチャンク追加。`PutDocuments` と同様ですが既存ドキュメントを削除しないため、同一論理ドキュメントの複数チャンクを追加する目的で ID をバッチ内で繰り返せます。 + +```protobuf +rpc AddDocuments(AddDocumentsRequest) returns (AddDocumentsResponse); +``` + +リクエスト/レスポンスのフィールドは `PutDocuments` と対になります。 + ### `GetDocuments` 指定された外部 ID に一致するすべてのドキュメントを取得します。 diff --git a/docs/ja/src/laurus-server/http_gateway.md b/docs/ja/src/laurus-server/http_gateway.md index 33bf64d2..77f044ec 100644 --- a/docs/ja/src/laurus-server/http_gateway.md +++ b/docs/ja/src/laurus-server/http_gateway.md @@ -38,6 +38,7 @@ laurus serve --config config.toml | POST | `/v1/documents/{id}` | `DocumentService/AddDocument` | ドキュメントの追加（チャンク） | | GET | `/v1/documents/{id}` | `DocumentService/GetDocuments` | ID でドキュメントを取得 | | DELETE | `/v1/documents/{id}` | `DocumentService/DeleteDocuments` | ID でドキュメントを削除 | +| POST | `/v1/documents:bulk` | `DocumentService/PutDocuments` / `AddDocuments` | ドキュメントのバルクインジェスト（`?mode=put\|add`、既定は `put`） | | POST | `/v1/commit` | `DocumentService/Commit` | 保留中の変更をコミット | | POST | `/v1/flush_wal` | `DocumentService/FlushWal` | full commit なしでバッファされた WAL レコードを durable 化 | | POST | `/v1/search` | `SearchService/Search` | 検索（単発） | @@ -141,6 +142,24 @@ curl -X POST http://localhost:8080/v1/documents/doc1 \ }' ``` +### ドキュメントのバルクインジェスト（POST） + +1 回の呼び出しで多数のドキュメントを適用します — エントリは入力順に逐次処理され、バッチ全体で WAL fsync は 1 回です。`?mode=put`（既定）は Upsert（重複 ID はデデュープ、最後が勝ち）、`?mode=add` はチャンク追加（繰り返した ID は蓄積）です。 + +```bash +curl -X POST 'http://localhost:8080/v1/documents:bulk?mode=put' \ + -H 'Content-Type: application/json' \ + -d '{ + "documents": [ + {"id": "doc1", "document": {"fields": {"title": "Hello"}}}, + {"id": "doc2", "document": {"fields": {"title": "World"}}} + ] + }' +# => {"applied": 2} +``` + +適用できない最初のエントリで fail-fast します。適用済みエントリはロールバックされず（次のコミットで永続化）、エラーには失敗位置が含まれるため、バッチまたはその suffix の再試行は冪等です。 + ### ドキュメントの取得 ```bash diff --git a/docs/src/laurus-cli/commands.md b/docs/src/laurus-cli/commands.md index 169b6a57..af74d72d 100644 --- a/docs/src/laurus-cli/commands.md +++ b/docs/src/laurus-cli/commands.md @@ -237,25 +237,45 @@ laurus add doc --id --data | `--id ` | Yes | External document ID (string) | | `--data ` | Yes | Document fields as a JSON string | -The JSON format is a flat object mapping field names to values: +The JSON is the document's serde shape: a `fields` object mapping each field name to an externally-tagged value (`Text`, `Int64`, `Float64`, `Bool`, `VectorValue`, ...): ```json { - "title": "Introduction to Rust", - "body": "Rust is a systems programming language.", - "category": "programming" + "fields": { + "title": {"Text": "Introduction to Rust"}, + "body": {"Text": "Rust is a systems programming language."}, + "year": {"Int64": 2024} + } } ``` **Example:** ```bash -laurus add doc --id doc1 --data '{"title":"Hello World","body":"This is a test document."}' +laurus add doc --id doc1 --data '{"fields":{"title":{"Text":"Hello World"},"body":{"Text":"This is a test document."}}}' # Document 'doc1' added. Run 'commit' to persist changes. ``` > **Tip:** Multiple documents can share the same external ID (chunking pattern). Use `add doc` for each chunk. +### `add docs` + +Bulk-add document chunks from a JSONL file — one `{"id": "...", "document": {"fields": {...}}}` entry per line, where `document` uses the same JSON shape as `add doc --data`. Entries are applied through the engine's batch API (one WAL fsync per batch) and, unlike `add doc`, the command **commits automatically**: every `--commit-every` applied documents and once at the end. + +```bash +laurus add docs --file [--batch-size 1000] [--commit-every 0] +``` + +**Arguments:** + +| Flag | Required | Description | +| :--- | :--- | :--- | +| `--file ` | Yes | Path to the JSONL file to ingest | +| `--batch-size ` | No | Documents per engine batch call (default `1000`) | +| `--commit-every ` | No | Commit every N applied documents; `0` = only the final commit (default) | + +Repeated IDs accumulate as chunks. On a mid-file failure the error names the offending line, the applied prefix is committed, and re-running the remaining lines continues the ingest. + --- ## `put` — Put (Upsert) a Resource @@ -278,12 +298,33 @@ laurus put doc --id --data **Example:** ```bash -laurus put doc --id doc1 --data '{"title":"Updated Title","body":"This replaces the existing document."}' +laurus put doc --id doc1 --data '{"fields":{"title":{"Text":"Updated Title"},"body":{"Text":"This replaces the existing document."}}}' # Document 'doc1' put (upserted). Run 'commit' to persist changes. ``` > **Note:** Unlike `add doc`, `put doc` replaces all existing chunks for the given ID. Use `add doc` when you want to append chunks, and `put doc` when you want to replace the entire document. +### `put docs` + +Bulk-upsert documents from a JSONL file — one `{"id": "...", "document": {"fields": {...}}}` entry per line, applied through the engine's batch API (one WAL fsync per batch). Duplicate IDs dedup in order (the last occurrence wins). Like `add docs`, the command **commits automatically**. + +```bash +laurus put docs --file [--batch-size 1000] [--commit-every 0] +``` + +Arguments are the same as `add docs`. On a mid-file failure the error names the offending line and the applied prefix is committed; because puts are idempotent, re-running the whole file (or its remaining suffix) is safe. + +**Example:** + +```bash +cat > docs.jsonl <<'JSONL' +{"id": "doc1", "document": {"fields": {"title": {"Text": "Hello"}}}} +{"id": "doc2", "document": {"fields": {"title": {"Text": "World"}}}} +JSONL +laurus put docs --file docs.jsonl +# 2 documents put (upserted) and committed. +``` + --- ### `add field` diff --git a/docs/src/laurus-mcp/tools.md b/docs/src/laurus-mcp/tools.md index c5b9c006..f7cfb79a 100644 --- a/docs/src/laurus-mcp/tools.md +++ b/docs/src/laurus-mcp/tools.md @@ -239,6 +239,54 @@ Result: `Document 'doc-1' added as chunk. Call commit to persist changes.` --- +## put_documents + +Put (upsert) many documents in one call. Entries are applied sequentially, in input order, with one WAL fsync for the whole batch — much faster than calling `put_document` per document. Duplicate IDs within one batch dedup (the last occurrence wins). Call `commit` afterwards. + +### Parameters + +| Name | Type | Required | Description | +| :--- | :--- | :--- | :--- | +| `documents` | array | Yes | Array of `{"id": "...", "document": {...}}` entries; each `document` has the same shape as `put_document`'s | + +### Example + +```text +Tool: put_documents +documents: [ + {"id": "doc-1", "document": {"title": "Hello"}}, + {"id": "doc-2", "document": {"title": "World"}} +] +``` + +Result: `2 documents put (upserted). Call commit to persist changes.` + +A failure aborts at the offending entry without rolling back the already-applied prefix, so retrying the batch (or its suffix) is idempotent. + +--- + +## add_documents + +Add many documents as new chunks in one call. Unlike `put_documents`, existing documents are never deleted, so repeating an ID adds multiple chunks of the same logical document. One WAL fsync covers the whole batch. Call `commit` afterwards. + +### Parameters + +Same as `put_documents`. + +### Example + +```text +Tool: add_documents +documents: [ + {"id": "doc-1", "document": {"title": "Part 1"}}, + {"id": "doc-1", "document": {"title": "Part 2"}} +] +``` + +Result: `2 documents added as chunks. Call commit to persist changes.` + +--- + ## get_documents Retrieve all stored documents (including chunks) by external ID. diff --git a/docs/src/laurus-server/grpc_api.md b/docs/src/laurus-server/grpc_api.md index ff3843ec..2b38f5f2 100644 --- a/docs/src/laurus-server/grpc_api.md +++ b/docs/src/laurus-server/grpc_api.md @@ -8,7 +8,7 @@ All services are defined under the `laurus.v1` protobuf package. | :--- | :--- | :--- | | `HealthService` | `Check` | Health checking | | `IndexService` | `CreateIndex`, `GetIndex`, `GetSchema`, `AddField`, `DeleteField` | Index lifecycle and schema | -| `DocumentService` | `PutDocument`, `AddDocument`, `GetDocuments`, `DeleteDocuments`, `Commit`, `FlushWal` | Document CRUD, commit, and WAL flush | +| `DocumentService` | `PutDocument`, `AddDocument`, `PutDocuments`, `AddDocuments`, `GetDocuments`, `DeleteDocuments`, `Commit`, `FlushWal` | Document CRUD, bulk ingestion, commit, and WAL flush | | `SearchService` | `Search`, `SearchStream` | Unary and streaming search | --- @@ -283,6 +283,39 @@ rpc AddDocument(AddDocumentRequest) returns (AddDocumentResponse); Request fields are the same as `PutDocument`. +### `PutDocuments` + +Batched upsert. Entries are applied sequentially, in input order, with one WAL fsync for the whole batch — much faster than one `PutDocument` call per document. Duplicate IDs within one batch dedup exactly like the same puts issued one by one (the last occurrence wins). + +```protobuf +rpc PutDocuments(PutDocumentsRequest) returns (PutDocumentsResponse); + +message DocumentEntry { + string id = 1; + Document document = 2; +} + +message PutDocumentsRequest { + repeated DocumentEntry documents = 1; +} + +message PutDocumentsResponse { + uint32 applied = 1; // equals the request size on success +} +``` + +The call fails fast at the first entry that cannot be applied. Already-applied entries are **not** rolled back — they are durable at the next commit — and the error status message carries the failing position, its ID, and the applied count, so retrying the batch (or its suffix) is idempotent. A batch that fails on a caller mistake (e.g. a schema violation) returns `INVALID_ARGUMENT`; storage failures return `INTERNAL`. + +### `AddDocuments` + +Batched chunk append. Like `PutDocuments` but never deletes existing documents, so a batch may legitimately repeat an ID to add multiple chunks of the same logical document. + +```protobuf +rpc AddDocuments(AddDocumentsRequest) returns (AddDocumentsResponse); +``` + +Request/response fields mirror `PutDocuments`. + ### `GetDocuments` Retrieve all documents matching the given external ID. diff --git a/docs/src/laurus-server/http_gateway.md b/docs/src/laurus-server/http_gateway.md index 452fa326..ccd102a0 100644 --- a/docs/src/laurus-server/http_gateway.md +++ b/docs/src/laurus-server/http_gateway.md @@ -38,6 +38,7 @@ If `http_port` is not set, only the gRPC server starts. | POST | `/v1/documents/{id}` | `DocumentService/AddDocument` | Add a document (chunk) | | GET | `/v1/documents/{id}` | `DocumentService/GetDocuments` | Get documents by ID | | DELETE | `/v1/documents/{id}` | `DocumentService/DeleteDocuments` | Delete documents by ID | +| POST | `/v1/documents:bulk` | `DocumentService/PutDocuments` / `AddDocuments` | Bulk-ingest documents (`?mode=put\|add`, default `put`) | | POST | `/v1/commit` | `DocumentService/Commit` | Commit pending changes | | POST | `/v1/flush_wal` | `DocumentService/FlushWal` | Force buffered WAL records durable without a full commit | | POST | `/v1/search` | `SearchService/Search` | Search (unary) | @@ -145,6 +146,30 @@ curl -X POST http://localhost:8080/v1/documents/doc1 \ }' ``` +### Bulk-Ingest Documents (POST) + +Applies many documents in one call — entries are processed sequentially, in +input order, with one WAL fsync for the whole batch. `?mode=put` (the +default) upserts (duplicate ids dedup, last wins); `?mode=add` appends +chunks, so repeated ids accumulate: + +```bash +curl -X POST 'http://localhost:8080/v1/documents:bulk?mode=put' \ + -H 'Content-Type: application/json' \ + -d '{ + "documents": [ + {"id": "doc1", "document": {"fields": {"title": "Hello"}}}, + {"id": "doc2", "document": {"fields": {"title": "World"}}} + ] + }' +# => {"applied": 2} +``` + +The call fails fast at the first entry that cannot be applied; +already-applied entries are not rolled back (they become durable at the next +commit), and the error names the failing position, so retrying the batch or +its suffix is idempotent. + ### Get Documents ```bash diff --git a/laurus-cli/Cargo.toml b/laurus-cli/Cargo.toml index 4bdef4c5..3e0eb650 100644 --- a/laurus-cli/Cargo.toml +++ b/laurus-cli/Cargo.toml @@ -30,6 +30,9 @@ tabled = { workspace = true } tokio = { workspace = true } toml = { workspace = true } +[dev-dependencies] +tempfile = { workspace = true } + [features] embeddings-candle = [ "laurus/embeddings-candle", diff --git a/laurus-cli/src/cli.rs b/laurus-cli/src/cli.rs index 463904d3..4f06c7b1 100644 --- a/laurus-cli/src/cli.rs +++ b/laurus-cli/src/cli.rs @@ -127,6 +127,23 @@ pub enum AddResource { #[arg(long)] data: String, }, + /// Bulk-add document chunks from a JSONL file (one entry per line). + /// + /// Each line is `{"id": "...", "document": {"fields": {...}}}` — the + /// same document JSON shape as `add doc --data`. Unlike `put docs`, + /// repeated ids accumulate as chunks. Commits automatically (per + /// `--commit-every` and once at the end). + Docs { + /// Path to the JSONL file to ingest. + #[arg(long)] + file: std::path::PathBuf, + /// Documents per engine batch call. + #[arg(long, default_value_t = 1000)] + batch_size: usize, + /// Commit every N applied documents (0 = only the final commit). + #[arg(long, default_value_t = 0)] + commit_every: usize, + }, /// Dynamically add a new field to an existing index. Field { /// The name of the new field. @@ -169,6 +186,23 @@ pub enum PutResource { #[arg(long)] data: String, }, + /// Bulk-upsert documents from a JSONL file (one entry per line). + /// + /// Each line is `{"id": "...", "document": {"fields": {...}}}` — the + /// same document JSON shape as `put doc --data`. Entries are applied in + /// order (duplicate ids dedup, last wins) with one WAL fsync per batch. + /// Commits automatically (per `--commit-every` and once at the end). + Docs { + /// Path to the JSONL file to ingest. + #[arg(long)] + file: std::path::PathBuf, + /// Documents per engine batch call. + #[arg(long, default_value_t = 1000)] + batch_size: usize, + /// Commit every N applied documents (0 = only the final commit). + #[arg(long, default_value_t = 0)] + commit_every: usize, + }, } // --- Delete --- diff --git a/laurus-cli/src/commands.rs b/laurus-cli/src/commands.rs index c7c60f85..d90868d7 100644 --- a/laurus-cli/src/commands.rs +++ b/laurus-cli/src/commands.rs @@ -14,6 +14,7 @@ //! - [`serve`] - gRPC (and optional HTTP gateway) server. pub mod add; +pub mod bulk; pub mod commit; pub mod create; pub mod delete; diff --git a/laurus-cli/src/commands/bulk.rs b/laurus-cli/src/commands/bulk.rs new file mode 100644 index 00000000..6fd3e7c5 --- /dev/null +++ b/laurus-cli/src/commands/bulk.rs @@ -0,0 +1,251 @@ +//! Implementations for the `put docs` / `add docs` bulk-ingest subcommands. +//! +//! Streams a JSONL file — one `{"id": "...", "document": {"fields": {...}}}` +//! entry per line — and applies it through the engine's batch API +//! ([`Engine::put_documents`](laurus::Engine::put_documents) / +//! [`Engine::add_documents`](laurus::Engine::add_documents)), which pays one +//! WAL fsync per batch instead of one per record. Unlike the singular +//! `put doc` / `add doc` commands, bulk ingestion commits automatically: +//! every `--commit-every` applied documents and once at the end. + +use std::io::BufRead; +use std::path::Path; + +use anyhow::{Context, Result, bail}; +use laurus::Document; + +use crate::context; + +/// Whether a bulk run upserts (`put docs`) or appends chunks (`add docs`). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BulkMode { + /// Upsert: duplicate ids replace earlier versions (last occurrence wins). + Put, + /// Chunk append: repeated ids accumulate as chunks. + Add, +} + +/// One parsed JSONL entry. +#[derive(serde::Deserialize)] +struct BulkEntry { + /// External document ID. + id: String, + /// The document, in the same serde JSON shape as `put doc --data` + /// (`{"fields": {"title": {"Text": "..."}}}`). + document: Document, +} + +/// Parse one JSONL line into an `(id, document)` pair. +/// +/// # Arguments +/// +/// * `line` - The raw line content (must not be blank). +/// * `line_no` - 1-based line number, used in error messages. +/// +/// # Errors +/// +/// Returns an error naming the line when the JSON does not parse into a +/// `{"id", "document"}` entry. +fn parse_entry(line: &str, line_no: usize) -> Result<(String, Document)> { + let entry: BulkEntry = serde_json::from_str(line) + .with_context(|| format!("line {line_no}: failed to parse JSONL entry"))?; + Ok((entry.id, entry.document)) +} + +/// Execute the `put docs` / `add docs` command. +/// +/// Reads `file` line by line (blank lines are skipped), groups entries into +/// batches of `batch_size`, and applies each batch via the engine's bulk API. +/// Commits after every `commit_every` applied documents (`0` disables the +/// periodic commits) and once at the end, then prints the applied count. +/// +/// On a mid-batch failure the engine's fail-fast semantics apply: the error +/// message names the offending **line** of the input file, the applied prefix +/// is committed (so the work is not lost), and re-running with the remaining +/// suffix of the file is idempotent under `put` mode. +/// +/// # Arguments +/// +/// * `file` - Path to the JSONL file to ingest. +/// * `mode` - [`BulkMode::Put`] (upsert) or [`BulkMode::Add`] (chunk append). +/// * `batch_size` - Documents per engine batch call (must be > 0). +/// * `commit_every` - Commit every N applied documents; `0` = final only. +/// * `index_dir` - Path to the index directory holding the index. +/// +/// # Errors +/// +/// Returns an error if the index cannot be opened, the file cannot be read, +/// a line fails to parse, or the engine rejects a batch (the message carries +/// the failing line number and the count applied before the failure). +pub async fn run( + file: &Path, + mode: BulkMode, + batch_size: usize, + commit_every: usize, + index_dir: &Path, +) -> Result<()> { + if batch_size == 0 { + bail!("--batch-size must be greater than 0"); + } + + let engine = context::open_index(index_dir).await?; + let reader = std::io::BufReader::new( + std::fs::File::open(file) + .with_context(|| format!("failed to open JSONL file '{}'", file.display()))?, + ); + + let mut batch: Vec<(String, Document)> = Vec::with_capacity(batch_size); + // 1-based input line number of each entry in `batch`, for error reports. + let mut batch_lines: Vec = Vec::with_capacity(batch_size); + let mut applied: usize = 0; + let mut since_commit: usize = 0; + + let flush = async |batch: Vec<(String, Document)>, + lines: Vec, + already_applied: usize| + -> Result { + if batch.is_empty() { + return Ok(0); + } + let size = batch.len(); + let result = match mode { + BulkMode::Put => engine.put_documents(batch).await, + BulkMode::Add => engine.add_documents(batch).await, + }; + if let Err(e) = result { + // Commit the applied prefix so completed work survives, then + // surface the failing input line. + engine.commit().await.ok(); + if let laurus::LaurusError::BatchIngest { + failed_index, + failed_id, + applied: batch_applied, + .. + } = &e + { + let line = lines.get(*failed_index).copied().unwrap_or(0); + bail!( + "line {line} (id '{failed_id}'): {e} — {} documents were applied and \ + committed before the failure", + already_applied + batch_applied + ); + } + return Err(e.into()); + } + Ok(size) + }; + + for (index, line) in reader.lines().enumerate() { + let line_no = index + 1; + let line = line.with_context(|| format!("line {line_no}: failed to read"))?; + if line.trim().is_empty() { + continue; + } + let (id, doc) = parse_entry(&line, line_no)?; + batch.push((id, doc)); + batch_lines.push(line_no); + + if batch.len() >= batch_size { + let docs = std::mem::take(&mut batch); + let lines = std::mem::take(&mut batch_lines); + let flushed = flush(docs, lines, applied).await?; + applied += flushed; + since_commit += flushed; + if commit_every > 0 && since_commit >= commit_every { + engine.commit().await?; + since_commit = 0; + println!("... {applied} documents applied (committed)"); + } + } + } + applied += flush(batch, batch_lines, applied).await?; + engine.commit().await?; + + let verb = match mode { + BulkMode::Put => "put (upserted)", + BulkMode::Add => "added as chunks", + }; + println!("{applied} documents {verb} and committed."); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_entry_accepts_the_put_doc_data_shape() { + let line = r#"{"id": "doc1", "document": {"fields": {"title": {"Text": "Hello"}}}}"#; + let (id, doc) = parse_entry(line, 1).unwrap(); + assert_eq!(id, "doc1"); + assert!(doc.fields.contains_key("title")); + } + + #[test] + fn parse_entry_names_the_failing_line() { + let err = parse_entry("{not json", 42).unwrap_err(); + assert!( + err.to_string().contains("line 42"), + "error must carry the line number: {err}" + ); + + let err = parse_entry(r#"{"document": {"fields": {}}}"#, 7).unwrap_err(); + assert!( + err.to_string().contains("line 7"), + "a missing id must also name the line: {err}" + ); + } + + /// End-to-end: `put docs` ingests a JSONL file (small batches + periodic + /// commits + in-batch dedup) and the docs are retrievable afterwards. + #[tokio::test] + async fn bulk_put_ingests_jsonl_end_to_end() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write( + dir.path().join("schema.toml"), + "[fields.title.Text]\nstored = true\nindexed = true\n", + ) + .unwrap(); + + let jsonl_path = dir.path().join("docs.jsonl"); + let mut jsonl = String::new(); + for i in 0..7 { + jsonl.push_str(&format!( + "{{\"id\": \"doc{i}\", \"document\": {{\"fields\": {{\"title\": {{\"Text\": \"t{i}\"}}}}}}}}\n", + )); + } + // Blank line + a duplicate id (must dedup, last wins under put mode). + jsonl.push('\n'); + jsonl.push_str( + "{\"id\": \"doc0\", \"document\": {\"fields\": {\"title\": {\"Text\": \"t0v2\"}}}}\n", + ); + std::fs::write(&jsonl_path, jsonl).unwrap(); + + // batch_size 3 exercises multiple flushes; commit_every 5 exercises + // the periodic commit path. + run(&jsonl_path, BulkMode::Put, 3, 5, dir.path()) + .await + .unwrap(); + + let engine = context::open_index(dir.path()).await.unwrap(); + let docs = engine.get_documents("doc0").await.unwrap(); + assert_eq!(docs.len(), 1, "duplicate id must dedup under put mode"); + let title = docs[0] + .fields + .get("title") + .and_then(|v| v.as_text()) + .unwrap(); + assert_eq!(title, "t0v2", "the last occurrence must win"); + for i in 1..7 { + assert_eq!( + engine + .get_documents(&format!("doc{i}")) + .await + .unwrap() + .len(), + 1, + "doc{i} must be ingested" + ); + } + } +} diff --git a/laurus-cli/src/main.rs b/laurus-cli/src/main.rs index bfc20682..e4c82a78 100644 --- a/laurus-cli/src/main.rs +++ b/laurus-cli/src/main.rs @@ -26,7 +26,7 @@ use clap::Parser; use crate::cli::{ AddResource, Cli, Command, CreateResource, DeleteResource, GetResource, McpCommand, PutResource, }; -use crate::commands::{add, commit, create, delete, get, mcp, put, repl, search, serve}; +use crate::commands::{add, bulk, commit, create, delete, get, mcp, put, repl, search, serve}; #[tokio::main] async fn main() -> Result<()> { @@ -48,12 +48,40 @@ async fn main() -> Result<()> { }, Command::Add(cmd) => match cmd.resource { AddResource::Doc { id, data } => add::run_doc(&id, &data, &index_dir).await, + AddResource::Docs { + file, + batch_size, + commit_every, + } => { + bulk::run( + &file, + bulk::BulkMode::Add, + batch_size, + commit_every, + &index_dir, + ) + .await + } AddResource::Field { name, field_option } => { add::run_field(&name, &field_option, &index_dir).await } }, Command::Put(cmd) => match cmd.resource { PutResource::Doc { id, data } => put::run_doc(&id, &data, &index_dir).await, + PutResource::Docs { + file, + batch_size, + commit_every, + } => { + bulk::run( + &file, + bulk::BulkMode::Put, + batch_size, + commit_every, + &index_dir, + ) + .await + } }, Command::Delete(cmd) => match cmd.resource { DeleteResource::Docs { id } => delete::run_docs(&id, &index_dir).await, diff --git a/laurus-mcp/src/server.rs b/laurus-mcp/src/server.rs index 06847deb..bc501985 100644 --- a/laurus-mcp/src/server.rs +++ b/laurus-mcp/src/server.rs @@ -19,9 +19,10 @@ use tonic::transport::Channel; use tracing::info; use laurus_server::proto::laurus::v1::{ - AddDocumentRequest, AddFieldRequest, CommitRequest, CreateIndexRequest, DeleteDocumentsRequest, - DeleteFieldRequest, GetDocumentsRequest, GetIndexRequest, GetSchemaRequest, PutDocumentRequest, - SearchBatchRequest, SearchRequest, document_service_client::DocumentServiceClient, + AddDocumentRequest, AddDocumentsRequest, AddFieldRequest, CommitRequest, CreateIndexRequest, + DeleteDocumentsRequest, DeleteFieldRequest, DocumentEntry, GetDocumentsRequest, + GetIndexRequest, GetSchemaRequest, PutDocumentRequest, PutDocumentsRequest, SearchBatchRequest, + SearchRequest, document_service_client::DocumentServiceClient, index_service_client::IndexServiceClient, search_service_client::SearchServiceClient, }; @@ -95,6 +96,17 @@ struct AddDocumentParams { document: Value, } +/// Parameters for the `put_documents` / `add_documents` tools. +#[derive(Debug, Deserialize, schemars::JsonSchema)] +struct BulkDocumentsParams { + /// Array of `{"id": "...", "document": {...}}` entries, applied + /// sequentially in input order with one WAL fsync for the whole batch. + /// + /// Each `document` is a JSON object of fields matching the index schema + /// (the same shape as the `put_document` tool's `document`). + documents: Vec, +} + /// Parameters for the `get_documents` tool. #[derive(Debug, Deserialize, schemars::JsonSchema)] struct GetDocumentsParams { @@ -557,6 +569,107 @@ impl LaurusMcpServer { } } + /// Convert the bulk tools' JSON entries into proto `DocumentEntry` + /// values, naming the offending position on the first invalid entry. + fn bulk_entries(documents: Vec) -> Result, String> { + documents + .into_iter() + .enumerate() + .map(|(index, entry)| { + let id = entry + .get("id") + .and_then(Value::as_str) + .ok_or_else(|| format!("documents[{index}]: missing string \"id\""))? + .to_string(); + let document = entry + .get("document") + .cloned() + .ok_or_else(|| format!("documents[{index}]: missing \"document\" key"))?; + let document = convert::json_to_document(document) + .map_err(|e| format!("documents[{index}]: {e}"))?; + Ok(DocumentEntry { + id, + document: Some(document), + }) + }) + .collect() + } + + /// Batched upsert of documents in one round trip. + /// + /// Entries are applied sequentially, in input order, with one WAL fsync + /// for the whole batch; a failure aborts at the offending entry without + /// rolling back the already-applied prefix (retrying is idempotent). + #[tool( + description = "Put (upsert) MANY documents in one call. Pass documents as an array of {\"id\": \"...\", \"document\": {...}} entries; they are applied in order (duplicate ids dedup, last wins) with one WAL fsync for the whole batch, which is much faster than calling put_document per document. Call commit afterwards to persist changes." + )] + async fn put_documents( + &self, + Parameters(params): Parameters, + ) -> Result { + let channel = match self.channel.read().await.clone() { + Some(ch) => ch, + None => { + return Ok(Self::tool_error( + "Not connected. Call the connect tool first.", + )); + } + }; + + let documents = match Self::bulk_entries(params.documents) { + Ok(entries) => entries, + Err(e) => return Ok(Self::tool_error(format!("Invalid batch: {e}"))), + }; + + match DocumentServiceClient::new(channel) + .put_documents(PutDocumentsRequest { documents }) + .await + { + Ok(resp) => Ok(CallToolResult::success(vec![Content::text(format!( + "{} documents put (upserted). Call commit to persist changes.", + resp.into_inner().applied + ))])), + Err(e) => Ok(Self::tool_error(format!("Failed to put documents: {e}"))), + } + } + + /// Batched chunk append of documents in one round trip. + /// + /// Like `put_documents` but never deletes existing documents, so a batch + /// may repeat an id to add multiple chunks of one logical document. + #[tool( + description = "Add MANY documents as new chunks in one call. Pass documents as an array of {\"id\": \"...\", \"document\": {...}} entries; unlike put_documents, existing documents are never deleted, so repeating an id adds multiple chunks. One WAL fsync covers the whole batch. Call commit afterwards to persist changes." + )] + async fn add_documents( + &self, + Parameters(params): Parameters, + ) -> Result { + let channel = match self.channel.read().await.clone() { + Some(ch) => ch, + None => { + return Ok(Self::tool_error( + "Not connected. Call the connect tool first.", + )); + } + }; + + let documents = match Self::bulk_entries(params.documents) { + Ok(entries) => entries, + Err(e) => return Ok(Self::tool_error(format!("Invalid batch: {e}"))), + }; + + match DocumentServiceClient::new(channel) + .add_documents(AddDocumentsRequest { documents }) + .await + { + Ok(resp) => Ok(CallToolResult::success(vec![Content::text(format!( + "{} documents added as chunks. Call commit to persist changes.", + resp.into_inner().applied + ))])), + Err(e) => Ok(Self::tool_error(format!("Failed to add documents: {e}"))), + } + } + /// Get all stored documents for a given ID. #[tool( description = "Retrieve all stored documents (including chunks) by external ID. Returns a JSON array of documents matching the ID." @@ -816,7 +929,8 @@ impl ServerHandler for LaurusMcpServer { .with_instructions( "Laurus search engine MCP server (gRPC client). \ Tools: connect, create_index, get_stats, get_schema, add_field, delete_field, \ - put_document, add_document, get_documents, delete_documents, commit, search. \ + put_document, add_document, put_documents, add_documents, get_documents, \ + delete_documents, commit, search. \ Start by calling connect(endpoint) to connect to a running laurus-server, \ then use the other tools to manage and search the index." .to_string(), diff --git a/laurus-server/proto/laurus/v1/document.proto b/laurus-server/proto/laurus/v1/document.proto index 691582bb..6736f2f6 100644 --- a/laurus-server/proto/laurus/v1/document.proto +++ b/laurus-server/proto/laurus/v1/document.proto @@ -11,6 +11,17 @@ service DocumentService { // Add a document as a new chunk (multiple chunks can share an ID). rpc AddDocument(AddDocumentRequest) returns (AddDocumentResponse); + // Batched upsert: apply (id, document) entries sequentially, in input + // order, with one WAL fsync for the whole batch. Fails fast at the first + // entry that cannot be applied; already-applied entries are not rolled + // back, so retrying the batch (or its suffix) is idempotent. + rpc PutDocuments(PutDocumentsRequest) returns (PutDocumentsResponse); + + // Batched chunk append: like PutDocuments but never deletes existing + // documents, so a batch may legitimately repeat an ID to add multiple + // chunks of the same logical document. + rpc AddDocuments(AddDocumentsRequest) returns (AddDocumentsResponse); + // Get all documents (including chunks) by external ID. rpc GetDocuments(GetDocumentsRequest) returns (GetDocumentsResponse); @@ -40,6 +51,32 @@ message AddDocumentRequest { message AddDocumentResponse {} +// One (external id, document) pair of a batched write. +message DocumentEntry { + string id = 1; + Document document = 2; +} + +message PutDocumentsRequest { + // Applied sequentially, in order. Duplicate ids within one batch dedup + // exactly like the same puts issued one by one (the last occurrence wins). + repeated DocumentEntry documents = 1; +} + +message PutDocumentsResponse { + // Number of documents applied (equals the request size on success). + uint32 applied = 1; +} + +message AddDocumentsRequest { + repeated DocumentEntry documents = 1; +} + +message AddDocumentsResponse { + // Number of documents applied (equals the request size on success). + uint32 applied = 1; +} + message GetDocumentsRequest { string id = 1; } diff --git a/laurus-server/src/convert/error.rs b/laurus-server/src/convert/error.rs index 08004375..0c4c3460 100644 --- a/laurus-server/src/convert/error.rs +++ b/laurus-server/src/convert/error.rs @@ -9,15 +9,26 @@ use tonic::Status; /// Convert a LaurusError into a tonic Status. pub fn to_status(err: LaurusError) -> Status { - match &err { - LaurusError::Schema(_) | LaurusError::Query(_) | LaurusError::Field(_) => { - Status::invalid_argument(err.to_string()) - } - LaurusError::SerializationError(_) | LaurusError::Json(_) => { - Status::invalid_argument(err.to_string()) - } - LaurusError::NotImplemented(_) => Status::unimplemented(err.to_string()), - _ => Status::internal(err.to_string()), + Status::new(code_for(&err), err.to_string()) +} + +/// Classify a [`LaurusError`] into a gRPC status code. +/// +/// [`LaurusError::BatchIngest`] is classified by its **source** — a batch +/// that failed on a caller mistake (e.g. a schema violation at position k) +/// surfaces as `INVALID_ARGUMENT`, while a storage failure stays `INTERNAL`. +/// The status *message* still carries the full batch context +/// (`failed_index`, `failed_id`, `applied`) from the error's `Display`. +fn code_for(err: &LaurusError) -> tonic::Code { + match err { + LaurusError::Schema(_) + | LaurusError::Query(_) + | LaurusError::Field(_) + | LaurusError::SerializationError(_) + | LaurusError::Json(_) => tonic::Code::InvalidArgument, + LaurusError::NotImplemented(_) => tonic::Code::Unimplemented, + LaurusError::BatchIngest { source, .. } => code_for(source), + _ => tonic::Code::Internal, } } @@ -25,3 +36,35 @@ pub fn to_status(err: LaurusError) -> Status { pub fn anyhow_to_status(err: anyhow::Error) -> Status { Status::internal(err.to_string()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn batch_ingest_is_classified_by_its_source() { + let caller_fault = LaurusError::BatchIngest { + failed_index: 3, + failed_id: "bad".into(), + applied: 3, + source: Box::new(LaurusError::schema("undeclared field")), + }; + let status = to_status(caller_fault); + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert!( + status.message().contains("doc 3") + && status.message().contains("'bad'") + && status.message().contains("3 documents were applied"), + "the status message must keep the batch context: {}", + status.message() + ); + + let server_fault = LaurusError::BatchIngest { + failed_index: 0, + failed_id: "x".into(), + applied: 0, + source: Box::new(LaurusError::storage("disk on fire")), + }; + assert_eq!(to_status(server_fault).code(), tonic::Code::Internal); + } +} diff --git a/laurus-server/src/gateway.rs b/laurus-server/src/gateway.rs index 2ba02200..d2eca634 100644 --- a/laurus-server/src/gateway.rs +++ b/laurus-server/src/gateway.rs @@ -55,6 +55,7 @@ pub fn create_router(state: GatewayState) -> Router { .get(document::get_documents) .delete(document::delete_documents), ) + .route("/v1/documents:bulk", post(document::bulk_documents)) .route("/v1/commit", post(document::commit)) .route("/v1/flush_wal", post(document::flush_wal)) .route("/v1/search", post(search::search)) @@ -62,3 +63,17 @@ pub fn create_router(state: GatewayState) -> Router { .route("/v1/search/batch", post(search::search_batch)) .with_state(state) } + +#[cfg(test)] +mod tests { + use super::*; + + /// The router must build without panicking — in particular the literal + /// `:bulk` suffix in `/v1/documents:bulk` must be accepted by axum 0.8's + /// `{param}` syntax (a bare colon is a plain path character there). + #[tokio::test] + async fn router_builds_with_bulk_route() { + let channel = tonic::transport::Endpoint::from_static("http://127.0.0.1:1").connect_lazy(); + let _router = create_router(GatewayState::new(channel)); + } +} diff --git a/laurus-server/src/gateway/document.rs b/laurus-server/src/gateway/document.rs index 2af159ae..d53cd063 100644 --- a/laurus-server/src/gateway/document.rs +++ b/laurus-server/src/gateway/document.rs @@ -1,7 +1,7 @@ //! Document CRUD endpoints. use axum::Json; -use axum::extract::{Path, State}; +use axum::extract::{Path, Query, State}; use axum::response::{IntoResponse, Response}; use serde_json::{Value, json}; @@ -58,6 +58,78 @@ pub async fn add_document( Ok(Json(json!({}))) } +/// `POST /v1/documents:bulk?mode=put|add` — Batched document ingestion. +/// +/// Body shape: `{"documents": [{"id": "...", "document": {"fields": {...}}}, ...]}`. +/// Entries are applied sequentially, in input order, with one WAL fsync for +/// the whole batch (see the core `Engine::put_documents` semantics): +/// `mode=put` (the default) upserts — duplicate ids within one batch dedup, +/// last occurrence wins — while `mode=add` appends chunks, so repeated ids +/// accumulate. Fails fast at the first entry that cannot be applied; +/// already-applied entries are not rolled back, and the error carries the +/// failing position, so retrying the batch (or its suffix) is idempotent. +/// Responds `{"applied": N}`. +pub async fn bulk_documents( + State(mut state): State, + Query(params): Query>, + Json(body): Json, +) -> Result, Response> { + let mode = params.get("mode").map(String::as_str).unwrap_or("put"); + if mode != "put" && mode != "add" { + return Err( + BadRequest(format!("invalid mode '{mode}' (expected 'put' or 'add')")).into_response(), + ); + } + + let entries = body + .get("documents") + .ok_or_else(|| BadRequest("missing \"documents\" key".to_string()).into_response())? + .as_array() + .ok_or_else(|| BadRequest("\"documents\" must be an array".to_string()).into_response())?; + + let documents = entries + .iter() + .enumerate() + .map(|(index, entry)| { + let id = entry + .get("id") + .and_then(Value::as_str) + .ok_or_else(|| format!("documents[{index}]: missing string \"id\""))?; + let document = convert::json_to_proto_document( + entry + .get("document") + .ok_or_else(|| format!("documents[{index}]: missing \"document\" key"))?, + ) + .map_err(|e| format!("documents[{index}]: {e}"))?; + Ok(v1::DocumentEntry { + id: id.to_string(), + document: Some(document), + }) + }) + .collect::, String>>() + .map_err(|e| BadRequest(e).into_response())?; + + let applied = if mode == "add" { + state + .document_client + .add_documents(v1::AddDocumentsRequest { documents }) + .await + .map_err(|s| GatewayError(s).into_response())? + .into_inner() + .applied + } else { + state + .document_client + .put_documents(v1::PutDocumentsRequest { documents }) + .await + .map_err(|s| GatewayError(s).into_response())? + .into_inner() + .applied + }; + + Ok(Json(json!({ "applied": applied }))) +} + /// `GET /v1/documents/:id` — Retrieves documents with the specified ID. pub async fn get_documents( State(mut state): State, diff --git a/laurus-server/src/service/document.rs b/laurus-server/src/service/document.rs index f17e33a3..988d534c 100644 --- a/laurus-server/src/service/document.rs +++ b/laurus-server/src/service/document.rs @@ -12,9 +12,10 @@ use laurus::Engine; use crate::convert::{document as doc_convert, error}; use crate::proto::laurus::v1::{ - AddDocumentRequest, AddDocumentResponse, CommitRequest, CommitResponse, DeleteDocumentsRequest, - DeleteDocumentsResponse, FlushWalRequest, FlushWalResponse, GetDocumentsRequest, - GetDocumentsResponse, PutDocumentRequest, PutDocumentResponse, + AddDocumentRequest, AddDocumentResponse, AddDocumentsRequest, AddDocumentsResponse, + CommitRequest, CommitResponse, DeleteDocumentsRequest, DeleteDocumentsResponse, DocumentEntry, + FlushWalRequest, FlushWalResponse, GetDocumentsRequest, GetDocumentsResponse, + PutDocumentRequest, PutDocumentResponse, PutDocumentsRequest, PutDocumentsResponse, document_service_server::DocumentService as DocumentServiceTrait, }; @@ -33,6 +34,32 @@ impl DocumentService { .as_ref() .ok_or_else(|| Status::failed_precondition("No index is open. Create an index first.")) } + + /// Convert a batched request's entries into the engine's + /// `(external_id, Document)` pairs, rejecting entries without a document. + /// + /// # Errors + /// + /// Returns `Status::invalid_argument` naming the offending position when + /// an entry carries no document. + #[allow(clippy::result_large_err)] + fn entries_to_docs( + entries: Vec, + ) -> Result, Status> { + entries + .into_iter() + .enumerate() + .map(|(index, entry)| { + let doc = entry.document.as_ref().ok_or_else(|| { + Status::invalid_argument(format!( + "documents[{index}] (id '{}'): document is required", + entry.id + )) + })?; + Ok((entry.id, doc_convert::from_proto(doc))) + }) + .collect() + } } #[tonic::async_trait] @@ -81,6 +108,45 @@ impl DocumentServiceTrait for DocumentService { Ok(Response::new(AddDocumentResponse {})) } + /// Batched upsert: applies the entries sequentially, in input order, with + /// one WAL fsync for the whole batch (see `Engine::put_documents`). + /// + /// Fails fast at the first entry that cannot be applied — already-applied + /// entries are not rolled back, and the returned status message carries + /// the failing position, its id, and the applied count, so clients can + /// retry the batch (or its suffix) idempotently. + async fn put_documents( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let docs = Self::entries_to_docs(req.documents)?; + let applied = docs.len() as u32; + + let guard = self.engine.read().await; + let engine = Self::get_engine_ref(&guard)?; + engine.put_documents(docs).await.map_err(error::to_status)?; + + Ok(Response::new(PutDocumentsResponse { applied })) + } + + /// Batched chunk append: like `put_documents` but never deletes existing + /// documents, so a batch may repeat an ID to add multiple chunks. + async fn add_documents( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let docs = Self::entries_to_docs(req.documents)?; + let applied = docs.len() as u32; + + let guard = self.engine.read().await; + let engine = Self::get_engine_ref(&guard)?; + engine.add_documents(docs).await.map_err(error::to_status)?; + + Ok(Response::new(AddDocumentsResponse { applied })) + } + /// Retrieves documents matching the given ID. async fn get_documents( &self, @@ -191,4 +257,101 @@ mod tests { .await .expect("flush_wal must act as a durability barrier under group policy"); } + + fn entry(id: &str, title: &str) -> DocumentEntry { + let doc = laurus::Document::builder() + .add_field("title", laurus::DataValue::Text(title.into())) + .build(); + DocumentEntry { + id: id.to_string(), + document: Some(doc_convert::to_proto(&doc)), + } + } + + /// #865: PutDocuments applies every entry (dedup included) and reports + /// the applied count; the docs are retrievable after a commit. + #[tokio::test] + async fn put_documents_applies_batch_and_reports_count() { + let service = service_with_engine(WalSyncPolicy::PerRecord).await; + + let documents = vec![entry("a", "one"), entry("b", "two"), entry("a", "one-v2")]; + let resp = service + .put_documents(Request::new(PutDocumentsRequest { documents })) + .await + .unwrap() + .into_inner(); + assert_eq!(resp.applied, 3); + + service + .commit(Request::new(CommitRequest {})) + .await + .unwrap(); + let docs = service + .get_documents(Request::new(GetDocumentsRequest { id: "a".into() })) + .await + .unwrap() + .into_inner(); + assert_eq!( + docs.documents.len(), + 1, + "the duplicate id within the batch must dedup (last wins)" + ); + } + + /// #865: AddDocuments never dedups — repeated ids accumulate as chunks. + #[tokio::test] + async fn add_documents_accumulates_chunks() { + let service = service_with_engine(WalSyncPolicy::PerRecord).await; + + let documents = vec![entry("doc", "chunk-0"), entry("doc", "chunk-1")]; + let resp = service + .add_documents(Request::new(AddDocumentsRequest { documents })) + .await + .unwrap() + .into_inner(); + assert_eq!(resp.applied, 2); + + service + .commit(Request::new(CommitRequest {})) + .await + .unwrap(); + let docs = service + .get_documents(Request::new(GetDocumentsRequest { id: "doc".into() })) + .await + .unwrap() + .into_inner(); + assert_eq!(docs.documents.len(), 2); + } + + /// #865: an empty batch succeeds with applied == 0; an entry without a + /// document is rejected up front naming its position. + #[tokio::test] + async fn put_documents_validates_entries() { + let service = service_with_engine(WalSyncPolicy::PerRecord).await; + + let resp = service + .put_documents(Request::new(PutDocumentsRequest { documents: vec![] })) + .await + .unwrap() + .into_inner(); + assert_eq!(resp.applied, 0); + + let documents = vec![ + entry("ok", "fine"), + DocumentEntry { + id: "broken".into(), + document: None, + }, + ]; + let status = service + .put_documents(Request::new(PutDocumentsRequest { documents })) + .await + .unwrap_err(); + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert!( + status.message().contains("documents[1]") && status.message().contains("'broken'"), + "the error must name the offending entry: {}", + status.message() + ); + } }