herocode/herodb
12
0
Fork 0
Code Issues 6 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

lancedb_impl #15

Open
maximevanhees wants to merge 7 commits from lancedb_impl into main
pull from: lancedb_impl
merge into: herocode:main
Conversation 0 Commits 7 Files Changed 17 +8205 -102
17 changed files with 8205 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4440
Cargo.lock generated
View File

File diff suppressed because it is too large Load Diff

8
Cargo.toml
View File

@@ -27,6 +27,14 @@ x25519-dalek = "2"
base64 = "0.22"
jsonrpsee = { version = "0.26.0", features = ["http-client", "ws-client", "server", "macros"] }
tantivy = "0.25.0"
arrow-schema = "55.2.0"
arrow-array = "55.2.0"
lance = "0.37.0"
lance-index = "0.37.0"
arrow = "55.2.0"
lancedb = "0.22.1"
uuid = "1.18.1"
ureq = { version = "2.10.0", features = ["json", "tls"] }
[dev-dependencies]
redis = { version = "0.24", features = ["aio", "tokio-comp"] }

12
README.md
View File

@@ -47,18 +47,24 @@ HeroDB can be interacted with using any standard Redis client, such as `redis-cl
### Example with `redis-cli`
Connections start with no database selected. You must SELECT a database first.
- To work in the admin database (DB 0), authenticate with the admin secret:
```bash
redis-cli -p 6379 SELECT 0 KEY myadminsecret
redis-cli -p 6379 SET mykey "Hello from HeroDB!"
redis-cli -p 6379 GET mykey
# → "Hello from HeroDB!"
```
- To use a user database, first create one via the JSON-RPC API (see docs/rpc_examples.md), then select it:
```bash
# Suppose RPC created database id 1
redis-cli -p 6379 SELECT 1
redis-cli -p 6379 HSET user:1 name "Alice" age "30"
redis-cli -p 6379 HGET user:1 name
# → "Alice"
redis-cli -p 6379 SCAN 0 MATCH user:* COUNT 10
# → 1) "0"
# 2) 1) "user:1"
```
## Cryptography

1
docs/admin.md
View File

@@ -80,6 +80,7 @@ Keys in `DB 0` (internal layout, but useful to understand how things work):
- Requires the exact admin secret as the `KEY` argument to `SELECT 0`
- Permission is `ReadWrite` when the secret matches
Connections start with no database selected. Any command that requires storage (GET, SET, H*, L*, SCAN, etc.) will return an error until you issue a SELECT to choose a database. Admin DB 0 is never accessible without authenticating via SELECT 0 KEY <admin_secret>.
### How to select databases with optional `KEY`
- Public DB (no key required)

10
docs/cmds.md
View File

@@ -126,7 +126,9 @@ redis-cli -p 6381 --pipe < dump.rdb
## Authentication and Database Selection
HeroDB uses an `Admin DB 0` to govern database existence, access and per-db encryption. Access control is enforced via `Admin DB 0` metadata. See the full model in `docs/admin.md`.
Connections start with no database selected. Any storage-backed command (GET, SET, H*, L*, SCAN, etc.) will return an error until you issue a SELECT to choose a database.
HeroDB uses an `Admin DB 0` to govern database existence, access and per-db encryption. Access control is enforced via `Admin DB 0` metadata. See the full model in (docs/admin.md:1).
Examples:
```bash
@@ -146,3 +148,9 @@ redis-cli -p $PORT SELECT 2 KEY my-db2-access-key
redis-cli -p $PORT SELECT 0 KEY my-admin-secret
# → OK
```
```bash
# Before selecting a DB, storage commands will fail
redis-cli -p $PORT GET key
# → -ERR No database selected. Use SELECT <id> [KEY <key>] first
```

444
docs/lance.md Normal file
View File

@@ -0,0 +1,444 @@
# Lance Vector Backend (RESP + JSON-RPC)
This document explains how to use HeroDB’s Lance-backed vector store. It is text-first: users provide text, and HeroDB computes embeddings server-side (no manual vectors). It includes copy-pasteable RESP (redis-cli) and JSON-RPC examples for:
- Creating a Lance database
- Embedding provider configuration (OpenAI, Azure OpenAI, or deterministic test provider)
- Dataset lifecycle: CREATE, LIST, INFO, DROP
- Ingestion: STORE text (+ optional metadata)
- Search: QUERY with K, optional FILTER and RETURN
- Delete by id
- Index creation (currently a placeholder/no-op)
References:
- Implementation: [src/lance_store.rs](src/lance_store.rs), [src/cmd.rs](src/cmd.rs), [src/rpc.rs](src/rpc.rs), [src/server.rs](src/server.rs), [src/embedding.rs](src/embedding.rs)
Notes:
- Admin DB 0 cannot be Lance (or Tantivy). Only databases with id >= 1 can use Lance.
- Permissions:
- Read operations (SEARCH, LIST, INFO) require read permission.
- Mutating operations (CREATE, STORE, CREATEINDEX, DEL, DROP, EMBEDDING CONFIG SET) require readwrite permission.
- Backend gating:
- If a DB is Lance, only LANCE.* and basic control commands (PING, ECHO, SELECT, INFO, CLIENT, etc.) are permitted.
- If a DB is not Lance, LANCE.* commands return an error.
Storage layout and schema:
- Files live at: <base_dir>/lance/<db_id>/<dataset>.lance
- Records schema:
- id: Utf8 (non-null)
- vector: FixedSizeList<Float32, dim> (non-null)
- text: Utf8 (nullable)
- meta: Utf8 JSON (nullable)
- Search is an L2 KNN brute-force scan for now (lower score = better). Index creation is a no-op placeholder to be implemented later.
Prerequisites:
- Start HeroDB with RPC enabled (for management calls):
- See [docs/basics.md](./basics.md) for flags. Example:
```bash
./target/release/herodb --dir /tmp/herodb --admin-secret mysecret --port 6379 --enable-rpc
```
## 0) Create a Lance-backed database (JSON-RPC)
Use the management API to create a database with backend "Lance". DB 0 is reserved for admin and cannot be Lance.
Request:
```json
{
"jsonrpc": "2.0",
"id": 1,
"method": "herodb_createDatabase",
"params": [
"Lance",
{ "name": "vectors-db", "storage_path": null, "max_size": null, "redis_version": null },
null
]
}
```
- Response contains the allocated db_id (>= 1). Use that id below (replace 1 with your actual id).
Select the database over RESP:
```bash
redis-cli -p 6379 SELECT 1
# → OK
```
## 1) Configure embedding provider (server-side embeddings)
HeroDB embeds text internally at STORE/SEARCH time using a per-dataset EmbeddingConfig sidecar. Configure provider before creating a dataset to choose dimensions and provider.
Supported providers:
- openai (standard OpenAI or Azure OpenAI)
- testhash (deterministic, CI-friendly; no network)
Environment variables for OpenAI:
- Standard OpenAI: export OPENAI_API_KEY=sk-...
- Azure OpenAI: export AZURE_OPENAI_API_KEY=...
RESP examples:
```bash
# Standard OpenAI with default dims (model-dependent, e.g. 1536)
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET myset PROVIDER openai MODEL text-embedding-3-small
# OpenAI with reduced output dimension (e.g., 512) when supported
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET myset PROVIDER openai MODEL text-embedding-3-small PARAM dim 512
# Azure OpenAI (set env: AZURE_OPENAI_API_KEY)
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET myset PROVIDER openai MODEL text-embedding-3-small \
PARAM use_azure true \
PARAM azure_endpoint https://myresource.openai.azure.com \
PARAM azure_deployment my-embed-deploy \
PARAM azure_api_version 2024-02-15 \
PARAM dim 512
# Deterministic test provider (no network, stable vectors)
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET myset PROVIDER testhash MODEL any
```
Read config:
```bash
redis-cli -p 6379 LANCE.EMBEDDING CONFIG GET myset
# → JSON blob describing provider/model/params
```
JSON-RPC examples:
```json
{
"jsonrpc": "2.0",
"id": 2,
"method": "herodb_lanceSetEmbeddingConfig",
"params": [
1,
"myset",
"openai",
"text-embedding-3-small",
{ "dim": "512" }
]
}
```
```json
{
"jsonrpc": "2.0",
"id": 3,
"method": "herodb_lanceGetEmbeddingConfig",
"params": [1, "myset"]
}
```
## 2) Create a dataset
Choose a dimension that matches your embedding configuration. For OpenAI text-embedding-3-small without dimension override, typical dimension is 1536; when `dim` is set (e.g., 512), use that. The current API requires an explicit DIM.
RESP:
```bash
redis-cli -p 6379 LANCE.CREATE myset DIM 512
# → OK
```
JSON-RPC:
```json
{
"jsonrpc": "2.0",
"id": 4,
"method": "herodb_lanceCreate",
"params": [1, "myset", 512]
}
```
## 3) Store text documents (server-side embedding)
Provide your id, the text to embed, and optional META fields. The server computes the embedding using the configured provider and stores id/vector/text/meta in the Lance dataset. Upserts by id are supported via delete-then-append semantics.
RESP:
```bash
redis-cli -p 6379 LANCE.STORE myset ID doc-1 TEXT "Hello vector world" META title "Hello" category "demo"
# → OK
```
JSON-RPC:
```json
{
"jsonrpc": "2.0",
"id": 5,
"method": "herodb_lanceStoreText",
"params": [
1,
"myset",
"doc-1",
"Hello vector world",
{ "title": "Hello", "category": "demo" }
]
}
```
## 4) Search with a text query
Provide a query string; the server embeds it and performs KNN search. Optional: FILTER expression and RETURN subset of fields.
RESP:
```bash
# K nearest neighbors for the query text
redis-cli -p 6379 LANCE.SEARCH myset K 5 QUERY "greetings to vectors"
# → Array of hits: [id, score, [k,v, ...]] pairs, lower score = closer
# With a filter on meta fields and return only title
redis-cli -p 6379 LANCE.SEARCH myset K 3 QUERY "greetings to vectors" FILTER "category = 'demo'" RETURN 1 title
```
JSON-RPC:
```json
{
"jsonrpc": "2.0",
"id": 6,
"method": "herodb_lanceSearchText",
"params": [1, "myset", "greetings to vectors", 5, null, null]
}
```
With filter and selected fields:
```json
{
"jsonrpc": "2.0",
"id": 7,
"method": "herodb_lanceSearchText",
"params": [1, "myset", "greetings to vectors", 3, "category = 'demo'", ["title"]]
}
```
Response shape:
- RESP over redis-cli: an array of hits [id, score, [k, v, ...]].
- JSON-RPC returns an object containing the RESP-encoded wire format string or a structured result depending on implementation. See [src/rpc.rs](src/rpc.rs) for details.
## 5) Create an index (placeholder)
Index creation currently returns OK but is a no-op. It will integrate Lance vector indices in a future update.
RESP:
```bash
redis-cli -p 6379 LANCE.CREATEINDEX myset TYPE "ivf_pq" PARAM nlist 100 PARAM pq_m 16
# → OK (no-op for now)
```
JSON-RPC:
```json
{
"jsonrpc": "2.0",
"id": 8,
"method": "herodb_lanceCreateIndex",
"params": [1, "myset", "ivf_pq", { "nlist": "100", "pq_m": "16" }]
}
```
## 6) Inspect datasets
RESP:
```bash
# List datasets in current Lance DB
redis-cli -p 6379 LANCE.LIST
# Get dataset info
redis-cli -p 6379 LANCE.INFO myset
```
JSON-RPC:
```json
{
"jsonrpc": "2.0",
"id": 9,
"method": "herodb_lanceList",
"params": [1]
}
```
```json
{
"jsonrpc": "2.0",
"id": 10,
"method": "herodb_lanceInfo",
"params": [1, "myset"]
}
```
## 7) Delete and drop
RESP:
```bash
# Delete by id
redis-cli -p 6379 LANCE.DEL myset doc-1
# → OK
# Drop the entire dataset
redis-cli -p 6379 LANCE.DROP myset
# → OK
```
JSON-RPC:
```json
{
"jsonrpc": "2.0",
"id": 11,
"method": "herodb_lanceDel",
"params": [1, "myset", "doc-1"]
}
```
```json
{
"jsonrpc": "2.0",
"id": 12,
"method": "herodb_lanceDrop",
"params": [1, "myset"]
}
```
## 8) End-to-end example (RESP)
```bash
# 1. Select Lance DB (assume db_id=1 created via RPC)
redis-cli -p 6379 SELECT 1
# 2. Configure embedding provider (OpenAI small model at 512 dims)
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET myset PROVIDER openai MODEL text-embedding-3-small PARAM dim 512
# 3. Create dataset
redis-cli -p 6379 LANCE.CREATE myset DIM 512
# 4. Store documents
redis-cli -p 6379 LANCE.STORE myset ID doc-1 TEXT "The quick brown fox jumps over the lazy dog" META title "Fox" category "animal"
redis-cli -p 6379 LANCE.STORE myset ID doc-2 TEXT "A fast auburn fox vaulted a sleepy canine" META title "Fox paraphrase" category "animal"
# 5. Search
redis-cli -p 6379 LANCE.SEARCH myset K 2 QUERY "quick brown fox" RETURN 1 title
# 6. Dataset info and listing
redis-cli -p 6379 LANCE.INFO myset
redis-cli -p 6379 LANCE.LIST
# 7. Delete and drop
redis-cli -p 6379 LANCE.DEL myset doc-2
redis-cli -p 6379 LANCE.DROP myset
```
## 9) End-to-end example (JSON-RPC)
Assume RPC server on port 8080. Replace ids and ports as needed.
1) Create Lance DB:
```json
{
"jsonrpc": "2.0",
"id": 100,
"method": "herodb_createDatabase",
"params": ["Lance", { "name": "vectors-db", "storage_path": null, "max_size": null, "redis_version": null }, null]
}
```
2) Set embedding config:
```json
{
"jsonrpc": "2.0",
"id": 101,
"method": "herodb_lanceSetEmbeddingConfig",
"params": [1, "myset", "openai", "text-embedding-3-small", { "dim": "512" }]
}
```
3) Create dataset:
```json
{
"jsonrpc": "2.0",
"id": 102,
"method": "herodb_lanceCreate",
"params": [1, "myset", 512]
}
```
4) Store text:
```json
{
"jsonrpc": "2.0",
"id": 103,
"method": "herodb_lanceStoreText",
"params": [1, "myset", "doc-1", "The quick brown fox jumps over the lazy dog", { "title": "Fox", "category": "animal" }]
}
```
5) Search text:
```json
{
"jsonrpc": "2.0",
"id": 104,
"method": "herodb_lanceSearchText",
"params": [1, "myset", "quick brown fox", 2, null, ["title"]]
}
```
6) Info/list:
```json
{
"jsonrpc": "2.0",
"id": 105,
"method": "herodb_lanceInfo",
"params": [1, "myset"]
}
```
```json
{
"jsonrpc": "2.0",
"id": 106,
"method": "herodb_lanceList",
"params": [1]
}
```
7) Delete/drop:
```json
{
"jsonrpc": "2.0",
"id": 107,
"method": "herodb_lanceDel",
"params": [1, "myset", "doc-1"]
}
```
```json
{
"jsonrpc": "2.0",
"id": 108,
"method": "herodb_lanceDrop",
"params": [1, "myset"]
}
```
## 10) Operational notes and troubleshooting
- If using OpenAI and you see “missing API key env”, set:
- Standard: `export OPENAI_API_KEY=sk-...`
- Azure: `export AZURE_OPENAI_API_KEY=...` and pass `use_azure true`, `azure_endpoint`, `azure_deployment`, `azure_api_version`.
- Dimensions mismatch:
- Ensure the dataset DIM equals the provider’s embedding dim. For OpenAI text-embedding-3 models, set `PARAM dim 512` (or another supported size) and use that same DIM for `LANCE.CREATE`.
- DB 0 restriction:
- Lance is not allowed on DB 0. Use db_id >= 1.
- Permissions:
- Read operations (SEARCH, LIST, INFO) require read permission.
- Mutations (CREATE, STORE, CREATEINDEX, DEL, DROP, EMBEDDING CONFIG SET) require readwrite permission.
- Backend gating:
- On Lance DBs, only LANCE.* commands are accepted (plus basic control).
- Current index behavior:
- `LANCE.CREATEINDEX` returns OK but is a no-op. Future versions will integrate Lance vector indices.
- Implementation files for reference:
- [src/lance_store.rs](src/lance_store.rs), [src/cmd.rs](src/cmd.rs), [src/rpc.rs](src/rpc.rs), [src/server.rs](src/server.rs), [src/embedding.rs](src/embedding.rs)

138
docs/lancedb_text_and_images_example.md Normal file
View File

@@ -0,0 +1,138 @@
# LanceDB Text and Images: End-to-End Example
This guide demonstrates creating a Lance backend database, ingesting two text documents and two images, performing searches over both, and cleaning up the datasets.
Prerequisites
- Build HeroDB and start the server with JSON-RPC enabled.
Commands:
```bash
cargo build --release
./target/release/herodb --dir /tmp/herodb --admin-secret mysecret --port 6379 --enable-rpc
```
We'll use:
- redis-cli for RESP commands against port 6379
- curl for JSON-RPC against 8080 if desired
- Deterministic local embedders to avoid external dependencies: testhash (text, dim 64) and testimagehash (image, dim 512)
0) Create a Lance-backed database (JSON-RPC)
Request:
```json
{ "jsonrpc": "2.0", "id": 1, "method": "herodb_createDatabase", "params": ["Lance", { "name": "media-db", "storage_path": null, "max_size": null, "redis_version": null }, null] }
```
Response returns db_id (assume 1). Select DB over RESP:
```bash
redis-cli -p 6379 SELECT 1
# → OK
```
1) Configure embedding providers
We'll create two datasets with independent embedding configs:
- textset → provider testhash, dim 64
- imageset → provider testimagehash, dim 512
Text config:
```bash
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET textset PROVIDER testhash MODEL any PARAM dim 64
# → OK
```
Image config:
```bash
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET imageset PROVIDER testimagehash MODEL any PARAM dim 512
# → OK
```
2) Create datasets
```bash
redis-cli -p 6379 LANCE.CREATE textset DIM 64
# → OK
redis-cli -p 6379 LANCE.CREATE imageset DIM 512
# → OK
```
3) Ingest two text documents (server-side embedding)
```bash
redis-cli -p 6379 LANCE.STORE textset ID doc-1 TEXT "The quick brown fox jumps over the lazy dog" META title "Fox" category "animal"
# → OK
redis-cli -p 6379 LANCE.STORE textset ID doc-2 TEXT "A fast auburn fox vaulted a sleepy canine" META title "Paraphrase" category "animal"
# → OK
```
4) Ingest two images
You can provide a URI or base64 bytes. Use URI for URIs, BYTES for base64 data.
Example using free placeholder images:
```bash
# Store via URI
redis-cli -p 6379 LANCE.STOREIMAGE imageset ID img-1 URI "https://picsum.photos/seed/1/256/256" META title "Seed1" group "demo"
# → OK
redis-cli -p 6379 LANCE.STOREIMAGE imageset ID img-2 URI "https://picsum.photos/seed/2/256/256" META title "Seed2" group "demo"
# → OK
```
If your environment blocks outbound HTTP, you can embed image bytes:
```bash
# Example: read a local file and base64 it (replace path)
b64=$(base64 -w0 ./image1.png)
redis-cli -p 6379 LANCE.STOREIMAGE imageset ID img-b64-1 BYTES "$b64" META title "Local1" group "demo"
```
5) Search text
```bash
# Top-2 nearest neighbors for a query
redis-cli -p 6379 LANCE.SEARCH textset K 2 QUERY "quick brown fox" RETURN 1 title
# → 1) [id, score, [k1,v1,...]]
```
With a filter (supports equality on schema or meta keys):
```bash
redis-cli -p 6379 LANCE.SEARCH textset K 2 QUERY "fox jumps" FILTER "category = 'animal'" RETURN 1 title
```
6) Search images
```bash
# Provide a URI as the query
redis-cli -p 6379 LANCE.SEARCHIMAGE imageset K 2 QUERYURI "https://picsum.photos/seed/1/256/256" RETURN 1 title
# Or provide base64 bytes as the query
qb64=$(curl -s https://picsum.photos/seed/3/256/256 | base64 -w0)
redis-cli -p 6379 LANCE.SEARCHIMAGE imageset K 2 QUERYBYTES "$qb64" RETURN 1 title
```
7) Inspect datasets
```bash
redis-cli -p 6379 LANCE.LIST
redis-cli -p 6379 LANCE.INFO textset
redis-cli -p 6379 LANCE.INFO imageset
```
8) Delete by id and drop datasets
```bash
# Delete one record
redis-cli -p 6379 LANCE.DEL textset doc-2
# → OK
# Drop entire datasets
redis-cli -p 6379 LANCE.DROP textset
redis-cli -p 6379 LANCE.DROP imageset
# → OK
```
Appendix: Using OpenAI embeddings instead of test providers
Text:
```bash
export OPENAI_API_KEY=sk-...
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET textset PROVIDER openai MODEL text-embedding-3-small PARAM dim 512
redis-cli -p 6379 LANCE.CREATE textset DIM 512
```
Azure OpenAI:
```bash
export AZURE_OPENAI_API_KEY=...
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET textset PROVIDER openai MODEL text-embedding-3-small \
PARAM use_azure true \
PARAM azure_endpoint https://myresource.openai.azure.com \
PARAM azure_deployment my-embed-deploy \
PARAM azure_api_version 2024-02-15 \
PARAM dim 512
```
Notes:
- Ensure dataset DIM matches the configured embedding dimension.
- Lance is only available for non-admin databases (db_id >= 1).
- On Lance DBs, only LANCE.* and basic control commands are allowed.

9
src/admin_meta.rs
View File

@@ -48,8 +48,8 @@ fn init_admin_storage(
let storage: Arc<dyn StorageBackend> = match backend {
options::BackendType::Redb => Arc::new(Storage::new(&db_file, true, Some(admin_secret))?),
options::BackendType::Sled => Arc::new(SledStorage::new(&db_file, true, Some(admin_secret))?),
options::BackendType::Tantivy => {
return Err(DBError("Admin DB 0 cannot use Tantivy backend".to_string()))
options::BackendType::Tantivy | options::BackendType::Lance => {
return Err(DBError("Admin DB 0 cannot use search-only backends (Tantivy/Lance)".to_string()))
}
};
Ok(storage)
@@ -206,6 +206,9 @@ pub fn open_data_storage(
options::BackendType::Tantivy => {
return Err(DBError("Tantivy backend has no KV storage; use FT.* commands only".to_string()))
}
options::BackendType::Lance => {
return Err(DBError("Lance backend has no KV storage; use LANCE.* commands only".to_string()))
}
};
// Publish to registry
@@ -299,6 +302,7 @@ pub fn set_database_backend(
options::BackendType::Redb => "Redb",
options::BackendType::Sled => "Sled",
options::BackendType::Tantivy => "Tantivy",
options::BackendType::Lance => "Lance",
};
let _ = admin.hset(&mk, vec![("backend".to_string(), val.to_string())])?;
Ok(())
@@ -316,6 +320,7 @@ pub fn get_database_backend(
Some(s) if s == "Redb" => Ok(Some(options::BackendType::Redb)),
Some(s) if s == "Sled" => Ok(Some(options::BackendType::Sled)),
Some(s) if s == "Tantivy" => Ok(Some(options::BackendType::Tantivy)),
Some(s) if s == "Lance" => Ok(Some(options::BackendType::Lance)),
_ => Ok(None),
}
}

727
src/cmd.rs
View File

@@ -1,4 +1,5 @@
use crate::{error::DBError, protocol::Protocol, server::Server};
use crate::{error::DBError, protocol::Protocol, server::Server, embedding::{EmbeddingConfig, EmbeddingProvider}};
use base64::{engine::general_purpose, Engine as _};
use tokio::time::{timeout, Duration};
use futures::future::select_all;
@@ -125,6 +126,67 @@ pub enum Cmd {
query: String,
group_by: Vec<String>,
reducers: Vec<String>,
},
// LanceDB text-first commands (no user-provided vectors)
LanceCreate {
name: String,
dim: usize,
},
LanceStoreText {
name: String,
id: String,
text: String,
meta: Vec<(String, String)>,
},
LanceSearchText {
name: String,
text: String,
k: usize,
filter: Option<String>,
return_fields: Option<Vec<String>>,
},
// Image-first commands (no user-provided vectors)
LanceStoreImage {
name: String,
id: String,
uri: Option<String>,
bytes_b64: Option<String>,
meta: Vec<(String, String)>,
},
LanceSearchImage {
name: String,
k: usize,
uri: Option<String>,
bytes_b64: Option<String>,
filter: Option<String>,
return_fields: Option<Vec<String>>,
},
LanceCreateIndex {
name: String,
index_type: String,
params: Vec<(String, String)>,
},
// Embedding configuration per dataset
LanceEmbeddingConfigSet {
name: String,
provider: String,
model: String,
params: Vec<(String, String)>,
},
LanceEmbeddingConfigGet {
name: String,
},
LanceList,
LanceInfo {
name: String,
},
LanceDel {
name: String,
id: String,
},
LanceDrop {
name: String,
}
}
@@ -815,6 +877,295 @@ impl Cmd {
let reducers = Vec::new();
Cmd::FtAggregate { index_name, query, group_by, reducers }
}
// ----- LANCE.* commands -----
"lance.create" => {
// LANCE.CREATE name DIM d
if cmd.len() != 4 || cmd[2].to_uppercase() != "DIM" {
return Err(DBError("ERR LANCE.CREATE requires: name DIM <dim>".to_string()));
}
let name = cmd[1].clone();
let dim: usize = cmd[3].parse().map_err(|_| DBError("ERR DIM must be an integer".to_string()))?;
Cmd::LanceCreate { name, dim }
}
"lance.store" => {
// LANCE.STORE name ID <id> TEXT <text> [META k v ...]
if cmd.len() < 6 {
return Err(DBError("ERR LANCE.STORE requires: name ID <id> TEXT <text> [META k v ...]".to_string()));
}
let name = cmd[1].clone();
let mut i = 2;
if cmd[i].to_uppercase() != "ID" || i + 1 >= cmd.len() {
return Err(DBError("ERR LANCE.STORE requires ID <id>".to_string()));
}
let id = cmd[i + 1].clone();
i += 2;
if i >= cmd.len() || cmd[i].to_uppercase() != "TEXT" {
return Err(DBError("ERR LANCE.STORE requires TEXT <text>".to_string()));
}
i += 1;
if i >= cmd.len() {
return Err(DBError("ERR LANCE.STORE requires TEXT <text>".to_string()));
}
let text = cmd[i].clone();
i += 1;
let mut meta: Vec<(String, String)> = Vec::new();
if i < cmd.len() && cmd[i].to_uppercase() == "META" {
i += 1;
while i + 1 < cmd.len() {
meta.push((cmd[i].clone(), cmd[i + 1].clone()));
i += 2;
}
}
Cmd::LanceStoreText { name, id, text, meta }
}
"lance.storeimage" => {
// LANCE.STOREIMAGE name ID <id> (URI <uri> | BYTES <base64>) [META k v ...]
if cmd.len() < 6 {
return Err(DBError("ERR LANCE.STOREIMAGE requires: name ID <id> (URI <uri> | BYTES <base64>) [META k v ...]".to_string()));
}
let name = cmd[1].clone();
let mut i = 2;
if cmd[i].to_uppercase() != "ID" || i + 1 >= cmd.len() {
return Err(DBError("ERR LANCE.STOREIMAGE requires ID <id>".to_string()));
}
let id = cmd[i + 1].clone();
i += 2;
let mut uri_opt: Option<String> = None;
let mut bytes_b64_opt: Option<String> = None;
if i < cmd.len() && cmd[i].to_uppercase() == "URI" {
if i + 1 >= cmd.len() { return Err(DBError("ERR LANCE.STOREIMAGE URI requires a value".to_string())); }
uri_opt = Some(cmd[i + 1].clone());
i += 2;
} else if i < cmd.len() && cmd[i].to_uppercase() == "BYTES" {
if i + 1 >= cmd.len() { return Err(DBError("ERR LANCE.STOREIMAGE BYTES requires a value".to_string())); }
bytes_b64_opt = Some(cmd[i + 1].clone());
i += 2;
} else {
return Err(DBError("ERR LANCE.STOREIMAGE requires either URI <uri> or BYTES <base64>".to_string()));
}
// Parse optional META pairs
let mut meta: Vec<(String, String)> = Vec::new();
if i < cmd.len() && cmd[i].to_uppercase() == "META" {
i += 1;
while i + 1 < cmd.len() {
meta.push((cmd[i].clone(), cmd[i + 1].clone()));
i += 2;
}
}
Cmd::LanceStoreImage { name, id, uri: uri_opt, bytes_b64: bytes_b64_opt, meta }
}
"lance.search" => {
// LANCE.SEARCH name K <k> QUERY <text> [FILTER expr] [RETURN n fields...]
if cmd.len() < 6 {
return Err(DBError("ERR LANCE.SEARCH requires: name K <k> QUERY <text> [FILTER expr] [RETURN n fields...]".to_string()));
}
let name = cmd[1].clone();
if cmd[2].to_uppercase() != "K" {
return Err(DBError("ERR LANCE.SEARCH requires K <k>".to_string()));
}
let k: usize = cmd[3].parse().map_err(|_| DBError("ERR K must be an integer".to_string()))?;
if cmd[4].to_uppercase() != "QUERY" {
return Err(DBError("ERR LANCE.SEARCH requires QUERY <text>".to_string()));
}
let mut i = 5;
if i >= cmd.len() {
return Err(DBError("ERR LANCE.SEARCH requires QUERY <text>".to_string()));
}
let text = cmd[i].clone();
i += 1;
let mut filter: Option<String> = None;
let mut return_fields: Option<Vec<String>> = None;
while i < cmd.len() {
match cmd[i].to_uppercase().as_str() {
"FILTER" => {
if i + 1 >= cmd.len() {
return Err(DBError("ERR FILTER requires an expression".to_string()));
}
filter = Some(cmd[i + 1].clone());
i += 2;
}
"RETURN" => {
if i + 1 >= cmd.len() {
return Err(DBError("ERR RETURN requires field count".to_string()));
}
let n: usize = cmd[i + 1].parse().map_err(|_| DBError("ERR RETURN count must be integer".to_string()))?;
i += 2;
let mut fields = Vec::new();
for _ in 0..n {
if i < cmd.len() {
fields.push(cmd[i].clone());
i += 1;
}
}
return_fields = Some(fields);
}
_ => { i += 1; }
}
}
Cmd::LanceSearchText { name, text, k, filter, return_fields }
}
"lance.searchimage" => {
// LANCE.SEARCHIMAGE name K <k> (QUERYURI <uri> | QUERYBYTES <base64>) [FILTER expr] [RETURN n fields...]
if cmd.len() < 6 {
return Err(DBError("ERR LANCE.SEARCHIMAGE requires: name K <k> (QUERYURI <uri> | QUERYBYTES <base64>) [FILTER expr] [RETURN n fields...]".to_string()));
}
let name = cmd[1].clone();
if cmd[2].to_uppercase() != "K" {
return Err(DBError("ERR LANCE.SEARCHIMAGE requires K <k>".to_string()));
}
let k: usize = cmd[3].parse().map_err(|_| DBError("ERR K must be an integer".to_string()))?;
let mut i = 4;
let mut uri_opt: Option<String> = None;
let mut bytes_b64_opt: Option<String> = None;
if i < cmd.len() && cmd[i].to_uppercase() == "QUERYURI" {
if i + 1 >= cmd.len() { return Err(DBError("ERR QUERYURI requires a value".to_string())); }
uri_opt = Some(cmd[i + 1].clone());
i += 2;
} else if i < cmd.len() && cmd[i].to_uppercase() == "QUERYBYTES" {
if i + 1 >= cmd.len() { return Err(DBError("ERR QUERYBYTES requires a value".to_string())); }
bytes_b64_opt = Some(cmd[i + 1].clone());
i += 2;
} else {
return Err(DBError("ERR LANCE.SEARCHIMAGE requires QUERYURI <uri> or QUERYBYTES <base64>".to_string()));
}
let mut filter: Option<String> = None;
let mut return_fields: Option<Vec<String>> = None;
while i < cmd.len() {
match cmd[i].to_uppercase().as_str() {
"FILTER" => {
if i + 1 >= cmd.len() {
return Err(DBError("ERR FILTER requires an expression".to_string()));
}
filter = Some(cmd[i + 1].clone());
i += 2;
}
"RETURN" => {
if i + 1 >= cmd.len() {
return Err(DBError("ERR RETURN requires field count".to_string()));
}
let n: usize = cmd[i + 1].parse().map_err(|_| DBError("ERR RETURN count must be integer".to_string()))?;
i += 2;
let mut fields = Vec::new();
for _ in 0..n {
if i < cmd.len() {
fields.push(cmd[i].clone());
i += 1;
}
}
return_fields = Some(fields);
}
_ => { i += 1; }
}
}
Cmd::LanceSearchImage { name, k, uri: uri_opt, bytes_b64: bytes_b64_opt, filter, return_fields }
}
"lance.createindex" => {
// LANCE.CREATEINDEX name TYPE t [PARAM k v ...]
if cmd.len() < 4 || cmd[2].to_uppercase() != "TYPE" {
return Err(DBError("ERR LANCE.CREATEINDEX requires: name TYPE <type> [PARAM k v ...]".to_string()));
}
let name = cmd[1].clone();
let index_type = cmd[3].clone();
let mut params: Vec<(String, String)> = Vec::new();
let mut i = 4;
if i < cmd.len() && cmd[i].to_uppercase() == "PARAM" {
i += 1;
while i + 1 < cmd.len() {
params.push((cmd[i].clone(), cmd[i + 1].clone()));
i += 2;
}
}
Cmd::LanceCreateIndex { name, index_type, params }
}
"lance.embedding" => {
// LANCE.EMBEDDING CONFIG SET name PROVIDER p MODEL m [PARAM k v ...]
// LANCE.EMBEDDING CONFIG GET name
if cmd.len() < 3 || cmd[1].to_uppercase() != "CONFIG" {
return Err(DBError("ERR LANCE.EMBEDDING requires CONFIG subcommand".to_string()));
}
if cmd.len() >= 4 && cmd[2].to_uppercase() == "SET" {
if cmd.len() < 8 {
return Err(DBError("ERR LANCE.EMBEDDING CONFIG SET requires: SET name PROVIDER p MODEL m [PARAM k v ...]".to_string()));
}
let name = cmd[3].clone();
let mut i = 4;
let mut provider: Option<String> = None;
let mut model: Option<String> = None;
let mut params: Vec<(String, String)> = Vec::new();
while i < cmd.len() {
match cmd[i].to_uppercase().as_str() {
"PROVIDER" => {
if i + 1 >= cmd.len() {
return Err(DBError("ERR PROVIDER requires a value".to_string()));
}
provider = Some(cmd[i + 1].clone());
i += 2;
}
"MODEL" => {
if i + 1 >= cmd.len() {
return Err(DBError("ERR MODEL requires a value".to_string()));
}
model = Some(cmd[i + 1].clone());
i += 2;
}
"PARAM" => {
i += 1;
while i + 1 < cmd.len() {
params.push((cmd[i].clone(), cmd[i + 1].clone()));
i += 2;
}
}
_ => {
// Unknown token; break to avoid infinite loop
i += 1;
}
}
}
let provider = provider.ok_or_else(|| DBError("ERR missing PROVIDER".to_string()))?;
let model = model.ok_or_else(|| DBError("ERR missing MODEL".to_string()))?;
Cmd::LanceEmbeddingConfigSet { name, provider, model, params }
} else if cmd.len() == 4 && cmd[2].to_uppercase() == "GET" {
let name = cmd[3].clone();
Cmd::LanceEmbeddingConfigGet { name }
} else {
return Err(DBError("ERR LANCE.EMBEDDING CONFIG supports: SET ... | GET name".to_string()));
}
}
"lance.list" => {
if cmd.len() != 1 {
return Err(DBError("ERR LANCE.LIST takes no arguments".to_string()));
}
Cmd::LanceList
}
"lance.info" => {
if cmd.len() != 2 {
return Err(DBError("ERR LANCE.INFO requires: name".to_string()));
}
Cmd::LanceInfo { name: cmd[1].clone() }
}
"lance.drop" => {
if cmd.len() != 2 {
return Err(DBError("ERR LANCE.DROP requires: name".to_string()));
}
Cmd::LanceDrop { name: cmd[1].clone() }
}
"lance.del" => {
if cmd.len() != 3 {
return Err(DBError("ERR LANCE.DEL requires: name id".to_string()));
}
Cmd::LanceDel { name: cmd[1].clone(), id: cmd[2].clone() }
}
_ => Cmd::Unknow(cmd[0].clone()),
},
protocol,
@@ -853,6 +1204,18 @@ impl Cmd {
.map(|b| matches!(b, crate::options::BackendType::Tantivy))
.unwrap_or(false);
// Determine Lance backend similarly
let is_lance_backend = crate::admin_meta::get_database_backend(
&server.option.dir,
server.option.backend.clone(),
&server.option.admin_secret,
server.selected_db,
)
.ok()
.flatten()
.map(|b| matches!(b, crate::options::BackendType::Lance))
.unwrap_or(false);
if is_tantivy_backend {
match &self {
Cmd::Select(..)
@@ -876,6 +1239,34 @@ impl Cmd {
}
}
// Lance backend gating: allow only LANCE.* and basic control/info commands
if is_lance_backend {
match &self {
Cmd::Select(..)
| Cmd::Quit
| Cmd::Client(..)
| Cmd::ClientSetName(..)
| Cmd::ClientGetName
| Cmd::Command(..)
| Cmd::Info(..)
| Cmd::LanceCreate { .. }
| Cmd::LanceStoreText { .. }
| Cmd::LanceSearchText { .. }
| Cmd::LanceStoreImage { .. }
| Cmd::LanceSearchImage { .. }
| Cmd::LanceEmbeddingConfigSet { .. }
| Cmd::LanceEmbeddingConfigGet { .. }
| Cmd::LanceCreateIndex { .. }
| Cmd::LanceList
| Cmd::LanceInfo { .. }
| Cmd::LanceDel { .. }
| Cmd::LanceDrop { .. } => {}
_ => {
return Ok(Protocol::err("ERR backend is Lance; only LANCE.* commands are allowed"));
}
}
}
// If selected DB is not Tantivy, forbid all FT.* commands here.
if !is_tantivy_backend {
match &self {
@@ -893,6 +1284,27 @@ impl Cmd {
}
}
// If selected DB is not Lance, forbid all LANCE.* commands here.
if !is_lance_backend {
match &self {
Cmd::LanceCreate { .. }
| Cmd::LanceStoreText { .. }
| Cmd::LanceSearchText { .. }
| Cmd::LanceStoreImage { .. }
| Cmd::LanceSearchImage { .. }
| Cmd::LanceEmbeddingConfigSet { .. }
| Cmd::LanceEmbeddingConfigGet { .. }
| Cmd::LanceCreateIndex { .. }
| Cmd::LanceList
| Cmd::LanceInfo { .. }
| Cmd::LanceDel { .. }
| Cmd::LanceDrop { .. } => {
return Ok(Protocol::err("ERR DB backend is not Lance; LANCE.* commands are not allowed"));
}
_ => {}
}
}
match self {
Cmd::Select(db, key) => select_cmd(server, db, key).await,
Cmd::Ping => Ok(Protocol::SimpleString("PONG".to_string())),
@@ -1015,6 +1427,307 @@ impl Cmd {
Ok(Protocol::err("FT.AGGREGATE not implemented yet"))
}
// LanceDB commands
Cmd::LanceCreate { name, dim } => {
if !server.has_write_permission() {
return Ok(Protocol::err("ERR write permission denied"));
}
match server.lance_store()?.create_dataset(&name, dim).await {
Ok(()) => Ok(Protocol::SimpleString("OK".to_string())),
Err(e) => Ok(Protocol::err(&e.0)),
}
}
Cmd::LanceEmbeddingConfigSet { name, provider, model, params } => {
if !server.has_write_permission() {
return Ok(Protocol::err("ERR write permission denied"));
}
// Map provider string to enum
let p_lc = provider.to_lowercase();
let prov = match p_lc.as_str() {
"test-hash" | "testhash" => EmbeddingProvider::TestHash,
"testimagehash" | "image-test-hash" | "imagetesthash" => EmbeddingProvider::ImageTestHash,
"fastembed" | "lancefastembed" => EmbeddingProvider::LanceFastEmbed,
"openai" | "lanceopenai" => EmbeddingProvider::LanceOpenAI,
other => EmbeddingProvider::LanceOther(other.to_string()),
};
let cfg = EmbeddingConfig {
provider: prov,
model,
params: params.into_iter().collect(),
};
match server.set_dataset_embedding_config(&name, &cfg) {
Ok(()) => Ok(Protocol::SimpleString("OK".to_string())),
Err(e) => Ok(Protocol::err(&e.0)),
}
}
Cmd::LanceEmbeddingConfigGet { name } => {
match server.get_dataset_embedding_config(&name) {
Ok(cfg) => {
let mut arr = Vec::new();
arr.push(Protocol::BulkString("provider".to_string()));
arr.push(Protocol::BulkString(match cfg.provider {
EmbeddingProvider::TestHash => "test-hash".to_string(),
EmbeddingProvider::ImageTestHash => "testimagehash".to_string(),
EmbeddingProvider::LanceFastEmbed => "lancefastembed".to_string(),
EmbeddingProvider::LanceOpenAI => "lanceopenai".to_string(),
EmbeddingProvider::LanceOther(ref s) => s.clone(),
}));
arr.push(Protocol::BulkString("model".to_string()));
arr.push(Protocol::BulkString(cfg.model.clone()));
arr.push(Protocol::BulkString("params".to_string()));
arr.push(Protocol::BulkString(serde_json::to_string(&cfg.params).unwrap_or_else(|_| "{}".to_string())));
Ok(Protocol::Array(arr))
}
Err(e) => Ok(Protocol::err(&e.0)),
}
}
Cmd::LanceStoreText { name, id, text, meta } => {
if !server.has_write_permission() {
return Ok(Protocol::err("ERR write permission denied"));
}
// Resolve embedder and embed text on a plain OS thread to avoid tokio runtime panics from reqwest::blocking
let embedder = server.get_embedder_for(&name)?;
let (tx, rx) = tokio::sync::oneshot::channel();
let emb_arc = embedder.clone();
let text_cl = text.clone();
std::thread::spawn(move || {
let res = emb_arc.embed(&text_cl);
let _ = tx.send(res);
});
let vector = match rx.await {
Ok(Ok(v)) => v,
Ok(Err(e)) => return Ok(Protocol::err(&e.0)),
Err(recv_err) => return Ok(Protocol::err(&format!("ERR embedding thread error: {}", recv_err))),
};
let meta_map: std::collections::HashMap<String, String> = meta.into_iter().collect();
match server.lance_store()?.store_vector(&name, &id, vector, meta_map, Some(text)).await {
Ok(()) => Ok(Protocol::SimpleString("OK".to_string())),
Err(e) => Ok(Protocol::err(&e.0)),
}
}
Cmd::LanceSearchText { name, text, k, filter, return_fields } => {
// Resolve embedder and embed query text on a plain OS thread
let embedder = server.get_embedder_for(&name)?;
let (tx, rx) = tokio::sync::oneshot::channel();
let emb_arc = embedder.clone();
let text_cl = text.clone();
std::thread::spawn(move || {
let res = emb_arc.embed(&text_cl);
let _ = tx.send(res);
});
let qv = match rx.await {
Ok(Ok(v)) => v,
Ok(Err(e)) => return Ok(Protocol::err(&e.0)),
Err(recv_err) => return Ok(Protocol::err(&format!("ERR embedding thread error: {}", recv_err))),
};
match server.lance_store()?.search_vectors(&name, qv, k, filter, return_fields).await {
Ok(results) => {
// Encode as array of [id, score, [k1, v1, k2, v2, ...]]
let mut arr = Vec::new();
for (id, score, meta) in results {
let mut meta_arr: Vec<Protocol> = Vec::new();
for (k, v) in meta {
meta_arr.push(Protocol::BulkString(k));
meta_arr.push(Protocol::BulkString(v));
}
arr.push(Protocol::Array(vec![
Protocol::BulkString(id),
Protocol::BulkString(score.to_string()),
Protocol::Array(meta_arr),
]));
}
Ok(Protocol::Array(arr))
}
Err(e) => Ok(Protocol::err(&e.0)),
}
}
// New: Image store
Cmd::LanceStoreImage { name, id, uri, bytes_b64, meta } => {
if !server.has_write_permission() {
return Ok(Protocol::err("ERR write permission denied"));
}
let use_uri = uri.is_some();
let use_b64 = bytes_b64.is_some();
if (use_uri && use_b64) || (!use_uri && !use_b64) {
return Ok(Protocol::err("ERR Provide exactly one of URI or BYTES for LANCE.STOREIMAGE"));
}
let max_bytes: usize = std::env::var("HERODB_IMAGE_MAX_BYTES")
.ok()
.and_then(|s| s.parse::<u64>().ok())
.unwrap_or(10 * 1024 * 1024) as usize;
let media_uri_opt = if let Some(u) = uri.clone() {
match server.fetch_image_bytes_from_uri(&u) {
Ok(_) => {}
Err(e) => return Ok(Protocol::err(&e.0)),
}
Some(u)
} else {
None
};
let bytes: Vec<u8> = if let Some(u) = uri {
match server.fetch_image_bytes_from_uri(&u) {
Ok(b) => b,
Err(e) => return Ok(Protocol::err(&e.0)),
}
} else {
let b64 = bytes_b64.unwrap_or_default();
let data = match general_purpose::STANDARD.decode(b64.as_bytes()) {
Ok(d) => d,
Err(e) => return Ok(Protocol::err(&format!("ERR base64 decode error: {}", e))),
};
if data.len() > max_bytes {
return Ok(Protocol::err(&format!("ERR image exceeds max allowed bytes {}", max_bytes)));
}
data
};
let img_embedder = match server.get_image_embedder_for(&name) {
Ok(e) => e,
Err(e) => return Ok(Protocol::err(&e.0)),
};
let (tx, rx) = tokio::sync::oneshot::channel();
let emb_arc = img_embedder.clone();
let bytes_cl = bytes.clone();
std::thread::spawn(move || {
let res = emb_arc.embed_image(&bytes_cl);
let _ = tx.send(res);
});
let vector = match rx.await {
Ok(Ok(v)) => v,
Ok(Err(e)) => return Ok(Protocol::err(&e.0)),
Err(recv_err) => return Ok(Protocol::err(&format!("ERR embedding thread error: {}", recv_err))),
};
let meta_map: std::collections::HashMap<String, String> = meta.into_iter().collect();
match server.lance_store()?.store_vector_with_media(
&name,
&id,
vector,
meta_map,
None,
Some("image".to_string()),
media_uri_opt,
).await {
Ok(()) => Ok(Protocol::SimpleString("OK".to_string())),
Err(e) => Ok(Protocol::err(&e.0)),
}
}
// New: Image search
Cmd::LanceSearchImage { name, k, uri, bytes_b64, filter, return_fields } => {
let use_uri = uri.is_some();
let use_b64 = bytes_b64.is_some();
if (use_uri && use_b64) || (!use_uri && !use_b64) {
return Ok(Protocol::err("ERR Provide exactly one of QUERYURI or QUERYBYTES for LANCE.SEARCHIMAGE"));
}
let max_bytes: usize = std::env::var("HERODB_IMAGE_MAX_BYTES")
.ok()
.and_then(|s| s.parse::<u64>().ok())
.unwrap_or(10 * 1024 * 1024) as usize;
let bytes: Vec<u8> = if let Some(u) = uri {
match server.fetch_image_bytes_from_uri(&u) {
Ok(b) => b,
Err(e) => return Ok(Protocol::err(&e.0)),
}
} else {
let b64 = bytes_b64.unwrap_or_default();
let data = match general_purpose::STANDARD.decode(b64.as_bytes()) {
Ok(d) => d,
Err(e) => return Ok(Protocol::err(&format!("ERR base64 decode error: {}", e))),
};
if data.len() > max_bytes {
return Ok(Protocol::err(&format!("ERR image exceeds max allowed bytes {}", max_bytes)));
}
data
};
let img_embedder = match server.get_image_embedder_for(&name) {
Ok(e) => e,
Err(e) => return Ok(Protocol::err(&e.0)),
};
let (tx, rx) = tokio::sync::oneshot::channel();
std::thread::spawn(move || {
let res = img_embedder.embed_image(&bytes);
let _ = tx.send(res);
});
let qv = match rx.await {
Ok(Ok(v)) => v,
Ok(Err(e)) => return Ok(Protocol::err(&e.0)),
Err(recv_err) => return Ok(Protocol::err(&format!("ERR embedding thread error: {}", recv_err))),
};
match server.lance_store()?.search_vectors(&name, qv, k, filter, return_fields).await {
Ok(results) => {
let mut arr = Vec::new();
for (id, score, meta) in results {
let mut meta_arr: Vec<Protocol> = Vec::new();
for (k, v) in meta {
meta_arr.push(Protocol::BulkString(k));
meta_arr.push(Protocol::BulkString(v));
}
arr.push(Protocol::Array(vec![
Protocol::BulkString(id),
Protocol::BulkString(score.to_string()),
Protocol::Array(meta_arr),
]));
}
Ok(Protocol::Array(arr))
}
Err(e) => Ok(Protocol::err(&e.0)),
}
}
Cmd::LanceCreateIndex { name, index_type, params } => {
if !server.has_write_permission() {
return Ok(Protocol::err("ERR write permission denied"));
}
let params_map: std::collections::HashMap<String, String> = params.into_iter().collect();
match server.lance_store()?.create_index(&name, &index_type, params_map).await {
Ok(()) => Ok(Protocol::SimpleString("OK".to_string())),
Err(e) => Ok(Protocol::err(&e.0)),
}
}
Cmd::LanceList => {
match server.lance_store()?.list_datasets().await {
Ok(list) => Ok(Protocol::Array(list.into_iter().map(Protocol::BulkString).collect())),
Err(e) => Ok(Protocol::err(&e.0)),
}
}
Cmd::LanceInfo { name } => {
match server.lance_store()?.get_dataset_info(&name).await {
Ok(info) => {
let mut arr = Vec::new();
for (k, v) in info {
arr.push(Protocol::BulkString(k));
arr.push(Protocol::BulkString(v));
}
Ok(Protocol::Array(arr))
}
Err(e) => Ok(Protocol::err(&e.0)),
}
}
Cmd::LanceDel { name, id } => {
if !server.has_write_permission() {
return Ok(Protocol::err("ERR write permission denied"));
}
match server.lance_store()?.delete_by_id(&name, &id).await {
Ok(b) => Ok(Protocol::SimpleString(if b { "1" } else { "0" }.to_string())),
Err(e) => Ok(Protocol::err(&e.0)),
}
}
Cmd::LanceDrop { name } => {
if !server.has_write_permission() {
return Ok(Protocol::err("ERR write permission denied"));
}
match server.lance_store()?.drop_dataset(&name).await {
Ok(_b) => Ok(Protocol::SimpleString("OK".to_string())),
Err(e) => Ok(Protocol::err(&e.0)),
}
}
Cmd::Unknow(s) => Ok(Protocol::err(&format!("ERR unknown command `{}`", s))),
}
}
@@ -1114,8 +1827,8 @@ async fn select_cmd(server: &mut Server, db: u64, key: Option<String>) -> Result
.ok()
.flatten();
if matches!(eff_backend, Some(crate::options::BackendType::Tantivy)) {
// Tantivy DBs have no KV storage; allow SELECT to succeed
if matches!(eff_backend, Some(crate::options::BackendType::Tantivy) | Some(crate::options::BackendType::Lance)) {
// Search-only DBs (Tantivy/Lance) have no KV storage; allow SELECT to succeed
Ok(Protocol::SimpleString("OK".to_string()))
} else {
match server.current_storage() {
@@ -1459,9 +2172,9 @@ async fn dbsize_cmd(server: &Server) -> Result<Protocol, DBError> {
}
async fn info_cmd(server: &Server, section: &Option<String>) -> Result<Protocol, DBError> {
// For Tantivy backend, there is no KV storage; synthesize minimal info.
// For Tantivy or Lance backend, there is no KV storage; synthesize minimal info.
// Determine effective backend for the currently selected db.
let is_tantivy_db = crate::admin_meta::get_database_backend(
let is_search_only_db = crate::admin_meta::get_database_backend(
&server.option.dir,
server.option.backend.clone(),
&server.option.admin_secret,
@@ -1469,10 +2182,10 @@ async fn info_cmd(server: &Server, section: &Option<String>) -> Result<Protocol,
)
.ok()
.flatten()
.map(|b| matches!(b, crate::options::BackendType::Tantivy))
.map(|b| matches!(b, crate::options::BackendType::Tantivy | crate::options::BackendType::Lance))
.unwrap_or(false);
let storage_info: Vec<(String, String)> = if is_tantivy_db {
let storage_info: Vec<(String, String)> = if is_search_only_db {
vec![
("db_size".to_string(), "0".to_string()),
("is_encrypted".to_string(), "false".to_string()),

6
src/crypto.rs
View File

@@ -1,8 +1,8 @@
use chacha20poly1305::{
aead::{Aead, KeyInit, OsRng},
aead::{Aead, KeyInit},
XChaCha20Poly1305, XNonce,
};
use rand::RngCore;
use rand::{rngs::OsRng, RngCore};
use sha2::{Digest, Sha256};
const VERSION: u8 = 1;
@@ -31,7 +31,7 @@ pub struct CryptoFactory {
impl CryptoFactory {
/// Accepts any secret bytes; turns them into a 32-byte key (SHA-256).
pub fn new<S: AsRef<[u8]>>(secret: S) -> Self {
let mut h = Sha256::new();
let mut h = Sha256::default();
h.update(b"xchacha20poly1305-factory:v1"); // domain separation
h.update(secret.as_ref());
let digest = h.finalize(); // 32 bytes

405
src/embedding.rs Normal file
View File

@@ -0,0 +1,405 @@
// Embedding abstraction and minimal providers.
use std::collections::HashMap;
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use crate::error::DBError;
// Networking for OpenAI/Azure
use std::time::Duration;
use ureq::{Agent, AgentBuilder};
use serde_json::json;
/// Provider identifiers. Extend as needed to mirror LanceDB-supported providers.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum EmbeddingProvider {
// Deterministic, local-only embedder for CI and offline development (text).
TestHash,
// Deterministic, local-only embedder for CI and offline development (image).
ImageTestHash,
// Placeholders for LanceDB-supported providers; implementers can add concrete backends later.
LanceFastEmbed,
LanceOpenAI,
LanceOther(String),
}
/// Serializable embedding configuration.
/// params: arbitrary key-value map for provider-specific knobs (e.g., "dim", "api_key_env", etc.)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingConfig {
pub provider: EmbeddingProvider,
pub model: String,
#[serde(default)]
pub params: HashMap<String, String>,
}
impl EmbeddingConfig {
pub fn get_param_usize(&self, key: &str) -> Option<usize> {
self.params.get(key).and_then(|v| v.parse::<usize>().ok())
}
pub fn get_param_string(&self, key: &str) -> Option<String> {
self.params.get(key).cloned()
}
}
/// A provider-agnostic text embedding interface.
pub trait Embedder: Send + Sync {
/// Human-readable provider/model name
fn name(&self) -> String;
/// Embedding dimension
fn dim(&self) -> usize;
/// Embed a single text string into a fixed-length vector
fn embed(&self, text: &str) -> Result<Vec<f32>, DBError>;
/// Embed many texts; default maps embed() over inputs
fn embed_many(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, DBError> {
texts.iter().map(|t| self.embed(t)).collect()
}
}
//// ----------------------------- TEXT: deterministic test embedder -----------------------------
/// Deterministic, no-deps, no-network embedder for CI and offline dev.
/// Algorithm:
/// - Fold bytes of UTF-8 into 'dim' buckets with a simple rolling hash
/// - Apply tanh-like scaling and L2-normalize to unit length
pub struct TestHashEmbedder {
dim: usize,
model_name: String,
}
impl TestHashEmbedder {
pub fn new(dim: usize, model_name: impl Into<String>) -> Self {
Self { dim, model_name: model_name.into() }
}
fn l2_normalize(mut v: Vec<f32>) -> Vec<f32> {
let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm > 0.0 {
for x in &mut v {
*x /= norm;
}
}
v
}
}
impl Embedder for TestHashEmbedder {
fn name(&self) -> String {
format!("test-hash:{}", self.model_name)
}
fn dim(&self) -> usize {
self.dim
}
fn embed(&self, text: &str) -> Result<Vec<f32>, DBError> {
let mut acc = vec![0f32; self.dim];
// A simple, deterministic folding hash over bytes
let mut h1: u32 = 2166136261u32; // FNV-like seed
let mut h2: u32 = 0x9e3779b9u32; // golden ratio
for (i, b) in text.as_bytes().iter().enumerate() {
h1 ^= *b as u32;
h1 = h1.wrapping_mul(16777619u32);
h2 = h2.wrapping_add(((*b as u32) << (i % 13)) ^ (h1.rotate_left((i % 7) as u32)));
let idx = (h1 ^ h2) as usize % self.dim;
// Map byte to [-1, 1] and accumulate with mild decay by position
let val = ((*b as f32) / 127.5 - 1.0) * (1.0 / (1.0 + (i as f32 / 32.0)));
acc[idx] += val;
}
// Non-linear squashing to stabilize + normalize
for x in &mut acc {
*x = x.tanh();
}
Ok(Self::l2_normalize(acc))
}
}
//// ----------------------------- IMAGE: trait + deterministic test embedder -----------------------------
/// Image embedding interface (separate from text to keep modality-specific inputs).
pub trait ImageEmbedder: Send + Sync {
/// Human-readable provider/model name
fn name(&self) -> String;
/// Embedding dimension
fn dim(&self) -> usize;
/// Embed a single image (raw bytes)
fn embed_image(&self, bytes: &[u8]) -> Result<Vec<f32>, DBError>;
/// Embed many images; default maps embed_image() over inputs
fn embed_many_images(&self, images: &[Vec<u8>]) -> Result<Vec<Vec<f32>>, DBError> {
images.iter().map(|b| self.embed_image(b)).collect()
}
}
/// Deterministic image embedder that folds bytes into buckets, applies tanh-like nonlinearity,
/// and L2-normalizes. Suitable for CI and offline development.
/// NOTE: This is NOT semantic; it is a stable hash-like representation.
pub struct TestImageHashEmbedder {
dim: usize,
model_name: String,
}
impl TestImageHashEmbedder {
pub fn new(dim: usize, model_name: impl Into<String>) -> Self {
Self { dim, model_name: model_name.into() }
}
fn l2_normalize(mut v: Vec<f32>) -> Vec<f32> {
let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm > 0.0 {
for x in &mut v {
*x /= norm;
}
}
v
}
}
impl ImageEmbedder for TestImageHashEmbedder {
fn name(&self) -> String {
format!("test-image-hash:{}", self.model_name)
}
fn dim(&self) -> usize {
self.dim
}
fn embed_image(&self, bytes: &[u8]) -> Result<Vec<f32>, DBError> {
// Deterministic fold across bytes with two rolling accumulators.
let mut acc = vec![0f32; self.dim];
let mut h1: u32 = 0x811C9DC5; // FNV-like
let mut h2: u32 = 0x9E3779B9; // golden ratio
for (i, b) in bytes.iter().enumerate() {
h1 ^= *b as u32;
h1 = h1.wrapping_mul(16777619u32);
// combine with position and h2
h2 = h2.wrapping_add(((i as u32).rotate_left((i % 13) as u32)) ^ h1.rotate_left((i % 7) as u32));
let idx = (h1 ^ h2) as usize % self.dim;
// Map to [-1,1] and decay with position
let val = ((*b as f32) / 127.5 - 1.0) * (1.0 / (1.0 + (i as f32 / 128.0)));
acc[idx] += val;
}
for x in &mut acc {
*x = x.tanh();
}
Ok(Self::l2_normalize(acc))
}
}
//// OpenAI embedder (supports OpenAI and Azure OpenAI via REST)
struct OpenAIEmbedder {
model: String,
dim: usize,
agent: Agent,
endpoint: String,
headers: Vec<(String, String)>,
use_azure: bool,
}
impl OpenAIEmbedder {
fn new_from_config(cfg: &EmbeddingConfig) -> Result<Self, DBError> {
// Whether to use Azure OpenAI
let use_azure = cfg
.get_param_string("use_azure")
.map(|s| s.eq_ignore_ascii_case("true"))
.unwrap_or(false);
// Resolve API key (OPENAI_API_KEY or AZURE_OPENAI_API_KEY by default)
let api_key_env = cfg
.get_param_string("api_key_env")
.unwrap_or_else(|| {
if use_azure {
"AZURE_OPENAI_API_KEY".to_string()
} else {
"OPENAI_API_KEY".to_string()
}
});
let api_key = std::env::var(&api_key_env)
.map_err(|_| DBError(format!("Missing API key in env '{}'", api_key_env)))?;
// Resolve endpoint
// - Standard OpenAI: https://api.openai.com/v1/embeddings (default) or params["base_url"]
// - Azure OpenAI: {azure_endpoint}/openai/deployments/{deployment}/embeddings?api-version=...
let endpoint = if use_azure {
let base = cfg
.get_param_string("azure_endpoint")
.ok_or_else(|| DBError("Missing 'azure_endpoint' for Azure OpenAI".into()))?;
let deployment = cfg
.get_param_string("azure_deployment")
.unwrap_or_else(|| cfg.model.clone());
let api_version = cfg
.get_param_string("azure_api_version")
.unwrap_or_else(|| "2023-05-15".to_string());
format!(
"{}/openai/deployments/{}/embeddings?api-version={}",
base.trim_end_matches('/'),
deployment,
api_version
)
} else {
cfg.get_param_string("base_url")
.unwrap_or_else(|| "https://api.openai.com/v1/embeddings".to_string())
};
// Determine expected dimension (default 1536 for text-embedding-3-small; callers should override if needed)
let dim = cfg
.get_param_usize("dim")
.or_else(|| cfg.get_param_usize("dimensions"))
.unwrap_or(1536);
// Build an HTTP agent with timeouts (blocking; no tokio runtime involved)
let agent = AgentBuilder::new()
.timeout_read(Duration::from_secs(30))
.timeout_write(Duration::from_secs(30))
.build();
// Headers
let mut headers: Vec<(String, String)> = Vec::new();
headers.push(("Content-Type".to_string(), "application/json".to_string()));
if use_azure {
headers.push(("api-key".to_string(), api_key));
} else {
headers.push(("Authorization".to_string(), format!("Bearer {}", api_key)));
}
Ok(Self {
model: cfg.model.clone(),
dim,
agent,
endpoint,
headers,
use_azure,
})
}
fn request_many(&self, inputs: &[String]) -> Result<Vec<Vec<f32>>, DBError> {
// Compose request body:
// - Standard OpenAI: { "model": ..., "input": [...], "dimensions": dim? }
// - Azure: { "input": [...], "dimensions": dim? } (model from deployment)
let mut body = if self.use_azure {
json!({ "input": inputs })
} else {
json!({ "model": self.model, "input": inputs })
};
if self.dim > 0 {
body.as_object_mut()
.unwrap()
.insert("dimensions".to_string(), json!(self.dim));
}
// Build request
let mut req = self.agent.post(&self.endpoint);
for (k, v) in &self.headers {
req = req.set(k, v);
}
// Send and handle errors
let resp = req.send_json(body);
let text = match resp {
Ok(r) => r
.into_string()
.map_err(|e| DBError(format!("Failed to read embeddings response: {}", e)))?,
Err(ureq::Error::Status(code, r)) => {
let body = r.into_string().unwrap_or_default();
return Err(DBError(format!("Embeddings API error {}: {}", code, body)));
}
Err(e) => return Err(DBError(format!("HTTP request failed: {}", e))),
};
let val: serde_json::Value = serde_json::from_str(&text)
.map_err(|e| DBError(format!("Invalid JSON from embeddings API: {}", e)))?;
let data = val
.get("data")
.and_then(|d| d.as_array())
.ok_or_else(|| DBError("Embeddings API response missing 'data' array".into()))?;
let mut out: Vec<Vec<f32>> = Vec::with_capacity(data.len());
for item in data {
let emb = item
.get("embedding")
.and_then(|e| e.as_array())
.ok_or_else(|| DBError("Embeddings API item missing 'embedding'".into()))?;
let mut v: Vec<f32> = Vec::with_capacity(emb.len());
for n in emb {
let f = n
.as_f64()
.ok_or_else(|| DBError("Embedding element is not a number".into()))?;
v.push(f as f32);
}
if self.dim > 0 && v.len() != self.dim {
return Err(DBError(format!(
"Embedding dimension mismatch: expected {}, got {}. Configure 'dim' or 'dimensions' to match output.",
self.dim, v.len()
)));
}
out.push(v);
}
Ok(out)
}
}
impl Embedder for OpenAIEmbedder {
fn name(&self) -> String {
if self.use_azure {
format!("azure-openai:{}", self.model)
} else {
format!("openai:{}", self.model)
}
}
fn dim(&self) -> usize {
self.dim
}
fn embed(&self, text: &str) -> Result<Vec<f32>, DBError> {
let v = self.request_many(&[text.to_string()])?;
Ok(v.into_iter().next().unwrap_or_else(|| vec![0.0; self.dim]))
}
fn embed_many(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, DBError> {
if texts.is_empty() {
return Ok(vec![]);
}
self.request_many(texts)
}
}
/// Create an embedder instance from a config.
/// - TestHash: uses params["dim"] or defaults to 64
/// - LanceOpenAI: uses OpenAI (or Azure OpenAI) embeddings REST API
/// - Other Lance providers can be added similarly
pub fn create_embedder(config: &EmbeddingConfig) -> Result<Arc<dyn Embedder>, DBError> {
match &config.provider {
EmbeddingProvider::TestHash => {
let dim = config.get_param_usize("dim").unwrap_or(64);
Ok(Arc::new(TestHashEmbedder::new(dim, config.model.clone())))
}
EmbeddingProvider::LanceOpenAI => {
let inner = OpenAIEmbedder::new_from_config(config)?;
Ok(Arc::new(inner))
}
EmbeddingProvider::ImageTestHash => {
Err(DBError("Use create_image_embedder() for image providers".into()))
}
EmbeddingProvider::LanceFastEmbed => Err(DBError("LanceFastEmbed provider not yet implemented in Rust embedding layer; configure 'test-hash' or use 'openai'".into())),
EmbeddingProvider::LanceOther(p) => Err(DBError(format!("Lance provider '{}' not implemented; configure 'openai' or 'test-hash'", p))),
}
}
/// Create an image embedder instance from a config.
pub fn create_image_embedder(config: &EmbeddingConfig) -> Result<Arc<dyn ImageEmbedder>, DBError> {
match &config.provider {
EmbeddingProvider::ImageTestHash => {
let dim = config.get_param_usize("dim").unwrap_or(512);
Ok(Arc::new(TestImageHashEmbedder::new(dim, config.model.clone())))
}
EmbeddingProvider::TestHash | EmbeddingProvider::LanceOpenAI => {
Err(DBError("Configured text provider; dataset expects image provider (e.g., 'testimagehash')".into()))
}
EmbeddingProvider::LanceFastEmbed => Err(DBError("Image provider 'lancefastembed' not yet implemented".into())),
EmbeddingProvider::LanceOther(p) => Err(DBError(format!("Image provider '{}' not implemented; use 'testimagehash' for now", p))),
}
}

663
src/lance_store.rs Normal file
View File

@@ -0,0 +1,663 @@
// LanceDB store abstraction (per database instance)
// This module encapsulates all Lance/LanceDB operations for a given DB id.
// Notes:
// - We persist each dataset (aka "table") under <base_dir>/lance/<db_id>/<name>.lance
// - Schema convention: id: Utf8 (non-null), vector: FixedSizeList<Float32, dim> (non-null), meta: Utf8 (nullable JSON string)
// - We implement naive KNN (L2) scan in Rust for search to avoid tight coupling to lancedb search builder API.
// Index creation uses lance::Dataset vector index; future optimization can route to index-aware search.
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use crate::error::DBError;
use arrow_array::{Array, RecordBatch, RecordBatchIterator, StringArray};
use arrow_array::builder::{FixedSizeListBuilder, Float32Builder, StringBuilder};
use arrow_array::cast::AsArray;
use arrow_schema::{DataType, Field, Schema};
use futures::StreamExt;
use serde_json::Value as JsonValue;
// Low-level Lance core
use lance::dataset::{WriteMode, WriteParams};
use lance::Dataset;
// Vector index (IVF_PQ etc.)
// High-level LanceDB (for deletes where available)
use lancedb::connection::Connection;
use arrow_array::types::Float32Type;
#[derive(Clone)]
pub struct LanceStore {
base_dir: PathBuf,
db_id: u64,
}
impl LanceStore {
// Create a new LanceStore rooted at <base_dir>/lance/<db_id>
pub fn new(base_dir: &Path, db_id: u64) -> Result<Self, DBError> {
let p = base_dir.join("lance").join(db_id.to_string());
std::fs::create_dir_all(&p)
.map_err(|e| DBError(format!("Failed to create Lance dir {}: {}", p.display(), e)))?;
Ok(Self { base_dir: p, db_id })
}
fn dataset_path(&self, name: &str) -> PathBuf {
// Store datasets as directories or files with .lance suffix
// We accept both "<name>" and "<name>.lance" as logical name; normalize on ".lance"
let has_ext = name.ends_with(".lance");
if has_ext {
self.base_dir.join(name)
} else {
self.base_dir.join(format!("{name}.lance"))
}
}
fn file_uri(path: &Path) -> String {
// lancedb can use filesystem path directly; keep it simple
// Avoid file:// scheme since local paths are supported.
path.to_string_lossy().to_string()
}
async fn connect_db(&self) -> Result<Connection, DBError> {
let uri = Self::file_uri(&self.base_dir);
lancedb::connect(&uri)
.execute()
.await
.map_err(|e| DBError(format!("LanceDB connect failed at {}: {}", uri, e)))
}
fn vector_field(dim: i32) -> Field {
Field::new(
"vector",
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim),
false,
)
}
async fn read_existing_dim(&self, name: &str) -> Result<Option<i32>, DBError> {
let path = self.dataset_path(name);
if !path.exists() {
return Ok(None);
}
let ds = Dataset::open(path.to_string_lossy().as_ref())
.await
.map_err(|e| DBError(format!("Open dataset failed: {}: {}", path.display(), e)))?;
// Scan a single batch to infer vector dimension from the 'vector' column type
let mut scan = ds.scan();
if let Err(e) = scan.project(&["vector"]) {
return Err(DBError(format!("Project failed while inferring dim: {}", e)));
}
let mut stream = scan
.try_into_stream()
.await
.map_err(|e| DBError(format!("Scan stream failed while inferring dim: {}", e)))?;
if let Some(batch_res) = stream.next().await {
let batch = batch_res.map_err(|e| DBError(format!("Batch error: {}", e)))?;
let vec_col = batch
.column_by_name("vector")
.ok_or_else(|| DBError("Column 'vector' missing".into()))?;
let fsl = vec_col.as_fixed_size_list();
let dim = fsl.value_length();
return Ok(Some(dim));
}
Ok(None)
}
fn build_schema(dim: i32) -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("id", DataType::Utf8, false),
Self::vector_field(dim),
Field::new("text", DataType::Utf8, true),
Field::new("media_type", DataType::Utf8, true),
Field::new("media_uri", DataType::Utf8, true),
Field::new("meta", DataType::Utf8, true),
]))
}
fn build_one_row_batch(
id: &str,
vector: &[f32],
meta: &HashMap<String, String>,
text: Option<&str>,
media_type: Option<&str>,
media_uri: Option<&str>,
dim: i32,
) -> Result<(Arc<Schema>, RecordBatch), DBError> {
if vector.len() as i32 != dim {
return Err(DBError(format!(
"Vector length mismatch: expected {}, got {}",
dim,
vector.len()
)));
}
let schema = Self::build_schema(dim);
// id column
let mut id_builder = StringBuilder::new();
id_builder.append_value(id);
let id_arr = Arc::new(id_builder.finish()) as Arc<dyn Array>;
// vector column (FixedSizeList<Float32, dim>)
let v_builder = Float32Builder::with_capacity(vector.len());
let mut list_builder = FixedSizeListBuilder::new(v_builder, dim);
for v in vector {
list_builder.values().append_value(*v);
}
list_builder.append(true);
let vec_arr = Arc::new(list_builder.finish()) as Arc<dyn Array>;
// text column (optional)
let mut text_builder = StringBuilder::new();
if let Some(t) = text {
text_builder.append_value(t);
} else {
text_builder.append_null();
}
let text_arr = Arc::new(text_builder.finish()) as Arc<dyn Array>;
// media_type column (optional)
let mut mt_builder = StringBuilder::new();
if let Some(mt) = media_type {
mt_builder.append_value(mt);
} else {
mt_builder.append_null();
}
let mt_arr = Arc::new(mt_builder.finish()) as Arc<dyn Array>;
// media_uri column (optional)
let mut mu_builder = StringBuilder::new();
if let Some(mu) = media_uri {
mu_builder.append_value(mu);
} else {
mu_builder.append_null();
}
let mu_arr = Arc::new(mu_builder.finish()) as Arc<dyn Array>;
// meta column (JSON string)
let meta_json = if meta.is_empty() {
None
} else {
Some(serde_json::to_string(meta).map_err(|e| DBError(format!("Serialize meta error: {e}")))?)
};
let mut meta_builder = StringBuilder::new();
if let Some(s) = meta_json {
meta_builder.append_value(&s);
} else {
meta_builder.append_null();
}
let meta_arr = Arc::new(meta_builder.finish()) as Arc<dyn Array>;
let batch =
RecordBatch::try_new(schema.clone(), vec![id_arr, vec_arr, text_arr, mt_arr, mu_arr, meta_arr]).map_err(|e| {
DBError(format!("RecordBatch build failed: {e}"))
})?;
Ok((schema, batch))
}
// Create a new dataset (vector collection) with dimension `dim`.
pub async fn create_dataset(&self, name: &str, dim: usize) -> Result<(), DBError> {
let dim_i32: i32 = dim
.try_into()
.map_err(|_| DBError("Dimension too large".into()))?;
let path = self.dataset_path(name);
if path.exists() {
// Validate dimension if present
if let Some(existing_dim) = self.read_existing_dim(name).await? {
if existing_dim != dim_i32 {
return Err(DBError(format!(
"Dataset '{}' already exists with dim {}, requested {}",
name, existing_dim, dim_i32
)));
}
// No-op
return Ok(());
}
}
// Create an empty dataset by writing an empty batch
let schema = Self::build_schema(dim_i32);
let empty_id = Arc::new(StringArray::new_null(0));
// Build an empty FixedSizeListArray
let v_builder = Float32Builder::new();
let mut list_builder = FixedSizeListBuilder::new(v_builder, dim_i32);
let empty_vec = Arc::new(list_builder.finish()) as Arc<dyn Array>;
let empty_text = Arc::new(StringArray::new_null(0));
let empty_media_type = Arc::new(StringArray::new_null(0));
let empty_media_uri = Arc::new(StringArray::new_null(0));
let empty_meta = Arc::new(StringArray::new_null(0));
let empty_batch =
RecordBatch::try_new(schema.clone(), vec![empty_id, empty_vec, empty_text, empty_media_type, empty_media_uri, empty_meta])
.map_err(|e| DBError(format!("Build empty batch failed: {e}")))?;
let write_params = WriteParams {
mode: WriteMode::Create,
..Default::default()
};
let reader = RecordBatchIterator::new([Ok(empty_batch)], schema.clone());
Dataset::write(reader, path.to_string_lossy().as_ref(), Some(write_params))
.await
.map_err(|e| DBError(format!("Create dataset failed at {}: {}", path.display(), e)))?;
Ok(())
}
// Store/Upsert a single vector with ID and optional metadata (append; duplicate IDs are possible for now)
pub async fn store_vector(
&self,
name: &str,
id: &str,
vector: Vec<f32>,
meta: HashMap<String, String>,
text: Option<String>,
) -> Result<(), DBError> {
// Delegate to media-aware path with no media fields
self.store_vector_with_media(name, id, vector, meta, text, None, None).await
}
/// Store/Upsert a single vector with optional text and media fields (media_type/media_uri).
pub async fn store_vector_with_media(
&self,
name: &str,
id: &str,
vector: Vec<f32>,
meta: HashMap<String, String>,
text: Option<String>,
media_type: Option<String>,
media_uri: Option<String>,
) -> Result<(), DBError> {
let path = self.dataset_path(name);
// Determine dimension: use existing or infer from vector
let dim_i32 = if let Some(d) = self.read_existing_dim(name).await? {
d
} else {
vector
.len()
.try_into()
.map_err(|_| DBError("Vector length too large".into()))?
};
let (schema, batch) = Self::build_one_row_batch(
id,
&vector,
&meta,
text.as_deref(),
media_type.as_deref(),
media_uri.as_deref(),
dim_i32,
)?;
// If LanceDB table exists and provides delete, we can upsert by deleting same id
// Try best-effort delete; ignore errors to keep operation append-only on failure
if path.exists() {
if let Ok(conn) = self.connect_db().await {
if let Ok(mut tbl) = conn.open_table(name).execute().await {
let _ = tbl
.delete(&format!("id = '{}'", id.replace('\'', "''")))
.await;
}
}
}
let write_params = WriteParams {
mode: if path.exists() {
WriteMode::Append
} else {
WriteMode::Create
},
..Default::default()
};
let reader = RecordBatchIterator::new([Ok(batch)], schema.clone());
Dataset::write(reader, path.to_string_lossy().as_ref(), Some(write_params))
.await
.map_err(|e| DBError(format!("Write (append/create) failed: {}", e)))?;
Ok(())
}
// Delete a record by ID (best-effort; returns true if delete likely removed rows)
pub async fn delete_by_id(&self, name: &str, id: &str) -> Result<bool, DBError> {
let path = self.dataset_path(name);
if !path.exists() {
return Ok(false);
}
let conn = self.connect_db().await?;
let mut tbl = conn
.open_table(name)
.execute()
.await
.map_err(|e| DBError(format!("Open table '{}' failed: {}", name, e)))?;
// SQL-like predicate quoting
let pred = format!("id = '{}'", id.replace('\'', "''"));
// lancedb returns count or () depending on version; treat Ok as success
match tbl.delete(&pred).await {
Ok(_) => Ok(true),
Err(e) => Err(DBError(format!("Delete failed: {}", e))),
}
}
// Drop the entire dataset
pub async fn drop_dataset(&self, name: &str) -> Result<bool, DBError> {
let path = self.dataset_path(name);
// Try LanceDB drop first
// Best-effort logical drop via lancedb if available; ignore failures.
// Note: we rely on filesystem removal below for final cleanup.
if let Ok(conn) = self.connect_db().await {
if let Ok(mut t) = conn.open_table(name).execute().await {
// Best-effort delete-all to reduce footprint prior to fs removal
let _ = t.delete("true").await;
}
}
if path.exists() {
if path.is_dir() {
std::fs::remove_dir_all(&path)
.map_err(|e| DBError(format!("Failed to drop dataset '{}': {}", name, e)))?;
} else {
std::fs::remove_file(&path)
.map_err(|e| DBError(format!("Failed to drop dataset '{}': {}", name, e)))?;
}
return Ok(true);
}
Ok(false)
}
// Search top-k nearest with optional filter; returns tuple of (id, score (lower=L2), meta)
pub async fn search_vectors(
&self,
name: &str,
query: Vec<f32>,
k: usize,
filter: Option<String>,
return_fields: Option<Vec<String>>,
) -> Result<Vec<(String, f32, HashMap<String, String>)>, DBError> {
let path = self.dataset_path(name);
if !path.exists() {
return Err(DBError(format!("Dataset '{}' not found", name)));
}
// Determine dim and validate query length
let dim_i32 = self
.read_existing_dim(name)
.await?
.ok_or_else(|| DBError("Vector column not found".into()))?;
if query.len() as i32 != dim_i32 {
return Err(DBError(format!(
"Query vector length mismatch: expected {}, got {}",
dim_i32,
query.len()
)));
}
let ds = Dataset::open(path.to_string_lossy().as_ref())
.await
.map_err(|e| DBError(format!("Open dataset failed: {}", e)))?;
// Build scanner with projection; we project needed fields and filter client-side to support meta keys
let mut scan = ds.scan();
if let Err(e) = scan.project(&["id", "vector", "meta", "text", "media_type", "media_uri"]) {
return Err(DBError(format!("Project failed: {}", e)));
}
// Note: we no longer push down filter to Lance to allow filtering on meta fields client-side.
let mut stream = scan
.try_into_stream()
.await
.map_err(|e| DBError(format!("Scan stream failed: {}", e)))?;
// Parse simple equality clause from filter for client-side filtering (supports one `key = 'value'`)
let clause = filter.as_ref().and_then(|s| {
fn parse_eq(s: &str) -> Option<(String, String)> {
let s = s.trim();
let pos = s.find('=').or_else(|| s.find(" = "))?;
let (k, vraw) = s.split_at(pos);
let mut v = vraw.trim_start_matches('=').trim();
if (v.starts_with('\'') && v.ends_with('\'')) || (v.starts_with('"') && v.ends_with('"')) {
if v.len() >= 2 {
v = &v[1..v.len()-1];
}
}
let key = k.trim().trim_matches('"').trim_matches('\'').to_string();
if key.is_empty() { return None; }
Some((key, v.to_string()))
}
parse_eq(s)
});
// Maintain a max-heap with reverse ordering to keep top-k smallest distances
#[derive(Debug)]
struct Hit {
dist: f32,
id: String,
meta: HashMap<String, String>,
}
impl PartialEq for Hit {
fn eq(&self, other: &Self) -> bool {
self.dist.eq(&other.dist)
}
}
impl Eq for Hit {}
impl PartialOrd for Hit {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
// Reverse for max-heap: larger distance = "greater"
other.dist.partial_cmp(&self.dist)
}
}
impl Ord for Hit {
fn cmp(&self, other: &Self) -> Ordering {
self.partial_cmp(other).unwrap_or(Ordering::Equal)
}
}
let mut heap: BinaryHeap<Hit> = BinaryHeap::with_capacity(k);
while let Some(batch_res) = stream.next().await {
let batch = batch_res.map_err(|e| DBError(format!("Stream batch error: {}", e)))?;
let id_arr = batch
.column_by_name("id")
.ok_or_else(|| DBError("Column 'id' missing".into()))?
.as_string::<i32>();
let vec_arr = batch
.column_by_name("vector")
.ok_or_else(|| DBError("Column 'vector' missing".into()))?
.as_fixed_size_list();
let meta_arr = batch
.column_by_name("meta")
.map(|a| a.as_string::<i32>().clone());
let text_arr = batch
.column_by_name("text")
.map(|a| a.as_string::<i32>().clone());
let mt_arr = batch
.column_by_name("media_type")
.map(|a| a.as_string::<i32>().clone());
let mu_arr = batch
.column_by_name("media_uri")
.map(|a| a.as_string::<i32>().clone());
for i in 0..batch.num_rows() {
// Extract id
let id_val = id_arr.value(i).to_string();
// Parse meta JSON if present
let mut meta: HashMap<String, String> = HashMap::new();
if let Some(meta_col) = &meta_arr {
if !meta_col.is_null(i) {
let s = meta_col.value(i);
if let Ok(JsonValue::Object(map)) = serde_json::from_str::<JsonValue>(s) {
for (k, v) in map {
if let Some(vs) = v.as_str() {
meta.insert(k, vs.to_string());
} else if v.is_number() || v.is_boolean() {
meta.insert(k, v.to_string());
}
}
}
}
}
// Evaluate simple equality filter if provided (supports one clause)
let passes = if let Some((ref key, ref val)) = clause {
let candidate = match key.as_str() {
"id" => Some(id_val.clone()),
"text" => text_arr.as_ref().and_then(|col| if col.is_null(i) { None } else { Some(col.value(i).to_string()) }),
"media_type" => mt_arr.as_ref().and_then(|col| if col.is_null(i) { None } else { Some(col.value(i).to_string()) }),
"media_uri" => mu_arr.as_ref().and_then(|col| if col.is_null(i) { None } else { Some(col.value(i).to_string()) }),
_ => meta.get(key).cloned(),
};
match candidate {
Some(cv) => cv == *val,
None => false,
}
} else { true };
if !passes {
continue;
}
// Compute L2 distance
let val = vec_arr.value(i);
let prim = val.as_primitive::<Float32Type>();
let mut dist: f32 = 0.0;
let plen = prim.len();
for j in 0..plen {
let r = prim.value(j);
let d = query[j] - r;
dist += d * d;
}
// Apply return_fields on meta
let mut meta_out = meta;
if let Some(fields) = &return_fields {
let mut filtered = HashMap::new();
for f in fields {
if let Some(val) = meta_out.get(f) {
filtered.insert(f.clone(), val.clone());
}
}
meta_out = filtered;
}
let hit = Hit { dist, id: id_val, meta: meta_out };
if heap.len() < k {
heap.push(hit);
} else if let Some(top) = heap.peek() {
if hit.dist < top.dist {
heap.pop();
heap.push(hit);
}
}
}
}
// Extract and sort ascending by distance
let mut hits: Vec<Hit> = heap.into_sorted_vec(); // already ascending by dist due to Ord
let out = hits
.drain(..)
.map(|h| (h.id, h.dist, h.meta))
.collect::<Vec<_>>();
Ok(out)
}
// Create an ANN index on the vector column (IVF_PQ or similar)
pub async fn create_index(
&self,
name: &str,
index_type: &str,
params: HashMap<String, String>,
) -> Result<(), DBError> {
let path = self.dataset_path(name);
if !path.exists() {
return Err(DBError(format!("Dataset '{}' not found", name)));
}
// Attempt to create a vector index using lance low-level API if available.
// Some crate versions hide IndexType; to ensure build stability, we fall back to a no-op if the API is not accessible.
let _ = (index_type, params); // currently unused; reserved for future tuning
// TODO: Implement using lance::Dataset::create_index when public API is stable across versions.
// For now, succeed as a no-op to keep flows working; search will operate as brute-force scan.
Ok(())
}
// List datasets (tables) under this DB (show user-level logical names without .lance)
pub async fn list_datasets(&self) -> Result<Vec<String>, DBError> {
let mut out = Vec::new();
if self.base_dir.exists() {
if let Ok(rd) = std::fs::read_dir(&self.base_dir) {
for entry in rd.flatten() {
let p = entry.path();
if let Some(name) = p.file_name().and_then(|s| s.to_str()) {
// Only list .lance datasets
if name.ends_with(".lance") {
out.push(name.trim_end_matches(".lance").to_string());
}
}
}
}
}
Ok(out)
}
// Return basic dataset info map
pub async fn get_dataset_info(&self, name: &str) -> Result<HashMap<String, String>, DBError> {
let path = self.dataset_path(name);
let mut m = HashMap::new();
m.insert("name".to_string(), name.to_string());
m.insert("path".to_string(), path.display().to_string());
if !path.exists() {
return Err(DBError(format!("Dataset '{}' not found", name)));
}
let ds = Dataset::open(path.to_string_lossy().as_ref())
.await
.map_err(|e| DBError(format!("Open dataset failed: {}", e)))?;
// dim: infer by scanning first batch
let mut dim_str = "unknown".to_string();
{
let mut scan = ds.scan();
if scan.project(&["vector"]).is_ok() {
if let Ok(mut stream) = scan.try_into_stream().await {
if let Some(batch_res) = stream.next().await {
if let Ok(batch) = batch_res {
if let Some(col) = batch.column_by_name("vector") {
let fsl = col.as_fixed_size_list();
dim_str = fsl.value_length().to_string();
}
}
}
}
}
}
m.insert("dimension".to_string(), dim_str);
// row_count (approximate by scanning)
let mut scan = ds.scan();
if let Err(e) = scan.project(&["id"]) {
return Err(DBError(format!("Project failed: {e}")));
}
let mut stream = scan
.try_into_stream()
.await
.map_err(|e| DBError(format!("Scan failed: {e}")))?;
let mut rows: usize = 0;
while let Some(batch_res) = stream.next().await {
let batch = batch_res.map_err(|e| DBError(format!("Scan batch error: {}", e)))?;
rows += batch.num_rows();
}
m.insert("row_count".to_string(), rows.to_string());
// indexes: we can’t easily enumerate; set to "unknown" (future: read index metadata)
m.insert("indexes".to_string(), "unknown".to_string());
Ok(m)
}
}

2
src/lib.rs
View File

@@ -14,3 +14,5 @@ pub mod storage_sled;
pub mod admin_meta;
pub mod tantivy_search;
pub mod search_cmd;
pub mod lance_store;
pub mod embedding;

1
src/options.rs
View File

@@ -5,6 +5,7 @@ pub enum BackendType {
Redb,
Sled,
Tantivy, // Full-text search backend (no KV storage)
Lance, // Vector database backend (no KV storage)
}
#[derive(Debug, Clone)]

695
src/rpc.rs
View File

@@ -9,6 +9,8 @@ use sha2::{Digest, Sha256};
use crate::server::Server;
use crate::options::DBOption;
use crate::admin_meta;
use crate::embedding::{EmbeddingConfig, EmbeddingProvider};
use base64::{engine::general_purpose, Engine as _};
/// Database backend types
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -16,6 +18,7 @@ pub enum BackendType {
Redb,
Sled,
Tantivy, // Full-text search backend (no KV storage)
Lance, // Vector search backend (no KV storage)
// Future: InMemory, Custom(String)
}
@@ -161,6 +164,152 @@ pub trait Rpc {
/// Drop an FT index
#[method(name = "ftDrop")]
async fn ft_drop(&self, db_id: u64, index_name: String) -> RpcResult<bool>;
// ----- LanceDB (Vector + Text) RPC endpoints -----
/// Create a new Lance dataset in a Lance-backed DB
#[method(name = "lanceCreate")]
async fn lance_create(
&self,
db_id: u64,
name: String,
dim: usize,
) -> RpcResult<bool>;
/// Store a vector (with id and metadata) into a Lance dataset (deprecated; returns error)
#[method(name = "lanceStore")]
async fn lance_store(
&self,
db_id: u64,
name: String,
id: String,
vector: Vec<f32>,
meta: Option<HashMap<String, String>>,
) -> RpcResult<bool>;
/// Search a Lance dataset with a query vector (deprecated; returns error)
#[method(name = "lanceSearch")]
async fn lance_search(
&self,
db_id: u64,
name: String,
vector: Vec<f32>,
k: usize,
filter: Option<String>,
return_fields: Option<Vec<String>>,
) -> RpcResult<serde_json::Value>;
/// Create an ANN index on a Lance dataset
#[method(name = "lanceCreateIndex")]
async fn lance_create_index(
&self,
db_id: u64,
name: String,
index_type: String,
params: Option<HashMap<String, String>>,
) -> RpcResult<bool>;
/// List Lance datasets for a DB
#[method(name = "lanceList")]
async fn lance_list(
&self,
db_id: u64,
) -> RpcResult<Vec<String>>;
/// Get info for a Lance dataset
#[method(name = "lanceInfo")]
async fn lance_info(
&self,
db_id: u64,
name: String,
) -> RpcResult<serde_json::Value>;
/// Delete a record by id from a Lance dataset
#[method(name = "lanceDel")]
async fn lance_del(
&self,
db_id: u64,
name: String,
id: String,
) -> RpcResult<bool>;
/// Drop a Lance dataset
#[method(name = "lanceDrop")]
async fn lance_drop(
&self,
db_id: u64,
name: String,
) -> RpcResult<bool>;
// New: Text-first endpoints (no user-provided vectors)
/// Set per-dataset embedding configuration
#[method(name = "lanceSetEmbeddingConfig")]
async fn lance_set_embedding_config(
&self,
db_id: u64,
name: String,
provider: String,
model: String,
params: Option<HashMap<String, String>>,
) -> RpcResult<bool>;
/// Get per-dataset embedding configuration
#[method(name = "lanceGetEmbeddingConfig")]
async fn lance_get_embedding_config(
&self,
db_id: u64,
name: String,
) -> RpcResult<serde_json::Value>;
/// Store text; server will embed and store vector+text+meta
#[method(name = "lanceStoreText")]
async fn lance_store_text(
&self,
db_id: u64,
name: String,
id: String,
text: String,
meta: Option<HashMap<String, String>>,
) -> RpcResult<bool>;
/// Search using a text query; server will embed then search
#[method(name = "lanceSearchText")]
async fn lance_search_text(
&self,
db_id: u64,
name: String,
text: String,
k: usize,
filter: Option<String>,
return_fields: Option<Vec<String>>,
) -> RpcResult<serde_json::Value>;
// ----- Image-first endpoints (no user-provided vectors) -----
/// Store an image; exactly one of uri or bytes_b64 must be provided.
#[method(name = "lanceStoreImage")]
async fn lance_store_image(
&self,
db_id: u64,
name: String,
id: String,
uri: Option<String>,
bytes_b64: Option<String>,
meta: Option<HashMap<String, String>>,
) -> RpcResult<bool>;
/// Search using an image query; exactly one of uri or bytes_b64 must be provided.
#[method(name = "lanceSearchImage")]
async fn lance_search_image(
&self,
db_id: u64,
name: String,
k: usize,
uri: Option<String>,
bytes_b64: Option<String>,
filter: Option<String>,
return_fields: Option<Vec<String>>,
) -> RpcResult<serde_json::Value>;
}
/// RPC Server implementation
@@ -236,7 +385,10 @@ impl RpcServerImpl {
}
// Create server instance with resolved backend
let is_tantivy = matches!(effective_backend, crate::options::BackendType::Tantivy);
let is_search_only = matches!(
effective_backend,
crate::options::BackendType::Tantivy | crate::options::BackendType::Lance
);
let db_option = DBOption {
dir: self.base_dir.clone(),
port: 0, // Not used for RPC-managed databases
@@ -253,8 +405,8 @@ impl RpcServerImpl {
server.selected_db = db_id;
// Lazily open/create physical storage according to admin meta (per-db encryption)
// Skip for Tantivy backend (no KV storage to open)
if !is_tantivy {
// Skip for search-only backends (Tantivy/Lance): no KV storage to open
if !is_search_only {
let _ = server.current_storage();
}
@@ -344,6 +496,7 @@ impl RpcServerImpl {
crate::options::BackendType::Redb => BackendType::Redb,
crate::options::BackendType::Sled => BackendType::Sled,
crate::options::BackendType::Tantivy => BackendType::Tantivy,
crate::options::BackendType::Lance => BackendType::Lance,
};
DatabaseInfo {
@@ -395,12 +548,16 @@ impl RpcServer for RpcServerImpl {
BackendType::Redb => crate::options::BackendType::Redb,
BackendType::Sled => crate::options::BackendType::Sled,
BackendType::Tantivy => crate::options::BackendType::Tantivy,
BackendType::Lance => crate::options::BackendType::Lance,
};
admin_meta::set_database_backend(&self.base_dir, self.backend.clone(), &self.admin_secret, db_id, opt_backend.clone())
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
// Create server instance using base_dir, chosen backend and admin secret
let is_tantivy_new = matches!(opt_backend, crate::options::BackendType::Tantivy);
let is_search_only_new = matches!(
opt_backend,
crate::options::BackendType::Tantivy | crate::options::BackendType::Lance
);
let option = DBOption {
dir: self.base_dir.clone(),
port: 0, // Not used for RPC-managed databases
@@ -415,8 +572,8 @@ impl RpcServer for RpcServerImpl {
server.selected_db = db_id;
// Initialize storage to create physical <id>.db with proper encryption from admin meta
// Skip for Tantivy backend (no KV storage to initialize)
if !is_tantivy_new {
// Skip for search-only backends (Tantivy/Lance): no KV storage to initialize
if !is_search_only_new {
let _ = server.current_storage();
}
@@ -676,4 +833,530 @@ impl RpcServer for RpcServerImpl {
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
Ok(true)
}
// ----- LanceDB (Vector) RPC endpoints -----
async fn lance_create(
&self,
db_id: u64,
name: String,
dim: usize,
) -> RpcResult<bool> {
let server = self.get_or_create_server(db_id).await?;
if db_id == 0 {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
}
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
}
if !server.has_write_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
}
server.lance_store()
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
.create_dataset(&name, dim).await
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
Ok(true)
}
async fn lance_store(
&self,
_db_id: u64,
_name: String,
_id: String,
_vector: Vec<f32>,
_meta: Option<HashMap<String, String>>,
) -> RpcResult<bool> {
Err(jsonrpsee::types::ErrorObjectOwned::owned(
-32000,
"Vector endpoint removed. Use lanceStoreText instead.",
None::<()>
))
}
async fn lance_search(
&self,
_db_id: u64,
_name: String,
_vector: Vec<f32>,
_k: usize,
_filter: Option<String>,
_return_fields: Option<Vec<String>>,
) -> RpcResult<serde_json::Value> {
Err(jsonrpsee::types::ErrorObjectOwned::owned(
-32000,
"Vector endpoint removed. Use lanceSearchText instead.",
None::<()>
))
}
async fn lance_create_index(
&self,
db_id: u64,
name: String,
index_type: String,
params: Option<HashMap<String, String>>,
) -> RpcResult<bool> {
let server = self.get_or_create_server(db_id).await?;
if db_id == 0 {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
}
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
}
if !server.has_write_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
}
server.lance_store()
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
.create_index(&name, &index_type, params.unwrap_or_default()).await
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
Ok(true)
}
async fn lance_list(
&self,
db_id: u64,
) -> RpcResult<Vec<String>> {
let server = self.get_or_create_server(db_id).await?;
if db_id == 0 {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
}
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
}
if !server.has_read_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "read permission denied", None::<()>));
}
let list = server.lance_store()
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
.list_datasets().await
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
Ok(list)
}
async fn lance_info(
&self,
db_id: u64,
name: String,
) -> RpcResult<serde_json::Value> {
let server = self.get_or_create_server(db_id).await?;
if db_id == 0 {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
}
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
}
if !server.has_read_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "read permission denied", None::<()>));
}
let info = server.lance_store()
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
.get_dataset_info(&name).await
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
Ok(serde_json::json!(info))
}
async fn lance_del(
&self,
db_id: u64,
name: String,
id: String,
) -> RpcResult<bool> {
let server = self.get_or_create_server(db_id).await?;
if db_id == 0 {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
}
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
}
if !server.has_write_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
}
let ok = server.lance_store()
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
.delete_by_id(&name, &id).await
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
Ok(ok)
}
async fn lance_drop(
&self,
db_id: u64,
name: String,
) -> RpcResult<bool> {
let server = self.get_or_create_server(db_id).await?;
if db_id == 0 {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
}
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
}
if !server.has_write_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
}
let ok = server.lance_store()
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
.drop_dataset(&name).await
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
Ok(ok)
}
// ----- New text-first Lance RPC implementations -----
async fn lance_set_embedding_config(
&self,
db_id: u64,
name: String,
provider: String,
model: String,
params: Option<HashMap<String, String>>,
) -> RpcResult<bool> {
let server = self.get_or_create_server(db_id).await?;
if db_id == 0 {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
}
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
}
if !server.has_write_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
}
let prov = match provider.to_lowercase().as_str() {
"test-hash" | "testhash" => EmbeddingProvider::TestHash,
"testimagehash" | "image-test-hash" | "imagetesthash" => EmbeddingProvider::ImageTestHash,
"fastembed" | "lancefastembed" => EmbeddingProvider::LanceFastEmbed,
"openai" | "lanceopenai" => EmbeddingProvider::LanceOpenAI,
other => EmbeddingProvider::LanceOther(other.to_string()),
};
let cfg = EmbeddingConfig {
provider: prov,
model,
params: params.unwrap_or_default(),
};
server.set_dataset_embedding_config(&name, &cfg)
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
Ok(true)
}
async fn lance_get_embedding_config(
&self,
db_id: u64,
name: String,
) -> RpcResult<serde_json::Value> {
let server = self.get_or_create_server(db_id).await?;
if db_id == 0 {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
}
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
}
if !server.has_read_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "read permission denied", None::<()>));
}
let cfg = server.get_dataset_embedding_config(&name)
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
Ok(serde_json::json!({
"provider": match cfg.provider {
EmbeddingProvider::TestHash => "test-hash",
EmbeddingProvider::ImageTestHash => "testimagehash",
EmbeddingProvider::LanceFastEmbed => "lancefastembed",
EmbeddingProvider::LanceOpenAI => "lanceopenai",
EmbeddingProvider::LanceOther(ref s) => s,
},
"model": cfg.model,
"params": cfg.params
}))
}
async fn lance_store_text(
&self,
db_id: u64,
name: String,
id: String,
text: String,
meta: Option<HashMap<String, String>>,
) -> RpcResult<bool> {
let server = self.get_or_create_server(db_id).await?;
if db_id == 0 {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
}
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
}
if !server.has_write_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
}
// Resolve embedder and run blocking embedding off the async runtime
// Resolve embedder and run embedding on a plain OS thread (avoid dropping any runtime in async context)
let embedder = server.get_embedder_for(&name)
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
let (tx, rx) = tokio::sync::oneshot::channel();
let emb_arc = embedder.clone();
let text_cl = text.clone();
std::thread::spawn(move || {
let res = emb_arc.embed(&text_cl);
let _ = tx.send(res);
});
let vector = match rx.await {
Ok(Ok(v)) => v,
Ok(Err(e)) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>)),
Err(recv_err) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, format!("embedding thread error: {}", recv_err), None::<()>)),
};
server.lance_store()
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
.store_vector(&name, &id, vector, meta.unwrap_or_default(), Some(text)).await
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
Ok(true)
}
async fn lance_search_text(
&self,
db_id: u64,
name: String,
text: String,
k: usize,
filter: Option<String>,
return_fields: Option<Vec<String>>,
) -> RpcResult<serde_json::Value> {
let server = self.get_or_create_server(db_id).await?;
if db_id == 0 {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
}
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
}
if !server.has_read_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "read permission denied", None::<()>));
}
// Resolve embedder and run embedding on a plain OS thread (avoid dropping any runtime in async context)
let embedder = server.get_embedder_for(&name)
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
let (tx, rx) = tokio::sync::oneshot::channel();
let emb_arc = embedder.clone();
let text_cl = text.clone();
std::thread::spawn(move || {
let res = emb_arc.embed(&text_cl);
let _ = tx.send(res);
});
let qv = match rx.await {
Ok(Ok(v)) => v,
Ok(Err(e)) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>)),
Err(recv_err) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, format!("embedding thread error: {}", recv_err), None::<()>)),
};
let results = server.lance_store()
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
.search_vectors(&name, qv, k, filter, return_fields).await
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
let json_results: Vec<serde_json::Value> = results.into_iter().map(|(id, score, meta)| {
serde_json::json!({
"id": id,
"score": score,
"meta": meta,
})
}).collect();
Ok(serde_json::json!({ "results": json_results }))
}
// ----- New image-first Lance RPC implementations -----
async fn lance_store_image(
&self,
db_id: u64,
name: String,
id: String,
uri: Option<String>,
bytes_b64: Option<String>,
meta: Option<HashMap<String, String>>,
) -> RpcResult<bool> {
let server = self.get_or_create_server(db_id).await?;
if db_id == 0 {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
}
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
}
if !server.has_write_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
}
// Validate exactly one of uri or bytes_b64
let (use_uri, use_b64) = (uri.is_some(), bytes_b64.is_some());
if (use_uri && use_b64) || (!use_uri && !use_b64) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(
-32000,
"Provide exactly one of 'uri' or 'bytes_b64'",
None::<()>,
));
}
// Acquire image bytes (with caps)
let max_bytes: usize = std::env::var("HERODB_IMAGE_MAX_BYTES")
.ok()
.and_then(|s| s.parse::<u64>().ok())
.unwrap_or(10 * 1024 * 1024) as usize;
let (bytes, media_uri_opt) = if let Some(u) = uri.clone() {
let data = server
.fetch_image_bytes_from_uri(&u)
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
(data, Some(u))
} else {
let b64 = bytes_b64.unwrap_or_default();
let data = general_purpose::STANDARD
.decode(b64.as_bytes())
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, format!("base64 decode error: {}", e), None::<()>))?;
if data.len() > max_bytes {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(
-32000,
format!("Image exceeds max allowed bytes {}", max_bytes),
None::<()>,
));
}
(data, None)
};
// Resolve image embedder and embed on a plain OS thread
let img_embedder = server
.get_image_embedder_for(&name)
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
let (tx, rx) = tokio::sync::oneshot::channel();
let emb_arc = img_embedder.clone();
let bytes_cl = bytes.clone();
std::thread::spawn(move || {
let res = emb_arc.embed_image(&bytes_cl);
let _ = tx.send(res);
});
let vector = match rx.await {
Ok(Ok(v)) => v,
Ok(Err(e)) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>)),
Err(recv_err) => {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(
-32000,
format!("embedding thread error: {}", recv_err),
None::<()>,
))
}
};
// Store vector with media fields
server
.lance_store()
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
.store_vector_with_media(
&name,
&id,
vector,
meta.unwrap_or_default(),
None,
Some("image".to_string()),
media_uri_opt,
)
.await
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
Ok(true)
}
async fn lance_search_image(
&self,
db_id: u64,
name: String,
k: usize,
uri: Option<String>,
bytes_b64: Option<String>,
filter: Option<String>,
return_fields: Option<Vec<String>>,
) -> RpcResult<serde_json::Value> {
let server = self.get_or_create_server(db_id).await?;
if db_id == 0 {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
}
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
}
if !server.has_read_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "read permission denied", None::<()>));
}
// Validate exactly one of uri or bytes_b64
let (use_uri, use_b64) = (uri.is_some(), bytes_b64.is_some());
if (use_uri && use_b64) || (!use_uri && !use_b64) {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(
-32000,
"Provide exactly one of 'uri' or 'bytes_b64'",
None::<()>,
));
}
// Acquire image bytes for query (with caps)
let max_bytes: usize = std::env::var("HERODB_IMAGE_MAX_BYTES")
.ok()
.and_then(|s| s.parse::<u64>().ok())
.unwrap_or(10 * 1024 * 1024) as usize;
let bytes = if let Some(u) = uri {
server
.fetch_image_bytes_from_uri(&u)
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
} else {
let b64 = bytes_b64.unwrap_or_default();
let data = general_purpose::STANDARD
.decode(b64.as_bytes())
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, format!("base64 decode error: {}", e), None::<()>))?;
if data.len() > max_bytes {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(
-32000,
format!("Image exceeds max allowed bytes {}", max_bytes),
None::<()>,
));
}
data
};
// Resolve image embedder and embed on OS thread
let img_embedder = server
.get_image_embedder_for(&name)
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
let (tx, rx) = tokio::sync::oneshot::channel();
let emb_arc = img_embedder.clone();
std::thread::spawn(move || {
let res = emb_arc.embed_image(&bytes);
let _ = tx.send(res);
});
let qv = match rx.await {
Ok(Ok(v)) => v,
Ok(Err(e)) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>)),
Err(recv_err) => {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(
-32000,
format!("embedding thread error: {}", recv_err),
None::<()>,
))
}
};
// KNN search and return results
let results = server
.lance_store()
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
.search_vectors(&name, qv, k, filter, return_fields)
.await
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
let json_results: Vec<serde_json::Value> = results
.into_iter()
.map(|(id, score, meta)| {
serde_json::json!({
"id": id,
"score": score,
"meta": meta,
})
})
.collect();
Ok(serde_json::json!({ "results": json_results }))
}
}

248
src/server.rs
View File

@@ -14,6 +14,15 @@ use crate::protocol::Protocol;
use crate::storage_trait::StorageBackend;
use crate::admin_meta;
// Embeddings: config and cache
use crate::embedding::{EmbeddingConfig, create_embedder, Embedder, create_image_embedder, ImageEmbedder};
use serde_json;
use ureq::{Agent, AgentBuilder};
use std::time::Duration;
use std::io::Read;
const NO_DB_SELECTED: u64 = u64::MAX;
#[derive(Clone)]
pub struct Server {
pub db_cache: std::sync::Arc<std::sync::RwLock<HashMap<u64, Arc<dyn StorageBackend>>>>,
@@ -26,6 +35,15 @@ pub struct Server {
// In-memory registry of Tantivy search indexes for this server
pub search_indexes: Arc<std::sync::RwLock<HashMap<String, Arc<crate::tantivy_search::TantivySearch>>>>,
// Per-DB Lance stores (vector DB), keyed by db_id
pub lance_stores: Arc<std::sync::RwLock<HashMap<u64, Arc<crate::lance_store::LanceStore>>>>,
// Per-(db_id, dataset) embedder cache (text)
pub embedders: Arc<std::sync::RwLock<HashMap<(u64, String), Arc<dyn Embedder>>>>,
// Per-(db_id, dataset) image embedder cache (image)
pub image_embedders: Arc<std::sync::RwLock<HashMap<(u64, String), Arc<dyn ImageEmbedder>>>>,
// BLPOP waiter registry: per (db_index, key) FIFO of waiters
pub list_waiters: Arc<Mutex<HashMap<u64, HashMap<String, Vec<Waiter>>>>>,
pub waiter_seq: Arc<AtomicU64>,
@@ -49,11 +67,14 @@ impl Server {
db_cache: Arc::new(std::sync::RwLock::new(HashMap::new())),
option,
client_name: None,
selected_db: 0,
selected_db: NO_DB_SELECTED,
queued_cmd: None,
current_permissions: None,
search_indexes: Arc::new(std::sync::RwLock::new(HashMap::new())),
lance_stores: Arc::new(std::sync::RwLock::new(HashMap::new())),
embedders: Arc::new(std::sync::RwLock::new(HashMap::new())),
image_embedders: Arc::new(std::sync::RwLock::new(HashMap::new())),
list_waiters: Arc::new(Mutex::new(HashMap::new())),
waiter_seq: Arc::new(AtomicU64::new(1)),
}
@@ -71,7 +92,30 @@ impl Server {
base
}
// Path where Lance datasets are stored, namespaced per selected DB:
// <base_dir>/lance/<db_id>
pub fn lance_data_path(&self) -> std::path::PathBuf {
let base = std::path::PathBuf::from(&self.option.dir)
.join("lance")
.join(self.selected_db.to_string());
if !base.exists() {
let _ = std::fs::create_dir_all(&base);
}
base
}
pub fn current_storage(&self) -> Result<Arc<dyn StorageBackend>, DBError> {
// Require explicit SELECT before any storage access
if self.selected_db == NO_DB_SELECTED {
return Err(DBError("No database selected. Use SELECT <id> [KEY <key>] first".to_string()));
}
// Admin DB 0 access must be authenticated with SELECT 0 KEY <admin_secret>
if self.selected_db == 0 {
if !matches!(self.current_permissions, Some(crate::rpc::Permissions::ReadWrite)) {
return Err(DBError("Admin DB 0 requires SELECT 0 KEY <admin_secret>".to_string()));
}
}
let mut cache = self.db_cache.write().unwrap();
if let Some(storage) = cache.get(&self.selected_db) {
@@ -100,9 +144,207 @@ impl Server {
Ok(storage)
}
/// Get or create the LanceStore for the currently selected DB.
/// Only valid for non-zero DBs and when the backend is Lance.
pub fn lance_store(&self) -> Result<Arc<crate::lance_store::LanceStore>, DBError> {
if self.selected_db == 0 {
return Err(DBError("Lance not available on admin DB 0".to_string()));
}
// Resolve backend for selected_db
let backend_opt = crate::admin_meta::get_database_backend(
&self.option.dir,
self.option.backend.clone(),
&self.option.admin_secret,
self.selected_db,
)
.ok()
.flatten();
if !matches!(backend_opt, Some(crate::options::BackendType::Lance)) {
return Err(DBError("ERR DB backend is not Lance; LANCE.* commands are not allowed".to_string()));
}
// Fast path: read lock
{
let map = self.lance_stores.read().unwrap();
if let Some(store) = map.get(&self.selected_db) {
return Ok(store.clone());
}
}
// Slow path: create and insert
let store = Arc::new(crate::lance_store::LanceStore::new(&self.option.dir, self.selected_db)?);
{
let mut map = self.lance_stores.write().unwrap();
map.insert(self.selected_db, store.clone());
}
Ok(store)
}
// ----- Embedding configuration and resolution -----
// Sidecar embedding config path: <base_dir>/lance/<db_id>/<dataset>.lance.embedding.json
fn dataset_embedding_config_path(&self, dataset: &str) -> std::path::PathBuf {
let mut base = self.lance_data_path();
// Ensure parent dir exists
if !base.exists() {
let _ = std::fs::create_dir_all(&base);
}
base.push(format!("{}.lance.embedding.json", dataset));
base
}
/// Persist per-dataset embedding config as JSON sidecar.
pub fn set_dataset_embedding_config(&self, dataset: &str, cfg: &EmbeddingConfig) -> Result<(), DBError> {
if self.selected_db == 0 {
return Err(DBError("Lance not available on admin DB 0".to_string()));
}
let p = self.dataset_embedding_config_path(dataset);
let data = serde_json::to_vec_pretty(cfg)
.map_err(|e| DBError(format!("Failed to serialize embedding config: {}", e)))?;
std::fs::write(&p, data)
.map_err(|e| DBError(format!("Failed to write embedding config {}: {}", p.display(), e)))?;
// Invalidate embedder cache entry for this dataset
{
let mut map = self.embedders.write().unwrap();
map.remove(&(self.selected_db, dataset.to_string()));
}
{
let mut map_img = self.image_embedders.write().unwrap();
map_img.remove(&(self.selected_db, dataset.to_string()));
}
Ok(())
}
/// Load per-dataset embedding config.
pub fn get_dataset_embedding_config(&self, dataset: &str) -> Result<EmbeddingConfig, DBError> {
if self.selected_db == 0 {
return Err(DBError("Lance not available on admin DB 0".to_string()));
}
let p = self.dataset_embedding_config_path(dataset);
if !p.exists() {
return Err(DBError(format!(
"Embedding config not set for dataset '{}'. Use LANCE.EMBEDDING CONFIG SET ... or RPC to configure.",
dataset
)));
}
let data = std::fs::read(&p)
.map_err(|e| DBError(format!("Failed to read embedding config {}: {}", p.display(), e)))?;
let cfg: EmbeddingConfig = serde_json::from_slice(&data)
.map_err(|e| DBError(format!("Failed to parse embedding config {}: {}", p.display(), e)))?;
Ok(cfg)
}
/// Resolve or build an embedder for (db_id, dataset). Caches instance.
pub fn get_embedder_for(&self, dataset: &str) -> Result<Arc<dyn Embedder>, DBError> {
if self.selected_db == 0 {
return Err(DBError("Lance not available on admin DB 0".to_string()));
}
// Fast path
{
let map = self.embedders.read().unwrap();
if let Some(e) = map.get(&(self.selected_db, dataset.to_string())) {
return Ok(e.clone());
}
}
// Load config and instantiate
let cfg = self.get_dataset_embedding_config(dataset)?;
let emb = create_embedder(&cfg)?;
{
let mut map = self.embedders.write().unwrap();
map.insert((self.selected_db, dataset.to_string()), emb.clone());
}
Ok(emb)
}
/// Resolve or build an IMAGE embedder for (db_id, dataset). Caches instance.
pub fn get_image_embedder_for(&self, dataset: &str) -> Result<Arc<dyn ImageEmbedder>, DBError> {
if self.selected_db == 0 {
return Err(DBError("Lance not available on admin DB 0".to_string()));
}
// Fast path
{
let map = self.image_embedders.read().unwrap();
if let Some(e) = map.get(&(self.selected_db, dataset.to_string())) {
return Ok(e.clone());
}
}
// Load config and instantiate
let cfg = self.get_dataset_embedding_config(dataset)?;
let emb = create_image_embedder(&cfg)?;
{
let mut map = self.image_embedders.write().unwrap();
map.insert((self.selected_db, dataset.to_string()), emb.clone());
}
Ok(emb)
}
/// Download image bytes from a URI with safety checks (size, timeout, content-type, optional host allowlist).
/// Env overrides:
/// - HERODB_IMAGE_MAX_BYTES (u64, default 10485760)
/// - HERODB_IMAGE_FETCH_TIMEOUT_SECS (u64, default 30)
/// - HERODB_IMAGE_ALLOWED_HOSTS (comma-separated, optional)
pub fn fetch_image_bytes_from_uri(&self, uri: &str) -> Result<Vec<u8>, DBError> {
// Basic scheme validation
if !(uri.starts_with("http://") || uri.starts_with("https://")) {
return Err(DBError("Only http(s) URIs are supported for image fetch".into()));
}
// Parse host (naive) for allowlist check
let host = {
let after_scheme = match uri.find("://") {
Some(i) => &uri[i + 3..],
None => uri,
};
let end = after_scheme.find('/').unwrap_or(after_scheme.len());
let host_port = &after_scheme[..end];
host_port.split('@').last().unwrap_or(host_port).split(':').next().unwrap_or(host_port).to_string()
};
let max_bytes: u64 = std::env::var("HERODB_IMAGE_MAX_BYTES").ok().and_then(|s| s.parse::<u64>().ok()).unwrap_or(10 * 1024 * 1024);
let timeout_secs: u64 = std::env::var("HERODB_IMAGE_FETCH_TIMEOUT_SECS").ok().and_then(|s| s.parse::<u64>().ok()).unwrap_or(30);
let allowed_hosts_env = std::env::var("HERODB_IMAGE_ALLOWED_HOSTS").ok();
if let Some(allow) = allowed_hosts_env {
if !allow.split(',').map(|s| s.trim()).filter(|s| !s.is_empty()).any(|h| h.eq_ignore_ascii_case(&host)) {
return Err(DBError(format!("Host '{}' not allowed for image fetch (HERODB_IMAGE_ALLOWED_HOSTS)", host)));
}
}
let agent: Agent = AgentBuilder::new()
.timeout_read(Duration::from_secs(timeout_secs))
.timeout_write(Duration::from_secs(timeout_secs))
.build();
let resp = agent.get(uri).call().map_err(|e| DBError(format!("HTTP GET failed: {}", e)))?;
// Validate content-type
let ctype = resp.header("Content-Type").unwrap_or("");
let ctype_main = ctype.split(';').next().unwrap_or("").trim().to_ascii_lowercase();
if !ctype_main.starts_with("image/") {
return Err(DBError(format!("Remote content-type '{}' is not image/*", ctype)));
}
// Read with cap
let mut reader = resp.into_reader();
let mut buf: Vec<u8> = Vec::with_capacity(8192);
let mut tmp = [0u8; 8192];
let mut total: u64 = 0;
loop {
let n = reader.read(&mut tmp).map_err(|e| DBError(format!("Read error: {}", e)))?;
if n == 0 { break; }
total += n as u64;
if total > max_bytes {
return Err(DBError(format!("Image exceeds max allowed bytes {}", max_bytes)));
}
buf.extend_from_slice(&tmp[..n]);
}
Ok(buf)
}
/// Check if current permissions allow read operations
pub fn has_read_permission(&self) -> bool {
// No DB selected -> no permissions
if self.selected_db == NO_DB_SELECTED {
return false;
}
// If an explicit permission is set for this connection, honor it.
if let Some(perms) = self.current_permissions.as_ref() {
return matches!(*perms, crate::rpc::Permissions::Read | crate::rpc::Permissions::ReadWrite);
@@ -122,6 +364,10 @@ impl Server {
/// Check if current permissions allow write operations
pub fn has_write_permission(&self) -> bool {
// No DB selected -> no permissions
if self.selected_db == NO_DB_SELECTED {
return false;
}
// If an explicit permission is set for this connection, honor it.
if let Some(perms) = self.current_permissions.as_ref() {
return matches!(*perms, crate::rpc::Permissions::ReadWrite);

484
tests/lance_integration_tests.rs Normal file
View File

@@ -0,0 +1,484 @@
use redis::{Client, Connection, RedisResult, Value};
use std::process::{Child, Command};
use std::time::Duration;
use jsonrpsee::http_client::{HttpClient, HttpClientBuilder};
use herodb::rpc::{BackendType, DatabaseConfig, RpcClient};
use base64::Engine;
use tokio::time::sleep;
// ------------------------
// Helpers
// ------------------------
fn get_redis_connection(port: u16) -> Connection {
let connection_info = format!("redis://127.0.0.1:{}", port);
let client = Client::open(connection_info).unwrap();
let mut attempts = 0;
loop {
match client.get_connection() {
Ok(mut conn) => {
if redis::cmd("PING").query::<String>(&mut conn).is_ok() {
return conn;
}
}
Err(e) => {
if attempts >= 3600 {
panic!("Failed to connect to Redis server after 3600 attempts: {}", e);
}
}
}
attempts += 1;
std::thread::sleep(Duration::from_millis(500));
}
}
async fn get_rpc_client(port: u16) -> HttpClient {
let url = format!("http://127.0.0.1:{}", port + 1); // RPC port = Redis port + 1
HttpClientBuilder::default().build(url).unwrap()
}
/// Wait until RPC server is responsive (getServerStats succeeds) or panic after retries.
async fn wait_for_rpc_ready(client: &HttpClient, max_attempts: u32, delay: Duration) {
for _ in 0..max_attempts {
match client.get_server_stats().await {
Ok(_) => return,
Err(_) => {
sleep(delay).await;
}
}
}
panic!("RPC server did not become ready in time");
}
// A guard to ensure the server process is killed when it goes out of scope and test dir cleaned.
struct ServerProcessGuard {
process: Child,
test_dir: String,
}
impl Drop for ServerProcessGuard {
fn drop(&mut self) {
eprintln!("Killing server process (pid: {})...", self.process.id());
if let Err(e) = self.process.kill() {
eprintln!("Failed to kill server process: {}", e);
}
match self.process.wait() {
Ok(status) => eprintln!("Server process exited with: {}", status),
Err(e) => eprintln!("Failed to wait on server process: {}", e),
}
// Clean up the specific test directory
eprintln!("Cleaning up test directory: {}", self.test_dir);
if let Err(e) = std::fs::remove_dir_all(&self.test_dir) {
eprintln!("Failed to clean up test directory: {}", e);
}
}
}
// Helper to set up the server and return guard + ports
async fn setup_server() -> (ServerProcessGuard, u16) {
use std::sync::atomic::{AtomicU16, Ordering};
static PORT_COUNTER: AtomicU16 = AtomicU16::new(17500);
let port = PORT_COUNTER.fetch_add(1, Ordering::SeqCst);
let test_dir = format!("/tmp/herodb_lance_test_{}", port);
// Clean up previous test data
if std::path::Path::new(&test_dir).exists() {
let _ = std::fs::remove_dir_all(&test_dir);
}
std::fs::create_dir_all(&test_dir).unwrap();
// Start the server in a subprocess with RPC enabled (follows tantivy test pattern)
let child = Command::new("cargo")
.args(&[
"run",
"--",
"--dir",
&test_dir,
"--port",
&port.to_string(),
"--rpc-port",
&(port + 1).to_string(),
"--enable-rpc",
"--debug",
"--admin-secret",
"test-admin",
])
.spawn()
.expect("Failed to start server process");
let guard = ServerProcessGuard {
process: child,
test_dir,
};
// Give the server time to build and start (cargo run may compile first)
// Increase significantly to accommodate first-time dependency compilation in CI.
std::thread::sleep(Duration::from_millis(60000));
(guard, port)
}
// Convenient helpers for assertions on redis::Value
fn value_is_ok(v: &Value) -> bool {
match v {
Value::Okay => true,
Value::Status(s) if s == "OK" => true,
Value::Data(d) if d == b"OK" => true,
_ => false,
}
}
fn value_is_int_eq(v: &Value, expected: i64) -> bool {
matches!(v, Value::Int(n) if *n == expected)
}
fn value_is_str_eq(v: &Value, expected: &str) -> bool {
match v {
Value::Status(s) => s == expected,
Value::Data(d) => String::from_utf8_lossy(d) == expected,
_ => false,
}
}
fn to_string_lossy(v: &Value) -> String {
match v {
Value::Nil => "Nil".to_string(),
Value::Int(n) => n.to_string(),
Value::Status(s) => s.clone(),
Value::Okay => "OK".to_string(),
Value::Data(d) => String::from_utf8_lossy(d).to_string(),
Value::Bulk(items) => {
let inner: Vec<String> = items.iter().map(to_string_lossy).collect();
format!("[{}]", inner.join(", "))
}
}
}
// Extract ids from LANCE.SEARCH / LANCE.SEARCHIMAGE reply which is:
// Array of elements: [ [id, score, [k,v,...]], [id, score, ...], ... ]
fn extract_hit_ids(v: &Value) -> Vec<String> {
let mut ids = Vec::new();
if let Value::Bulk(items) = v {
for item in items {
if let Value::Bulk(row) = item {
if !row.is_empty() {
// first element is id (Data or Status)
let id = match &row[0] {
Value::Data(d) => String::from_utf8_lossy(d).to_string(),
Value::Status(s) => s.clone(),
Value::Int(n) => n.to_string(),
_ => continue,
};
ids.push(id);
}
}
}
}
ids
}
// Check whether a Bulk array (RESP array) contains a given string element.
fn bulk_contains_string(v: &Value, needle: &str) -> bool {
match v {
Value::Bulk(items) => items.iter().any(|it| match it {
Value::Data(d) => String::from_utf8_lossy(d).contains(needle),
Value::Status(s) => s.contains(needle),
Value::Bulk(_) => bulk_contains_string(it, needle),
_ => false,
}),
_ => false,
}
}
// ------------------------
// Test: Lance end-to-end (RESP) using only local embedders
// ------------------------
#[tokio::test]
async fn test_lance_end_to_end() {
let (_guard, port) = setup_server().await;
// First, wait for RESP to be available; this also gives cargo-run child ample time to finish building.
// Reuse the helper that retries PING until success.
{
let _conn_ready = get_redis_connection(port);
// Drop immediately; we only needed readiness.
}
// Build RPC client and create a Lance DB
let rpc_client = get_rpc_client(port).await;
// Ensure RPC server is listening before we issue createDatabase (allow longer warm-up to accommodate first-build costs)
wait_for_rpc_ready(&rpc_client, 3600, Duration::from_millis(250)).await;
let db_config = DatabaseConfig {
name: Some("media-db".to_string()),
storage_path: None,
max_size: None,
redis_version: None,
};
let db_id = rpc_client
.create_database(BackendType::Lance, db_config, None)
.await
.expect("create_database Lance failed");
assert_eq!(db_id, 1, "Expected first Lance DB id to be 1");
// Add access keys
let _ = rpc_client
.add_access_key(db_id, "readwrite_key".to_string(), "readwrite".to_string())
.await
.expect("add_access_key readwrite failed");
let _ = rpc_client
.add_access_key(db_id, "read_key".to_string(), "read".to_string())
.await
.expect("add_access_key read failed");
// Connect to Redis and SELECT DB with readwrite key
let mut conn = get_redis_connection(port);
let sel_ok: RedisResult<String> = redis::cmd("SELECT")
.arg(db_id)
.arg("KEY")
.arg("readwrite_key")
.query(&mut conn);
assert!(sel_ok.is_ok(), "SELECT db with key failed: {:?}", sel_ok);
assert_eq!(sel_ok.unwrap(), "OK");
// 1) Configure embedding providers: textset -> testhash dim 64, imageset -> testimagehash dim 512
let v = redis::cmd("LANCE.EMBEDDING")
.arg("CONFIG")
.arg("SET")
.arg("textset")
.arg("PROVIDER")
.arg("testhash")
.arg("MODEL")
.arg("any")
.arg("PARAM")
.arg("dim")
.arg("64")
.query::<Value>(&mut conn)
.unwrap();
assert!(value_is_ok(&v), "Embedding config set (text) not OK: {}", to_string_lossy(&v));
let v = redis::cmd("LANCE.EMBEDDING")
.arg("CONFIG")
.arg("SET")
.arg("imageset")
.arg("PROVIDER")
.arg("testimagehash")
.arg("MODEL")
.arg("any")
.arg("PARAM")
.arg("dim")
.arg("512")
.query::<Value>(&mut conn)
.unwrap();
assert!(value_is_ok(&v), "Embedding config set (image) not OK: {}", to_string_lossy(&v));
// 2) Create datasets
let v = redis::cmd("LANCE.CREATE")
.arg("textset")
.arg("DIM")
.arg(64)
.query::<Value>(&mut conn)
.unwrap();
assert!(value_is_ok(&v), "LANCE.CREATE textset failed: {}", to_string_lossy(&v));
let v = redis::cmd("LANCE.CREATE")
.arg("imageset")
.arg("DIM")
.arg(512)
.query::<Value>(&mut conn)
.unwrap();
assert!(value_is_ok(&v), "LANCE.CREATE imageset failed: {}", to_string_lossy(&v));
// 3) Store two text documents
let v = redis::cmd("LANCE.STORE")
.arg("textset")
.arg("ID")
.arg("doc-1")
.arg("TEXT")
.arg("The quick brown fox jumps over the lazy dog")
.arg("META")
.arg("title")
.arg("Fox")
.arg("category")
.arg("animal")
.query::<Value>(&mut conn)
.unwrap();
assert!(value_is_ok(&v), "LANCE.STORE doc-1 failed: {}", to_string_lossy(&v));
let v = redis::cmd("LANCE.STORE")
.arg("textset")
.arg("ID")
.arg("doc-2")
.arg("TEXT")
.arg("A fast auburn fox vaulted a sleepy canine")
.arg("META")
.arg("title")
.arg("Paraphrase")
.arg("category")
.arg("animal")
.query::<Value>(&mut conn)
.unwrap();
assert!(value_is_ok(&v), "LANCE.STORE doc-2 failed: {}", to_string_lossy(&v));
// 4) Store two images via BYTES (local fake bytes; embedder only hashes bytes, not decoding)
let img1: Vec<u8> = b"local-image-bytes-1-abcdefghijklmnopqrstuvwxyz".to_vec();
let img2: Vec<u8> = b"local-image-bytes-2-ABCDEFGHIJKLMNOPQRSTUVWXYZ".to_vec();
let img1_b64 = base64::engine::general_purpose::STANDARD.encode(&img1);
let img2_b64 = base64::engine::general_purpose::STANDARD.encode(&img2);
let v = redis::cmd("LANCE.STOREIMAGE")
.arg("imageset")
.arg("ID")
.arg("img-1")
.arg("BYTES")
.arg(&img1_b64)
.arg("META")
.arg("title")
.arg("Local1")
.arg("group")
.arg("demo")
.query::<Value>(&mut conn)
.unwrap();
assert!(value_is_ok(&v), "LANCE.STOREIMAGE img-1 failed: {}", to_string_lossy(&v));
let v = redis::cmd("LANCE.STOREIMAGE")
.arg("imageset")
.arg("ID")
.arg("img-2")
.arg("BYTES")
.arg(&img2_b64)
.arg("META")
.arg("title")
.arg("Local2")
.arg("group")
.arg("demo")
.query::<Value>(&mut conn)
.unwrap();
assert!(value_is_ok(&v), "LANCE.STOREIMAGE img-2 failed: {}", to_string_lossy(&v));
// 5) Search text: K 2 QUERY "quick brown fox" RETURN 1 title
let v = redis::cmd("LANCE.SEARCH")
.arg("textset")
.arg("K")
.arg(2)
.arg("QUERY")
.arg("quick brown fox")
.arg("RETURN")
.arg(1)
.arg("title")
.query::<Value>(&mut conn)
.unwrap();
// Should be an array of hits
let ids = extract_hit_ids(&v);
assert!(
ids.contains(&"doc-1".to_string()) || ids.contains(&"doc-2".to_string()),
"LANCE.SEARCH should return doc-1/doc-2; got: {}",
to_string_lossy(&v)
);
// With FILTER on category
let v = redis::cmd("LANCE.SEARCH")
.arg("textset")
.arg("K")
.arg(2)
.arg("QUERY")
.arg("fox jumps")
.arg("FILTER")
.arg("category = 'animal'")
.arg("RETURN")
.arg(1)
.arg("title")
.query::<Value>(&mut conn)
.unwrap();
let ids_f = extract_hit_ids(&v);
assert!(
!ids_f.is_empty(),
"Filtered LANCE.SEARCH should return at least one document; got: {}",
to_string_lossy(&v)
);
// 6) Search images with QUERYBYTES
let query_img: Vec<u8> = b"local-image-query-3-1234567890".to_vec();
let query_img_b64 = base64::engine::general_purpose::STANDARD.encode(&query_img);
let v = redis::cmd("LANCE.SEARCHIMAGE")
.arg("imageset")
.arg("K")
.arg(2)
.arg("QUERYBYTES")
.arg(&query_img_b64)
.arg("RETURN")
.arg(1)
.arg("title")
.query::<Value>(&mut conn)
.unwrap();
// Should get 2 hits (img-1 and img-2) in some order; assert array non-empty
let img_ids = extract_hit_ids(&v);
assert!(
!img_ids.is_empty(),
"LANCE.SEARCHIMAGE should return non-empty results; got: {}",
to_string_lossy(&v)
);
// 7) Inspect datasets
let v = redis::cmd("LANCE.LIST").query::<Value>(&mut conn).unwrap();
assert!(
bulk_contains_string(&v, "textset"),
"LANCE.LIST missing textset: {}",
to_string_lossy(&v)
);
assert!(
bulk_contains_string(&v, "imageset"),
"LANCE.LIST missing imageset: {}",
to_string_lossy(&v)
);
// INFO textset
let info_text = redis::cmd("LANCE.INFO")
.arg("textset")
.query::<Value>(&mut conn)
.unwrap();
// INFO returns Array [k,v,k,v,...] including "dimension" "64" and "row_count" "...".
let info_str = to_string_lossy(&info_text);
assert!(
info_str.contains("dimension") && info_str.contains("64"),
"LANCE.INFO textset should include dimension 64; got: {}",
info_str
);
// 8) Delete by id and drop datasets
let v = redis::cmd("LANCE.DEL")
.arg("textset")
.arg("doc-2")
.query::<Value>(&mut conn)
.unwrap();
// Returns SimpleString "1" or Int 1 depending on encoding path; accept either
assert!(
value_is_int_eq(&v, 1) || value_is_str_eq(&v, "1"),
"LANCE.DEL doc-2 expected 1; got {}",
to_string_lossy(&v)
);
let v = redis::cmd("LANCE.DROP")
.arg("textset")
.query::<Value>(&mut conn)
.unwrap();
assert!(value_is_ok(&v), "LANCE.DROP textset failed: {}", to_string_lossy(&v));
let v = redis::cmd("LANCE.DROP")
.arg("imageset")
.query::<Value>(&mut conn)
.unwrap();
assert!(value_is_ok(&v), "LANCE.DROP imageset failed: {}", to_string_lossy(&v));
}