Compare commits
50
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9e6002c4d4 | ||
|
|
3bc82833ac | ||
|
|
fd4b6836ae | ||
|
|
59e72fa566 | ||
|
|
2592f8fcb3 | ||
|
|
17a0259cd0 | ||
|
|
7c1222d240 | ||
|
|
0d037cfac1 | ||
|
|
4e1c400f5d | ||
|
|
4d5c49090c | ||
|
|
a293502265 | ||
|
|
d87f2ccac6 | ||
|
|
09fd446c8d | ||
|
|
24f1249a62 | ||
|
|
650f5e97b6 | ||
|
|
97a376482c | ||
|
|
76f6fba5de | ||
|
|
dbe00d018f | ||
|
|
84407eeb51 | ||
|
|
037555802e | ||
|
|
fdb2e9cda8 | ||
|
|
238cb84d26 | ||
|
|
ed17216005 | ||
|
|
0bd753294b | ||
|
|
34c03894ae | ||
|
|
ddcbb9590c | ||
|
|
2c5aeaeada | ||
|
|
399fc81ab5 | ||
|
|
9992b5f135 | ||
|
|
294f3f6a49 | ||
|
|
04770c97eb | ||
|
|
8282842eaf | ||
|
|
35f0559d78 | ||
|
|
d61096e26f | ||
|
|
e5b7663c20 | ||
|
|
b2cd54e941 | ||
|
|
6f0f92a9e4 | ||
|
|
6cd1daeb51 | ||
|
|
6d49f81ebf | ||
|
|
9e53caaec3 | ||
|
|
bf7cf96911 | ||
|
|
3b695cd592 | ||
|
|
6a532cb248 | ||
|
|
354c9c4a04 | ||
|
|
0b9da45647 | ||
|
|
6df904a03c | ||
|
|
eae51977ab | ||
|
|
26fc98c885 | ||
|
|
bb05c4412e | ||
|
|
6740a017c7 |
Generated
+3453
-364
File diff suppressed because it is too large
Load Diff
+11
@@ -7,8 +7,13 @@ members = [
|
|||||||
"crates/catalogd",
|
"crates/catalogd",
|
||||||
"crates/queryd",
|
"crates/queryd",
|
||||||
"crates/aibridge",
|
"crates/aibridge",
|
||||||
|
"crates/ingestd",
|
||||||
|
"crates/vectord",
|
||||||
|
"crates/journald",
|
||||||
"crates/gateway",
|
"crates/gateway",
|
||||||
"crates/ui",
|
"crates/ui",
|
||||||
|
"crates/lance-bench",
|
||||||
|
"crates/vectord-lance",
|
||||||
]
|
]
|
||||||
|
|
||||||
[workspace.dependencies]
|
[workspace.dependencies]
|
||||||
@@ -38,3 +43,9 @@ opentelemetry_sdk = { version = "0.28", features = ["rt-tokio"] }
|
|||||||
opentelemetry-stdout = { version = "0.28", features = ["trace"] }
|
opentelemetry-stdout = { version = "0.28", features = ["trace"] }
|
||||||
tracing-opentelemetry = "0.29"
|
tracing-opentelemetry = "0.29"
|
||||||
toml = "0.8"
|
toml = "0.8"
|
||||||
|
csv = "1"
|
||||||
|
lopdf = "0.35"
|
||||||
|
encoding_rs = "0.8"
|
||||||
|
instant-distance = "0.6"
|
||||||
|
tokio-postgres = { version = "0.7", features = ["with-serde_json-1", "with-chrono-0_4", "with-uuid-1"] }
|
||||||
|
mysql_async = { version = "0.34", default-features = false, features = ["minimal"] }
|
||||||
|
|||||||
@@ -134,4 +134,47 @@ impl AiClient {
|
|||||||
}
|
}
|
||||||
resp.json().await.map_err(|e| format!("rerank parse error: {e}"))
|
resp.json().await.map_err(|e| format!("rerank parse error: {e}"))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Force Ollama to unload the named model from VRAM (keep_alive=0).
|
||||||
|
/// Used for predictable profile swaps — without this, Ollama holds a
|
||||||
|
/// model for its configured TTL (default 5min) and the previous
|
||||||
|
/// profile's model can linger in VRAM next to the new one.
|
||||||
|
pub async fn unload_model(&self, model: &str) -> Result<serde_json::Value, String> {
|
||||||
|
let resp = self.client
|
||||||
|
.post(format!("{}/admin/unload", self.base_url))
|
||||||
|
.json(&serde_json::json!({ "model": model }))
|
||||||
|
.send().await
|
||||||
|
.map_err(|e| format!("unload request failed: {e}"))?;
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
let text = resp.text().await.unwrap_or_default();
|
||||||
|
return Err(format!("unload error: {text}"));
|
||||||
|
}
|
||||||
|
resp.json().await.map_err(|e| format!("unload parse error: {e}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Ask Ollama to load the named model into VRAM proactively. Makes
|
||||||
|
/// the first real request after profile activation fast (no cold-load
|
||||||
|
/// latency).
|
||||||
|
pub async fn preload_model(&self, model: &str) -> Result<serde_json::Value, String> {
|
||||||
|
let resp = self.client
|
||||||
|
.post(format!("{}/admin/preload", self.base_url))
|
||||||
|
.json(&serde_json::json!({ "model": model }))
|
||||||
|
.send().await
|
||||||
|
.map_err(|e| format!("preload request failed: {e}"))?;
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
let text = resp.text().await.unwrap_or_default();
|
||||||
|
return Err(format!("preload error: {text}"));
|
||||||
|
}
|
||||||
|
resp.json().await.map_err(|e| format!("preload parse error: {e}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GPU + loaded-model snapshot from the sidecar. Combines nvidia-smi
|
||||||
|
/// output (if available) with Ollama's /api/ps.
|
||||||
|
pub async fn vram_snapshot(&self) -> Result<serde_json::Value, String> {
|
||||||
|
let resp = self.client
|
||||||
|
.get(format!("{}/admin/vram", self.base_url))
|
||||||
|
.send().await
|
||||||
|
.map_err(|e| format!("vram request failed: {e}"))?;
|
||||||
|
resp.json().await.map_err(|e| format!("vram parse error: {e}"))
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -16,9 +16,17 @@ pub fn router(client: AiClient) -> Router {
|
|||||||
.route("/embed", post(embed))
|
.route("/embed", post(embed))
|
||||||
.route("/generate", post(generate))
|
.route("/generate", post(generate))
|
||||||
.route("/rerank", post(rerank))
|
.route("/rerank", post(rerank))
|
||||||
|
.route("/vram", get(vram))
|
||||||
.with_state(client)
|
.with_state(client)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn vram(State(client): State<AiClient>) -> impl IntoResponse {
|
||||||
|
match client.vram_snapshot().await {
|
||||||
|
Ok(snap) => Ok(Json(snap)),
|
||||||
|
Err(e) => Err((StatusCode::BAD_GATEWAY, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
async fn health(State(client): State<AiClient>) -> impl IntoResponse {
|
async fn health(State(client): State<AiClient>) -> impl IntoResponse {
|
||||||
match client.health().await {
|
match client.health().await {
|
||||||
Ok(info) => Ok(Json(info)),
|
Ok(info) => Ok(Json(info)),
|
||||||
|
|||||||
@@ -15,5 +15,6 @@ bytes = { workspace = true }
|
|||||||
chrono = { workspace = true }
|
chrono = { workspace = true }
|
||||||
uuid = { workspace = true }
|
uuid = { workspace = true }
|
||||||
object_store = { workspace = true }
|
object_store = { workspace = true }
|
||||||
|
arrow = { workspace = true }
|
||||||
proto = { path = "../proto" }
|
proto = { path = "../proto" }
|
||||||
tonic = { workspace = true }
|
tonic = { workspace = true }
|
||||||
|
|||||||
@@ -1,3 +1,4 @@
|
|||||||
pub mod registry;
|
pub mod registry;
|
||||||
pub mod service;
|
pub mod service;
|
||||||
pub mod grpc;
|
pub mod grpc;
|
||||||
|
pub mod tombstones;
|
||||||
|
|||||||
@@ -1,4 +1,9 @@
|
|||||||
use shared::types::{DatasetId, DatasetManifest, ObjectRef, SchemaFingerprint};
|
use shared::types::{
|
||||||
|
AiView, ColumnMeta, DatasetId, DatasetManifest, FreshnessContract, Lineage, ModelProfile,
|
||||||
|
ObjectRef, RefreshPolicy, SchemaFingerprint, Sensitivity, Tombstone,
|
||||||
|
};
|
||||||
|
|
||||||
|
use crate::tombstones::TombstoneStore;
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tokio::sync::RwLock;
|
use tokio::sync::RwLock;
|
||||||
@@ -6,12 +11,51 @@ use tokio::sync::RwLock;
|
|||||||
use storaged::ops;
|
use storaged::ops;
|
||||||
use object_store::ObjectStore;
|
use object_store::ObjectStore;
|
||||||
|
|
||||||
|
/// Make a view name safe for use as an object storage key.
|
||||||
|
/// Allows letters, digits, `_`, `-`, `.`. Anything else becomes `_`.
|
||||||
|
fn sanitize_view_name(name: &str) -> String {
|
||||||
|
name.chars()
|
||||||
|
.map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.' { c } else { '_' })
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Default, serde::Serialize)]
|
||||||
|
pub struct MigrateBucketsReport {
|
||||||
|
pub refs_examined: usize,
|
||||||
|
pub refs_renamed: usize, // legacy "data"/"local" → "primary"
|
||||||
|
pub refs_stamped: usize, // empty → "primary"
|
||||||
|
pub refs_unchanged: usize, // already canonical
|
||||||
|
pub manifests_persisted: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Partial metadata update — only set fields are applied.
|
||||||
|
#[derive(Debug, Clone, Default, serde::Deserialize)]
|
||||||
|
pub struct MetadataUpdate {
|
||||||
|
pub description: Option<String>,
|
||||||
|
pub owner: Option<String>,
|
||||||
|
pub sensitivity: Option<Sensitivity>,
|
||||||
|
pub tags: Option<Vec<String>>,
|
||||||
|
pub columns: Option<Vec<ColumnMeta>>,
|
||||||
|
pub lineage: Option<Lineage>,
|
||||||
|
pub freshness: Option<FreshnessContract>,
|
||||||
|
pub row_count: Option<u64>,
|
||||||
|
// Phase C embedding freshness
|
||||||
|
pub embedding_refresh_policy: Option<RefreshPolicy>,
|
||||||
|
}
|
||||||
|
|
||||||
const MANIFEST_PREFIX: &str = "_catalog/manifests";
|
const MANIFEST_PREFIX: &str = "_catalog/manifests";
|
||||||
|
const VIEW_PREFIX: &str = "_catalog/views";
|
||||||
|
const PROFILE_PREFIX: &str = "_catalog/profiles";
|
||||||
|
|
||||||
/// In-memory dataset registry backed by manifest persistence in object storage.
|
/// In-memory dataset registry backed by manifest persistence in object storage.
|
||||||
|
/// Also tracks AiViews (Phase D) — safe projections over base datasets.
|
||||||
|
/// And tombstones (Phase E) — soft-delete markers applied at query time.
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
pub struct Registry {
|
pub struct Registry {
|
||||||
datasets: Arc<RwLock<HashMap<DatasetId, DatasetManifest>>>,
|
datasets: Arc<RwLock<HashMap<DatasetId, DatasetManifest>>>,
|
||||||
|
views: Arc<RwLock<HashMap<String, AiView>>>,
|
||||||
|
profiles: Arc<RwLock<HashMap<String, ModelProfile>>>,
|
||||||
|
tombstones: TombstoneStore,
|
||||||
store: Arc<dyn ObjectStore>,
|
store: Arc<dyn ObjectStore>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -19,11 +63,47 @@ impl Registry {
|
|||||||
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||||
Self {
|
Self {
|
||||||
datasets: Arc::new(RwLock::new(HashMap::new())),
|
datasets: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
views: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
profiles: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
tombstones: TombstoneStore::new(store.clone()),
|
||||||
store,
|
store,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Rebuild in-memory index from persisted manifests on startup.
|
pub fn tombstones(&self) -> &TombstoneStore {
|
||||||
|
&self.tombstones
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Add a tombstone for a dataset row. Validates that the dataset
|
||||||
|
/// exists. Idempotent at the endpoint layer — callers dedup if needed.
|
||||||
|
pub async fn add_tombstone(
|
||||||
|
&self,
|
||||||
|
dataset: &str,
|
||||||
|
row_key_column: &str,
|
||||||
|
row_key_value: &str,
|
||||||
|
actor: &str,
|
||||||
|
reason: &str,
|
||||||
|
) -> Result<Tombstone, String> {
|
||||||
|
if self.get_by_name(dataset).await.is_none() {
|
||||||
|
return Err(format!("dataset not found: {dataset}"));
|
||||||
|
}
|
||||||
|
let ts = Tombstone {
|
||||||
|
dataset: dataset.to_string(),
|
||||||
|
row_key_column: row_key_column.to_string(),
|
||||||
|
row_key_value: row_key_value.to_string(),
|
||||||
|
deleted_at: chrono::Utc::now(),
|
||||||
|
actor: actor.to_string(),
|
||||||
|
reason: reason.to_string(),
|
||||||
|
};
|
||||||
|
self.tombstones.append(&ts).await?;
|
||||||
|
Ok(ts)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_tombstones(&self, dataset: &str) -> Result<Vec<Tombstone>, String> {
|
||||||
|
self.tombstones.list(dataset).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Rebuild in-memory index from persisted manifests + views on startup.
|
||||||
pub async fn rebuild(&self) -> Result<usize, String> {
|
pub async fn rebuild(&self) -> Result<usize, String> {
|
||||||
let keys = ops::list(&self.store, Some(MANIFEST_PREFIX)).await?;
|
let keys = ops::list(&self.store, Some(MANIFEST_PREFIX)).await?;
|
||||||
let mut datasets = self.datasets.write().await;
|
let mut datasets = self.datasets.write().await;
|
||||||
@@ -36,9 +116,102 @@ impl Registry {
|
|||||||
}
|
}
|
||||||
let count = datasets.len();
|
let count = datasets.len();
|
||||||
tracing::info!("catalog rebuilt: {count} datasets loaded");
|
tracing::info!("catalog rebuilt: {count} datasets loaded");
|
||||||
|
|
||||||
|
// Phase D: load AiView definitions alongside manifests.
|
||||||
|
let view_keys = ops::list(&self.store, Some(VIEW_PREFIX)).await.unwrap_or_default();
|
||||||
|
let mut views = self.views.write().await;
|
||||||
|
views.clear();
|
||||||
|
for key in &view_keys {
|
||||||
|
if !key.ends_with(".json") { continue; }
|
||||||
|
let data = match ops::get(&self.store, key).await {
|
||||||
|
Ok(d) => d,
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!("view '{key}': read failed: {e}");
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
match serde_json::from_slice::<AiView>(&data) {
|
||||||
|
Ok(view) => { views.insert(view.name.clone(), view); }
|
||||||
|
Err(e) => tracing::warn!("view '{key}': parse failed: {e}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !views.is_empty() {
|
||||||
|
tracing::info!("catalog: {} views loaded", views.len());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Phase 17: load model profiles.
|
||||||
|
let profile_keys = ops::list(&self.store, Some(PROFILE_PREFIX)).await.unwrap_or_default();
|
||||||
|
let mut profiles = self.profiles.write().await;
|
||||||
|
profiles.clear();
|
||||||
|
for key in &profile_keys {
|
||||||
|
if !key.ends_with(".json") { continue; }
|
||||||
|
let data = match ops::get(&self.store, key).await {
|
||||||
|
Ok(d) => d,
|
||||||
|
Err(e) => { tracing::warn!("profile '{key}': read failed: {e}"); continue; }
|
||||||
|
};
|
||||||
|
match serde_json::from_slice::<ModelProfile>(&data) {
|
||||||
|
Ok(p) => { profiles.insert(p.id.clone(), p); }
|
||||||
|
Err(e) => tracing::warn!("profile '{key}': parse failed: {e}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !profiles.is_empty() {
|
||||||
|
tracing::info!("catalog: {} model profiles loaded", profiles.len());
|
||||||
|
}
|
||||||
|
|
||||||
Ok(count)
|
Ok(count)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// --- Phase 17: Model profiles ---
|
||||||
|
|
||||||
|
/// Create or replace a model profile. Validates id slug + ensures
|
||||||
|
/// every bound_dataset exists (as either a raw dataset or an AiView).
|
||||||
|
pub async fn put_profile(&self, profile: ModelProfile) -> Result<ModelProfile, String> {
|
||||||
|
if profile.id.is_empty() {
|
||||||
|
return Err("profile id is empty".into());
|
||||||
|
}
|
||||||
|
if !profile.id.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_') {
|
||||||
|
return Err(format!(
|
||||||
|
"profile id '{}' must be alphanumeric + '-' or '_' only", profile.id,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
for binding in &profile.bound_datasets {
|
||||||
|
let exists = self.get_by_name(binding).await.is_some()
|
||||||
|
|| self.get_view(binding).await.is_some();
|
||||||
|
if !exists {
|
||||||
|
return Err(format!(
|
||||||
|
"bound dataset '{}' not found as dataset or view", binding,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let key = format!("{PROFILE_PREFIX}/{}.json", profile.id);
|
||||||
|
let json = serde_json::to_vec_pretty(&profile).map_err(|e| e.to_string())?;
|
||||||
|
ops::put(&self.store, &key, json.into()).await?;
|
||||||
|
|
||||||
|
let mut profiles = self.profiles.write().await;
|
||||||
|
profiles.insert(profile.id.clone(), profile.clone());
|
||||||
|
tracing::info!(
|
||||||
|
"profile registered: {} -> ollama={} bindings={:?}",
|
||||||
|
profile.id, profile.ollama_name, profile.bound_datasets,
|
||||||
|
);
|
||||||
|
Ok(profile)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_profile(&self, id: &str) -> Option<ModelProfile> {
|
||||||
|
self.profiles.read().await.get(id).cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_profiles(&self) -> Vec<ModelProfile> {
|
||||||
|
self.profiles.read().await.values().cloned().collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn delete_profile(&self, id: &str) -> Result<(), String> {
|
||||||
|
let key = format!("{PROFILE_PREFIX}/{id}.json");
|
||||||
|
ops::delete(&self.store, &key).await?;
|
||||||
|
self.profiles.write().await.remove(id);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
/// Register a new dataset. Persists manifest to storage before updating memory.
|
/// Register a new dataset. Persists manifest to storage before updating memory.
|
||||||
pub async fn register(
|
pub async fn register(
|
||||||
&self,
|
&self,
|
||||||
@@ -54,6 +227,17 @@ impl Registry {
|
|||||||
objects,
|
objects,
|
||||||
created_at: now,
|
created_at: now,
|
||||||
updated_at: now,
|
updated_at: now,
|
||||||
|
description: String::new(),
|
||||||
|
owner: String::new(),
|
||||||
|
sensitivity: None,
|
||||||
|
columns: vec![],
|
||||||
|
lineage: None,
|
||||||
|
freshness: None,
|
||||||
|
tags: vec![],
|
||||||
|
row_count: None,
|
||||||
|
last_embedded_at: None,
|
||||||
|
embedding_stale_since: None,
|
||||||
|
embedding_refresh_policy: None,
|
||||||
};
|
};
|
||||||
|
|
||||||
// Write-ahead: persist before in-memory update
|
// Write-ahead: persist before in-memory update
|
||||||
@@ -68,6 +252,37 @@ impl Registry {
|
|||||||
Ok(manifest)
|
Ok(manifest)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Update metadata on an existing dataset (owner, description, tags, sensitivity, etc.)
|
||||||
|
pub async fn update_metadata(
|
||||||
|
&self,
|
||||||
|
name: &str,
|
||||||
|
updates: MetadataUpdate,
|
||||||
|
) -> Result<DatasetManifest, String> {
|
||||||
|
let mut datasets = self.datasets.write().await;
|
||||||
|
let manifest = datasets.values_mut()
|
||||||
|
.find(|d| d.name == name)
|
||||||
|
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||||
|
|
||||||
|
if let Some(desc) = updates.description { manifest.description = desc; }
|
||||||
|
if let Some(owner) = updates.owner { manifest.owner = owner; }
|
||||||
|
if let Some(sens) = updates.sensitivity { manifest.sensitivity = Some(sens); }
|
||||||
|
if let Some(tags) = updates.tags { manifest.tags = tags; }
|
||||||
|
if let Some(cols) = updates.columns { manifest.columns = cols; }
|
||||||
|
if let Some(lineage) = updates.lineage { manifest.lineage = Some(lineage); }
|
||||||
|
if let Some(freshness) = updates.freshness { manifest.freshness = Some(freshness); }
|
||||||
|
if let Some(count) = updates.row_count { manifest.row_count = Some(count); }
|
||||||
|
if let Some(policy) = updates.embedding_refresh_policy { manifest.embedding_refresh_policy = Some(policy); }
|
||||||
|
manifest.updated_at = chrono::Utc::now();
|
||||||
|
|
||||||
|
// Persist
|
||||||
|
let manifest_key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
|
||||||
|
let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
|
||||||
|
ops::put(&self.store, &manifest_key, json.into()).await?;
|
||||||
|
|
||||||
|
let result = manifest.clone();
|
||||||
|
Ok(result)
|
||||||
|
}
|
||||||
|
|
||||||
/// Get a dataset by ID.
|
/// Get a dataset by ID.
|
||||||
pub async fn get(&self, id: &DatasetId) -> Option<DatasetManifest> {
|
pub async fn get(&self, id: &DatasetId) -> Option<DatasetManifest> {
|
||||||
let datasets = self.datasets.read().await;
|
let datasets = self.datasets.read().await;
|
||||||
@@ -86,6 +301,264 @@ impl Registry {
|
|||||||
datasets.values().cloned().collect()
|
datasets.values().cloned().collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Re-read the parquet footer(s) for a dataset and repopulate `row_count`
|
||||||
|
/// and `columns` from reality. Use this to repair manifests whose
|
||||||
|
/// metadata was lost (e.g. migrated from a pre-Phase 10 catalog).
|
||||||
|
///
|
||||||
|
/// Does NOT touch owner/description/sensitivity/lineage/tags — only
|
||||||
|
/// the structural facts that parquet can tell us authoritatively.
|
||||||
|
/// The existing `schema_fingerprint` is updated if the recomputed one
|
||||||
|
/// differs; a warning is logged so drift is visible.
|
||||||
|
pub async fn resync_from_parquet(&self, name: &str) -> Result<DatasetManifest, String> {
|
||||||
|
use shared::arrow_helpers::{fingerprint_schema, parquet_to_record_batches};
|
||||||
|
|
||||||
|
// Snapshot the target manifest so we don't hold the write lock during IO.
|
||||||
|
let (id, objects, old_fp) = {
|
||||||
|
let datasets = self.datasets.read().await;
|
||||||
|
let m = datasets
|
||||||
|
.values()
|
||||||
|
.find(|d| d.name == name)
|
||||||
|
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||||
|
(m.id.clone(), m.objects.clone(), m.schema_fingerprint.clone())
|
||||||
|
};
|
||||||
|
|
||||||
|
if objects.is_empty() {
|
||||||
|
return Err(format!("dataset '{name}' has no object references to resync from"));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut total_rows: u64 = 0;
|
||||||
|
let mut first_schema: Option<arrow::datatypes::SchemaRef> = None;
|
||||||
|
for obj in &objects {
|
||||||
|
let data = ops::get(&self.store, &obj.key).await
|
||||||
|
.map_err(|e| format!("read {}: {e}", obj.key))?;
|
||||||
|
let (schema, batches) = parquet_to_record_batches(&data)
|
||||||
|
.map_err(|e| format!("parse {}: {e}", obj.key))?;
|
||||||
|
let rows: u64 = batches.iter().map(|b| b.num_rows() as u64).sum();
|
||||||
|
total_rows += rows;
|
||||||
|
if first_schema.is_none() {
|
||||||
|
first_schema = Some(schema);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let schema = first_schema.ok_or("no schema recovered")?;
|
||||||
|
let new_fp = fingerprint_schema(&schema);
|
||||||
|
if new_fp != old_fp {
|
||||||
|
tracing::warn!(
|
||||||
|
"dataset '{}' schema fingerprint drift: {} -> {} (updating to match parquet reality)",
|
||||||
|
name, old_fp.0, new_fp.0,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let columns: Vec<ColumnMeta> = schema
|
||||||
|
.fields()
|
||||||
|
.iter()
|
||||||
|
.map(|f| ColumnMeta {
|
||||||
|
name: f.name().clone(),
|
||||||
|
data_type: f.data_type().to_string(),
|
||||||
|
sensitivity: None,
|
||||||
|
description: String::new(),
|
||||||
|
is_pii: false,
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// Apply updates.
|
||||||
|
let mut datasets = self.datasets.write().await;
|
||||||
|
let manifest = datasets
|
||||||
|
.get_mut(&id)
|
||||||
|
.ok_or_else(|| format!("dataset disappeared during resync: {name}"))?;
|
||||||
|
manifest.row_count = Some(total_rows);
|
||||||
|
manifest.columns = columns;
|
||||||
|
manifest.schema_fingerprint = new_fp;
|
||||||
|
manifest.updated_at = chrono::Utc::now();
|
||||||
|
|
||||||
|
// Persist.
|
||||||
|
let manifest_key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
|
||||||
|
let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
|
||||||
|
ops::put(&self.store, &manifest_key, json.into()).await?;
|
||||||
|
|
||||||
|
tracing::info!("resynced '{name}': row_count={total_rows}, {} columns", manifest.columns.len());
|
||||||
|
Ok(manifest.clone())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resync every dataset that currently has a null row_count.
|
||||||
|
/// Returns (successes, failures) where each entry is (name, detail).
|
||||||
|
pub async fn resync_missing(&self) -> (Vec<(String, u64)>, Vec<(String, String)>) {
|
||||||
|
let names: Vec<String> = {
|
||||||
|
let datasets = self.datasets.read().await;
|
||||||
|
datasets
|
||||||
|
.values()
|
||||||
|
.filter(|d| d.row_count.is_none() || d.columns.is_empty())
|
||||||
|
.map(|d| d.name.clone())
|
||||||
|
.collect()
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut ok = Vec::new();
|
||||||
|
let mut err = Vec::new();
|
||||||
|
for name in names {
|
||||||
|
match self.resync_from_parquet(&name).await {
|
||||||
|
Ok(m) => ok.push((name, m.row_count.unwrap_or(0))),
|
||||||
|
Err(e) => err.push((name, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
(ok, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mark a dataset's embeddings as stale (row-level data has been written
|
||||||
|
/// since the last embedding refresh). Idempotent — setting stale when
|
||||||
|
/// already stale is a no-op. Only marks stale if the dataset has been
|
||||||
|
/// embedded before — a never-embedded dataset doesn't need a stale flag
|
||||||
|
/// (it just needs an initial index build). Called from the ingest path.
|
||||||
|
pub async fn mark_embeddings_stale(&self, name: &str) -> Result<(), String> {
|
||||||
|
let mut datasets = self.datasets.write().await;
|
||||||
|
let manifest = datasets
|
||||||
|
.values_mut()
|
||||||
|
.find(|d| d.name == name)
|
||||||
|
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||||
|
|
||||||
|
if manifest.last_embedded_at.is_none() {
|
||||||
|
return Ok(()); // never embedded -> no stale semantics yet
|
||||||
|
}
|
||||||
|
if manifest.embedding_stale_since.is_none() {
|
||||||
|
manifest.embedding_stale_since = Some(chrono::Utc::now());
|
||||||
|
manifest.updated_at = chrono::Utc::now();
|
||||||
|
|
||||||
|
let key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
|
||||||
|
let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
|
||||||
|
ops::put(&self.store, &key, json.into()).await?;
|
||||||
|
tracing::info!("marked embeddings stale for dataset '{name}'");
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Clear the stale marker and set `last_embedded_at = now`.
|
||||||
|
/// Called by the embedding refresh pipeline once it finishes.
|
||||||
|
pub async fn clear_embeddings_stale(&self, name: &str) -> Result<(), String> {
|
||||||
|
let mut datasets = self.datasets.write().await;
|
||||||
|
let manifest = datasets
|
||||||
|
.values_mut()
|
||||||
|
.find(|d| d.name == name)
|
||||||
|
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||||
|
|
||||||
|
let now = chrono::Utc::now();
|
||||||
|
manifest.embedding_stale_since = None;
|
||||||
|
manifest.last_embedded_at = Some(now);
|
||||||
|
manifest.updated_at = now;
|
||||||
|
|
||||||
|
let key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
|
||||||
|
let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
|
||||||
|
ops::put(&self.store, &key, json.into()).await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Federation layer 2: stamp `bucket = "primary"` on every ObjectRef
|
||||||
|
/// whose `bucket` field is empty or matches a legacy value (`"data"`,
|
||||||
|
/// `"local"`). One-shot migration; re-running is a safe no-op once
|
||||||
|
/// every ref is canonical.
|
||||||
|
pub async fn migrate_buckets_to_primary(&self) -> Result<MigrateBucketsReport, String> {
|
||||||
|
let mut datasets = self.datasets.write().await;
|
||||||
|
let mut report = MigrateBucketsReport::default();
|
||||||
|
let mut to_persist: Vec<(DatasetId, DatasetManifest)> = Vec::new();
|
||||||
|
|
||||||
|
for manifest in datasets.values_mut() {
|
||||||
|
let mut changed = false;
|
||||||
|
for obj in manifest.objects.iter_mut() {
|
||||||
|
report.refs_examined += 1;
|
||||||
|
if obj.bucket.is_empty() {
|
||||||
|
obj.bucket = "primary".to_string();
|
||||||
|
report.refs_stamped += 1;
|
||||||
|
changed = true;
|
||||||
|
} else if obj.bucket == "data" || obj.bucket == "local" {
|
||||||
|
obj.bucket = "primary".to_string();
|
||||||
|
report.refs_renamed += 1;
|
||||||
|
changed = true;
|
||||||
|
} else {
|
||||||
|
report.refs_unchanged += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if changed {
|
||||||
|
manifest.updated_at = chrono::Utc::now();
|
||||||
|
to_persist.push((manifest.id.clone(), manifest.clone()));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Persist updated manifests after we've finished mutating the map.
|
||||||
|
for (id, manifest) in to_persist {
|
||||||
|
let key = format!("{MANIFEST_PREFIX}/{}.json", id);
|
||||||
|
let json = serde_json::to_vec_pretty(&manifest).map_err(|e| e.to_string())?;
|
||||||
|
ops::put(&self.store, &key, json.into()).await?;
|
||||||
|
report.manifests_persisted += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
"bucket migration: examined {} refs, renamed {}, stamped {}, unchanged {}, persisted {} manifests",
|
||||||
|
report.refs_examined,
|
||||||
|
report.refs_renamed,
|
||||||
|
report.refs_stamped,
|
||||||
|
report.refs_unchanged,
|
||||||
|
report.manifests_persisted,
|
||||||
|
);
|
||||||
|
Ok(report)
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Phase D: AI-safe views ---
|
||||||
|
|
||||||
|
/// Create or replace a named view. Validates that the base dataset
|
||||||
|
/// exists and that the column whitelist is non-empty. Persists to
|
||||||
|
/// `_catalog/views/{name}.json`.
|
||||||
|
pub async fn put_view(&self, mut view: AiView) -> Result<AiView, String> {
|
||||||
|
if view.name.is_empty() {
|
||||||
|
return Err("view name is empty".into());
|
||||||
|
}
|
||||||
|
if view.columns.is_empty() {
|
||||||
|
return Err("view must whitelist at least one column".into());
|
||||||
|
}
|
||||||
|
// Base dataset must exist (read-side will fail anyway, fail-fast here
|
||||||
|
// so operators don't end up with dangling views).
|
||||||
|
if self.get_by_name(&view.base_dataset).await.is_none() {
|
||||||
|
return Err(format!(
|
||||||
|
"base dataset '{}' not found in catalog",
|
||||||
|
view.base_dataset
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if view.created_at.timestamp() == 0 {
|
||||||
|
view.created_at = chrono::Utc::now();
|
||||||
|
}
|
||||||
|
|
||||||
|
let key = format!("{VIEW_PREFIX}/{}.json", sanitize_view_name(&view.name));
|
||||||
|
let json = serde_json::to_vec_pretty(&view).map_err(|e| e.to_string())?;
|
||||||
|
ops::put(&self.store, &key, json.into()).await?;
|
||||||
|
|
||||||
|
let mut views = self.views.write().await;
|
||||||
|
views.insert(view.name.clone(), view.clone());
|
||||||
|
tracing::info!("view registered: {} over '{}'", view.name, view.base_dataset);
|
||||||
|
Ok(view)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_view(&self, name: &str) -> Option<AiView> {
|
||||||
|
self.views.read().await.get(name).cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_views(&self) -> Vec<AiView> {
|
||||||
|
self.views.read().await.values().cloned().collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn delete_view(&self, name: &str) -> Result<(), String> {
|
||||||
|
let key = format!("{VIEW_PREFIX}/{}.json", sanitize_view_name(name));
|
||||||
|
ops::delete(&self.store, &key).await?;
|
||||||
|
self.views.write().await.remove(name);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List datasets whose `embedding_stale_since` is set — they need a refresh.
|
||||||
|
pub async fn stale_datasets(&self) -> Vec<DatasetManifest> {
|
||||||
|
let datasets = self.datasets.read().await;
|
||||||
|
datasets
|
||||||
|
.values()
|
||||||
|
.filter(|d| d.embedding_stale_since.is_some())
|
||||||
|
.cloned()
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
/// Add objects to an existing dataset.
|
/// Add objects to an existing dataset.
|
||||||
pub async fn add_objects(
|
pub async fn add_objects(
|
||||||
&self,
|
&self,
|
||||||
|
|||||||
@@ -18,6 +18,18 @@ pub fn router(registry: Registry) -> Router {
|
|||||||
.route("/datasets", get(list_datasets))
|
.route("/datasets", get(list_datasets))
|
||||||
.route("/datasets/{id}", get(get_dataset))
|
.route("/datasets/{id}", get(get_dataset))
|
||||||
.route("/datasets/by-name/{name}", get(get_dataset_by_name))
|
.route("/datasets/by-name/{name}", get(get_dataset_by_name))
|
||||||
|
.route("/datasets/by-name/{name}/metadata", post(update_metadata))
|
||||||
|
.route("/datasets/by-name/{name}/resync", post(resync_dataset))
|
||||||
|
.route("/resync-missing", post(resync_all_missing))
|
||||||
|
.route("/migrate-buckets", post(migrate_buckets))
|
||||||
|
// Phase D: AI-safe views
|
||||||
|
.route("/views", post(create_view).get(list_views))
|
||||||
|
.route("/views/{name}", get(get_view).delete(delete_view))
|
||||||
|
// Phase E: soft-delete tombstones
|
||||||
|
.route("/datasets/by-name/{name}/tombstone", post(tombstone_rows).get(list_tombstones))
|
||||||
|
// Phase 17: model profiles
|
||||||
|
.route("/profiles", post(create_profile).get(list_profiles))
|
||||||
|
.route("/profiles/{id}", get(get_profile).delete(delete_profile))
|
||||||
.with_state(registry)
|
.with_state(registry)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -47,6 +59,15 @@ struct DatasetResponse {
|
|||||||
objects: Vec<ObjectRefResponse>,
|
objects: Vec<ObjectRefResponse>,
|
||||||
created_at: String,
|
created_at: String,
|
||||||
updated_at: String,
|
updated_at: String,
|
||||||
|
// Rich metadata
|
||||||
|
description: String,
|
||||||
|
owner: String,
|
||||||
|
sensitivity: Option<shared::types::Sensitivity>,
|
||||||
|
columns: Vec<shared::types::ColumnMeta>,
|
||||||
|
lineage: Option<shared::types::Lineage>,
|
||||||
|
freshness: Option<shared::types::FreshnessContract>,
|
||||||
|
tags: Vec<String>,
|
||||||
|
row_count: Option<u64>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Serialize)]
|
#[derive(Serialize)]
|
||||||
@@ -71,6 +92,14 @@ impl From<&shared::types::DatasetManifest> for DatasetResponse {
|
|||||||
}).collect(),
|
}).collect(),
|
||||||
created_at: m.created_at.to_rfc3339(),
|
created_at: m.created_at.to_rfc3339(),
|
||||||
updated_at: m.updated_at.to_rfc3339(),
|
updated_at: m.updated_at.to_rfc3339(),
|
||||||
|
description: m.description.clone(),
|
||||||
|
owner: m.owner.clone(),
|
||||||
|
sensitivity: m.sensitivity.clone(),
|
||||||
|
columns: m.columns.clone(),
|
||||||
|
lineage: m.lineage.clone(),
|
||||||
|
freshness: m.freshness.clone(),
|
||||||
|
tags: m.tags.clone(),
|
||||||
|
row_count: m.row_count,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -123,3 +152,265 @@ async fn get_dataset_by_name(
|
|||||||
None => Err((StatusCode::NOT_FOUND, format!("dataset not found: {name}"))),
|
None => Err((StatusCode::NOT_FOUND, format!("dataset not found: {name}"))),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn update_metadata(
|
||||||
|
State(registry): State<Registry>,
|
||||||
|
Path(name): Path<String>,
|
||||||
|
Json(updates): Json<crate::registry::MetadataUpdate>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match registry.update_metadata(&name, updates).await {
|
||||||
|
Ok(manifest) => Ok(Json(DatasetResponse::from(&manifest))),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Re-read parquet footers for a single dataset and repopulate row_count
|
||||||
|
/// and columns from reality. Useful for repairing manifests whose metadata
|
||||||
|
/// was lost or never backfilled.
|
||||||
|
async fn resync_dataset(
|
||||||
|
State(registry): State<Registry>,
|
||||||
|
Path(name): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match registry.resync_from_parquet(&name).await {
|
||||||
|
Ok(manifest) => Ok(Json(DatasetResponse::from(&manifest))),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize)]
|
||||||
|
struct ResyncAllResponse {
|
||||||
|
succeeded: Vec<ResyncOk>,
|
||||||
|
failed: Vec<ResyncErr>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize)]
|
||||||
|
struct ResyncOk {
|
||||||
|
name: String,
|
||||||
|
row_count: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize)]
|
||||||
|
struct ResyncErr {
|
||||||
|
name: String,
|
||||||
|
error: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resync every dataset that currently has null row_count or empty columns.
|
||||||
|
async fn resync_all_missing(State(registry): State<Registry>) -> impl IntoResponse {
|
||||||
|
let (ok, err) = registry.resync_missing().await;
|
||||||
|
Json(ResyncAllResponse {
|
||||||
|
succeeded: ok.into_iter().map(|(name, row_count)| ResyncOk { name, row_count }).collect(),
|
||||||
|
failed: err.into_iter().map(|(name, error)| ResyncErr { name, error }).collect(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Federation layer 2 one-shot: normalize every ObjectRef.bucket field
|
||||||
|
/// to the canonical "primary" value. Idempotent — re-running once
|
||||||
|
/// everything is canonical is a safe no-op.
|
||||||
|
async fn migrate_buckets(State(registry): State<Registry>) -> impl IntoResponse {
|
||||||
|
match registry.migrate_buckets_to_primary().await {
|
||||||
|
Ok(report) => Ok(Json(report)),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Phase D: AI-safe views ---
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct CreateViewRequest {
|
||||||
|
name: String,
|
||||||
|
base_dataset: String,
|
||||||
|
columns: Vec<String>,
|
||||||
|
#[serde(default)]
|
||||||
|
row_filter: Option<String>,
|
||||||
|
#[serde(default)]
|
||||||
|
column_redactions: std::collections::HashMap<String, shared::types::Redaction>,
|
||||||
|
#[serde(default)]
|
||||||
|
description: String,
|
||||||
|
#[serde(default)]
|
||||||
|
created_by: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn create_view(
|
||||||
|
State(registry): State<Registry>,
|
||||||
|
Json(req): Json<CreateViewRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let view = shared::types::AiView {
|
||||||
|
name: req.name,
|
||||||
|
base_dataset: req.base_dataset,
|
||||||
|
columns: req.columns,
|
||||||
|
row_filter: req.row_filter,
|
||||||
|
column_redactions: req.column_redactions,
|
||||||
|
created_at: chrono::Utc::now(),
|
||||||
|
created_by: req.created_by,
|
||||||
|
description: req.description,
|
||||||
|
};
|
||||||
|
match registry.put_view(view).await {
|
||||||
|
Ok(v) => Ok((StatusCode::CREATED, Json(v))),
|
||||||
|
Err(e) => Err((StatusCode::BAD_REQUEST, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_views(State(registry): State<Registry>) -> impl IntoResponse {
|
||||||
|
Json(registry.list_views().await)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_view(
|
||||||
|
State(registry): State<Registry>,
|
||||||
|
Path(name): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match registry.get_view(&name).await {
|
||||||
|
Some(v) => Ok(Json(v)),
|
||||||
|
None => Err((StatusCode::NOT_FOUND, format!("view not found: {name}"))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_view(
|
||||||
|
State(registry): State<Registry>,
|
||||||
|
Path(name): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match registry.delete_view(&name).await {
|
||||||
|
Ok(()) => Ok(StatusCode::NO_CONTENT),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Phase E: soft-delete tombstones ---
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct TombstoneRequest {
|
||||||
|
row_key_column: String,
|
||||||
|
row_key_values: Vec<String>,
|
||||||
|
#[serde(default)]
|
||||||
|
actor: String,
|
||||||
|
#[serde(default)]
|
||||||
|
reason: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize)]
|
||||||
|
struct TombstoneResponse {
|
||||||
|
dataset: String,
|
||||||
|
row_key_column: String,
|
||||||
|
rows_tombstoned: usize,
|
||||||
|
failures: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn tombstone_rows(
|
||||||
|
State(registry): State<Registry>,
|
||||||
|
Path(name): Path<String>,
|
||||||
|
Json(req): Json<TombstoneRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
if req.row_key_values.is_empty() {
|
||||||
|
return Err((StatusCode::BAD_REQUEST, "row_key_values is empty".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut ok = 0;
|
||||||
|
let mut failures = Vec::new();
|
||||||
|
for value in &req.row_key_values {
|
||||||
|
match registry
|
||||||
|
.add_tombstone(&name, &req.row_key_column, value, &req.actor, &req.reason)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(_) => ok += 1,
|
||||||
|
Err(e) => failures.push(format!("{value}: {e}")),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let status = if ok > 0 && failures.is_empty() {
|
||||||
|
StatusCode::CREATED
|
||||||
|
} else if ok > 0 {
|
||||||
|
StatusCode::MULTI_STATUS
|
||||||
|
} else {
|
||||||
|
StatusCode::BAD_REQUEST
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok((status, Json(TombstoneResponse {
|
||||||
|
dataset: name,
|
||||||
|
row_key_column: req.row_key_column,
|
||||||
|
rows_tombstoned: ok,
|
||||||
|
failures,
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_tombstones(
|
||||||
|
State(registry): State<Registry>,
|
||||||
|
Path(name): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match registry.list_tombstones(&name).await {
|
||||||
|
Ok(ts) => Ok(Json(ts)),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Phase 17: Model profiles ---
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct CreateProfileRequest {
|
||||||
|
id: String,
|
||||||
|
ollama_name: String,
|
||||||
|
#[serde(default)]
|
||||||
|
description: String,
|
||||||
|
bound_datasets: Vec<String>,
|
||||||
|
#[serde(default)]
|
||||||
|
hnsw_config: shared::types::ProfileHnswConfig,
|
||||||
|
#[serde(default = "default_embed_model_req")]
|
||||||
|
embed_model: String,
|
||||||
|
#[serde(default)]
|
||||||
|
created_by: String,
|
||||||
|
/// Federation: optional per-profile bucket (`profile:{id}` by convention).
|
||||||
|
/// Omitting keeps artifacts in primary.
|
||||||
|
#[serde(default)]
|
||||||
|
bucket: Option<String>,
|
||||||
|
/// ADR-019 hybrid: which vector backend to route this profile's
|
||||||
|
/// indexes to. Defaults to Parquet+HNSW.
|
||||||
|
#[serde(default)]
|
||||||
|
vector_backend: shared::types::VectorBackend,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_embed_model_req() -> String { "nomic-embed-text".to_string() }
|
||||||
|
|
||||||
|
async fn create_profile(
|
||||||
|
State(registry): State<Registry>,
|
||||||
|
Json(req): Json<CreateProfileRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let profile = shared::types::ModelProfile {
|
||||||
|
id: req.id,
|
||||||
|
ollama_name: req.ollama_name,
|
||||||
|
description: req.description,
|
||||||
|
bound_datasets: req.bound_datasets,
|
||||||
|
hnsw_config: req.hnsw_config,
|
||||||
|
embed_model: req.embed_model,
|
||||||
|
created_at: chrono::Utc::now(),
|
||||||
|
created_by: req.created_by,
|
||||||
|
bucket: req.bucket,
|
||||||
|
vector_backend: req.vector_backend,
|
||||||
|
};
|
||||||
|
match registry.put_profile(profile).await {
|
||||||
|
Ok(p) => Ok((StatusCode::CREATED, Json(p))),
|
||||||
|
Err(e) => Err((StatusCode::BAD_REQUEST, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_profiles(State(registry): State<Registry>) -> impl IntoResponse {
|
||||||
|
Json(registry.list_profiles().await)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_profile(
|
||||||
|
State(registry): State<Registry>,
|
||||||
|
Path(id): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match registry.get_profile(&id).await {
|
||||||
|
Some(p) => Ok(Json(p)),
|
||||||
|
None => Err((StatusCode::NOT_FOUND, format!("profile not found: {id}"))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_profile(
|
||||||
|
State(registry): State<Registry>,
|
||||||
|
Path(id): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match registry.delete_profile(&id).await {
|
||||||
|
Ok(()) => Ok(StatusCode::NO_CONTENT),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,151 @@
|
|||||||
|
//! Soft-delete tombstone storage (Phase E).
|
||||||
|
//!
|
||||||
|
//! One append-log per dataset at `_catalog/tombstones/{dataset}/batch_*.jsonl`.
|
||||||
|
//! Uses the shared `storaged::append_log::AppendLog` pattern so appends
|
||||||
|
//! are write-once (never rewrites existing files) and can be compacted.
|
||||||
|
//!
|
||||||
|
//! The store exposes a per-dataset cache of active tombstones for the
|
||||||
|
//! hot path (queryd filter construction) — that's why events are pulled
|
||||||
|
//! into an in-memory map on every call rather than scanning object
|
||||||
|
//! storage repeatedly.
|
||||||
|
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use shared::types::Tombstone;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use storaged::append_log::{AppendLog, CompactStats};
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
const TOMBSTONE_PREFIX: &str = "_catalog/tombstones";
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct TombstoneStore {
|
||||||
|
store: Arc<dyn ObjectStore>,
|
||||||
|
logs: Arc<RwLock<HashMap<String, Arc<AppendLog>>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TombstoneStore {
|
||||||
|
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||||
|
Self {
|
||||||
|
store,
|
||||||
|
logs: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn prefix_for(dataset: &str) -> String {
|
||||||
|
// Sanitize dataset name for filesystem safety.
|
||||||
|
let safe: String = dataset
|
||||||
|
.chars()
|
||||||
|
.map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' })
|
||||||
|
.collect();
|
||||||
|
format!("{TOMBSTONE_PREFIX}/{}", safe)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn log_for(&self, dataset: &str) -> Arc<AppendLog> {
|
||||||
|
if let Some(log) = self.logs.read().await.get(dataset) {
|
||||||
|
return log.clone();
|
||||||
|
}
|
||||||
|
let mut guard = self.logs.write().await;
|
||||||
|
if let Some(log) = guard.get(dataset) {
|
||||||
|
return log.clone();
|
||||||
|
}
|
||||||
|
// Threshold of 1 — every tombstone is high-value. Compliance/audit
|
||||||
|
// doesn't tolerate "lost on restart"; we trade a small file count
|
||||||
|
// for guaranteed durability. Compaction merges later if volume grows.
|
||||||
|
let log = Arc::new(
|
||||||
|
AppendLog::new(self.store.clone(), Self::prefix_for(dataset))
|
||||||
|
.with_flush_threshold(1),
|
||||||
|
);
|
||||||
|
guard.insert(dataset.to_string(), log.clone());
|
||||||
|
log
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append one tombstone. Validates that the `row_key_column` matches
|
||||||
|
/// the column already used for this dataset (all tombstones for a
|
||||||
|
/// dataset share one key column so the query filter is well-defined).
|
||||||
|
/// Forces a flush so the tombstone is durable before this call returns.
|
||||||
|
pub async fn append(&self, ts: &Tombstone) -> Result<(), String> {
|
||||||
|
let existing = self.list(&ts.dataset).await?;
|
||||||
|
if let Some(prior) = existing.first() {
|
||||||
|
if prior.row_key_column != ts.row_key_column {
|
||||||
|
return Err(format!(
|
||||||
|
"dataset '{}' already uses '{}' as tombstone key; cannot mix with '{}'",
|
||||||
|
ts.dataset, prior.row_key_column, ts.row_key_column,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let line = serde_json::to_vec(ts).map_err(|e| e.to_string())?;
|
||||||
|
let log = self.log_for(&ts.dataset).await;
|
||||||
|
log.append(line).await?;
|
||||||
|
// Belt-and-suspenders: explicit flush in case the threshold is
|
||||||
|
// ever raised. Tombstones must be durable on return.
|
||||||
|
log.flush().await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// All tombstones for a dataset (chronological).
|
||||||
|
pub async fn list(&self, dataset: &str) -> Result<Vec<Tombstone>, String> {
|
||||||
|
let log = self.log_for(dataset).await;
|
||||||
|
let lines = log.read_all().await?;
|
||||||
|
let mut out = Vec::with_capacity(lines.len());
|
||||||
|
for line in lines {
|
||||||
|
match serde_json::from_slice::<Tombstone>(&line) {
|
||||||
|
Ok(t) => out.push(t),
|
||||||
|
Err(e) => tracing::warn!("tombstones/{}: skip malformed entry: {e}", dataset),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(out)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Per-dataset grouped view used by queryd — returns a map of
|
||||||
|
/// `{dataset -> (row_key_column, set_of_values)}` for every dataset
|
||||||
|
/// that has any tombstones.
|
||||||
|
pub async fn all_grouped(
|
||||||
|
&self,
|
||||||
|
datasets: &[String],
|
||||||
|
) -> Result<HashMap<String, (String, Vec<String>)>, String> {
|
||||||
|
let mut grouped = HashMap::new();
|
||||||
|
for dataset in datasets {
|
||||||
|
let ts = match self.list(dataset).await {
|
||||||
|
Ok(ts) => ts,
|
||||||
|
Err(_) => continue,
|
||||||
|
};
|
||||||
|
if ts.is_empty() { continue; }
|
||||||
|
let col = ts[0].row_key_column.clone();
|
||||||
|
let values: Vec<String> = ts.iter().map(|t| t.row_key_value.clone()).collect();
|
||||||
|
grouped.insert(dataset.clone(), (col, values));
|
||||||
|
}
|
||||||
|
Ok(grouped)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn compact(&self, dataset: &str) -> Result<CompactStats, String> {
|
||||||
|
let log = self.log_for(dataset).await;
|
||||||
|
log.compact().await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Remove every tombstone for a dataset. Called after a successful
|
||||||
|
/// parquet compaction has physically deleted those rows — the
|
||||||
|
/// tombstones have done their job and the journal can be cleared.
|
||||||
|
///
|
||||||
|
/// Implementation: drop the per-dataset AppendLog from the cache and
|
||||||
|
/// delete all its batch files. Next write starts fresh.
|
||||||
|
pub async fn clear(&self, dataset: &str) -> Result<usize, String> {
|
||||||
|
let prefix = format!("{}/", Self::prefix_for(dataset));
|
||||||
|
let keys = storaged::ops::list(&self.store, Some(&prefix)).await?;
|
||||||
|
let matching: Vec<String> = keys
|
||||||
|
.into_iter()
|
||||||
|
.filter(|k| {
|
||||||
|
let basename = k.rsplit('/').next().unwrap_or(k);
|
||||||
|
basename.starts_with("batch_") && basename.ends_with(".jsonl")
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let count = matching.len();
|
||||||
|
for key in &matching {
|
||||||
|
let _ = storaged::ops::delete(&self.store, key).await;
|
||||||
|
}
|
||||||
|
self.logs.write().await.remove(dataset);
|
||||||
|
if count > 0 {
|
||||||
|
tracing::info!("cleared {count} tombstone batch files for '{}'", dataset);
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -9,6 +9,9 @@ storaged = { path = "../storaged" }
|
|||||||
catalogd = { path = "../catalogd" }
|
catalogd = { path = "../catalogd" }
|
||||||
queryd = { path = "../queryd" }
|
queryd = { path = "../queryd" }
|
||||||
aibridge = { path = "../aibridge" }
|
aibridge = { path = "../aibridge" }
|
||||||
|
ingestd = { path = "../ingestd" }
|
||||||
|
vectord = { path = "../vectord" }
|
||||||
|
journald = { path = "../journald" }
|
||||||
tokio = { workspace = true }
|
tokio = { workspace = true }
|
||||||
axum = { workspace = true }
|
axum = { workspace = true }
|
||||||
serde = { workspace = true }
|
serde = { workspace = true }
|
||||||
@@ -23,3 +26,5 @@ opentelemetry = { workspace = true }
|
|||||||
opentelemetry_sdk = { workspace = true }
|
opentelemetry_sdk = { workspace = true }
|
||||||
opentelemetry-stdout = { workspace = true }
|
opentelemetry-stdout = { workspace = true }
|
||||||
tracing-opentelemetry = { workspace = true }
|
tracing-opentelemetry = { workspace = true }
|
||||||
|
arrow = { workspace = true }
|
||||||
|
chrono = { workspace = true }
|
||||||
|
|||||||
@@ -0,0 +1,155 @@
|
|||||||
|
/// Access control: field-level sensitivity enforcement, column masking, query audit.
|
||||||
|
/// Evaluates policies at query time — not at storage level.
|
||||||
|
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use shared::types::Sensitivity;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
/// An agent's role determines what sensitivity levels they can see.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct AgentRole {
|
||||||
|
pub agent_name: String,
|
||||||
|
pub role: String, // "recruiter", "admin", "analyst", "agent"
|
||||||
|
pub allowed_sensitivity: Vec<Sensitivity>, // what they can see unmasked
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A query audit entry.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct QueryAudit {
|
||||||
|
pub id: String,
|
||||||
|
pub agent: String,
|
||||||
|
pub sql: String,
|
||||||
|
pub datasets_accessed: Vec<String>,
|
||||||
|
pub pii_fields_accessed: Vec<String>,
|
||||||
|
pub timestamp: DateTime<Utc>,
|
||||||
|
pub row_count: usize,
|
||||||
|
pub allowed: bool,
|
||||||
|
pub masked_fields: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Access control manager.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct AccessControl {
|
||||||
|
roles: Arc<RwLock<HashMap<String, AgentRole>>>,
|
||||||
|
audit_log: Arc<RwLock<Vec<QueryAudit>>>,
|
||||||
|
enabled: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AccessControl {
|
||||||
|
pub fn new(enabled: bool) -> Self {
|
||||||
|
let ac = Self {
|
||||||
|
roles: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
audit_log: Arc::new(RwLock::new(Vec::new())),
|
||||||
|
enabled,
|
||||||
|
};
|
||||||
|
ac
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Register default roles.
|
||||||
|
pub async fn register_defaults(&self) {
|
||||||
|
let defaults = vec![
|
||||||
|
AgentRole {
|
||||||
|
agent_name: "admin".into(),
|
||||||
|
role: "admin".into(),
|
||||||
|
allowed_sensitivity: vec![
|
||||||
|
Sensitivity::Public, Sensitivity::Internal,
|
||||||
|
Sensitivity::Pii, Sensitivity::Phi, Sensitivity::Financial,
|
||||||
|
],
|
||||||
|
},
|
||||||
|
AgentRole {
|
||||||
|
agent_name: "recruiter".into(),
|
||||||
|
role: "recruiter".into(),
|
||||||
|
allowed_sensitivity: vec![
|
||||||
|
Sensitivity::Public, Sensitivity::Internal, Sensitivity::Pii,
|
||||||
|
],
|
||||||
|
},
|
||||||
|
AgentRole {
|
||||||
|
agent_name: "analyst".into(),
|
||||||
|
role: "analyst".into(),
|
||||||
|
allowed_sensitivity: vec![
|
||||||
|
Sensitivity::Public, Sensitivity::Internal, Sensitivity::Financial,
|
||||||
|
],
|
||||||
|
},
|
||||||
|
AgentRole {
|
||||||
|
agent_name: "agent".into(),
|
||||||
|
role: "agent".into(),
|
||||||
|
allowed_sensitivity: vec![
|
||||||
|
Sensitivity::Public, Sensitivity::Internal,
|
||||||
|
],
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
let mut roles = self.roles.write().await;
|
||||||
|
for role in defaults {
|
||||||
|
roles.insert(role.agent_name.clone(), role);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Register or update an agent role.
|
||||||
|
pub async fn set_role(&self, role: AgentRole) {
|
||||||
|
self.roles.write().await.insert(role.agent_name.clone(), role);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get an agent's role.
|
||||||
|
pub async fn get_role(&self, agent: &str) -> Option<AgentRole> {
|
||||||
|
self.roles.read().await.get(agent).cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all roles.
|
||||||
|
pub async fn list_roles(&self) -> Vec<AgentRole> {
|
||||||
|
self.roles.read().await.values().cloned().collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if an agent can see a field with given sensitivity.
|
||||||
|
pub async fn can_access(&self, agent: &str, sensitivity: &Sensitivity) -> bool {
|
||||||
|
if !self.enabled { return true; }
|
||||||
|
match self.roles.read().await.get(agent) {
|
||||||
|
Some(role) => role.allowed_sensitivity.contains(sensitivity),
|
||||||
|
None => false, // unknown agent = no access
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Determine which fields should be masked for an agent.
|
||||||
|
pub async fn masked_fields(
|
||||||
|
&self,
|
||||||
|
agent: &str,
|
||||||
|
columns: &[shared::types::ColumnMeta],
|
||||||
|
) -> Vec<String> {
|
||||||
|
if !self.enabled { return vec![]; }
|
||||||
|
|
||||||
|
let role = match self.roles.read().await.get(agent) {
|
||||||
|
Some(r) => r.clone(),
|
||||||
|
None => return columns.iter().filter(|c| c.sensitivity.is_some()).map(|c| c.name.clone()).collect(),
|
||||||
|
};
|
||||||
|
|
||||||
|
columns.iter()
|
||||||
|
.filter(|col| {
|
||||||
|
if let Some(ref sens) = col.sensitivity {
|
||||||
|
!role.allowed_sensitivity.contains(sens)
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.map(|col| col.name.clone())
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Log a query for audit.
|
||||||
|
pub async fn log_query(&self, audit: QueryAudit) {
|
||||||
|
self.audit_log.write().await.push(audit);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get recent audit entries.
|
||||||
|
pub async fn recent_audit(&self, limit: usize) -> Vec<QueryAudit> {
|
||||||
|
let log = self.audit_log.read().await;
|
||||||
|
let start = log.len().saturating_sub(limit);
|
||||||
|
log[start..].iter().rev().cloned().collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn is_enabled(&self) -> bool {
|
||||||
|
self.enabled
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,62 @@
|
|||||||
|
use axum::{
|
||||||
|
Json, Router,
|
||||||
|
extract::{Query, State},
|
||||||
|
http::StatusCode,
|
||||||
|
response::IntoResponse,
|
||||||
|
routing::{get, post},
|
||||||
|
};
|
||||||
|
use serde::Deserialize;
|
||||||
|
|
||||||
|
use crate::access::{AccessControl, AgentRole};
|
||||||
|
|
||||||
|
pub fn router(ac: AccessControl) -> Router {
|
||||||
|
Router::new()
|
||||||
|
.route("/roles", get(list_roles))
|
||||||
|
.route("/roles", post(set_role))
|
||||||
|
.route("/audit", get(query_audit))
|
||||||
|
.route("/check", post(check_access))
|
||||||
|
.with_state(ac)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_roles(State(ac): State<AccessControl>) -> impl IntoResponse {
|
||||||
|
Json(ac.list_roles().await)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn set_role(
|
||||||
|
State(ac): State<AccessControl>,
|
||||||
|
Json(role): Json<AgentRole>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let name = role.agent_name.clone();
|
||||||
|
ac.set_role(role).await;
|
||||||
|
(StatusCode::OK, format!("role set: {name}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct AuditQuery {
|
||||||
|
limit: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn query_audit(
|
||||||
|
State(ac): State<AccessControl>,
|
||||||
|
Query(q): Query<AuditQuery>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
Json(ac.recent_audit(q.limit.unwrap_or(50)).await)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct CheckRequest {
|
||||||
|
agent: String,
|
||||||
|
sensitivity: shared::types::Sensitivity,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn check_access(
|
||||||
|
State(ac): State<AccessControl>,
|
||||||
|
Json(req): Json<CheckRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let allowed = ac.can_access(&req.agent, &req.sensitivity).await;
|
||||||
|
Json(serde_json::json!({
|
||||||
|
"agent": req.agent,
|
||||||
|
"sensitivity": req.sensitivity,
|
||||||
|
"allowed": allowed,
|
||||||
|
}))
|
||||||
|
}
|
||||||
+148
-8
@@ -1,7 +1,10 @@
|
|||||||
|
mod access;
|
||||||
|
mod access_service;
|
||||||
mod auth;
|
mod auth;
|
||||||
mod observability;
|
mod observability;
|
||||||
|
mod tools;
|
||||||
|
|
||||||
use axum::{Router, routing::get};
|
use axum::{Router, extract::DefaultBodyLimit, routing::get};
|
||||||
use proto::lakehouse::catalog_service_server::CatalogServiceServer;
|
use proto::lakehouse::catalog_service_server::CatalogServiceServer;
|
||||||
use shared::config::Config;
|
use shared::config::Config;
|
||||||
use tower_http::cors::{Any, CorsLayer};
|
use tower_http::cors::{Any, CorsLayer};
|
||||||
@@ -21,8 +24,30 @@ async fn main() {
|
|||||||
tracing::info!("config loaded: gateway={}:{}, storage={}",
|
tracing::info!("config loaded: gateway={}:{}, storage={}",
|
||||||
config.gateway.host, config.gateway.port, config.storage.root);
|
config.gateway.host, config.gateway.port, config.storage.root);
|
||||||
|
|
||||||
// Storage backend
|
// Secrets provider
|
||||||
let store = storaged::backend::init_local(&config.storage.root);
|
let secrets = std::sync::Arc::new(shared::secrets::FileSecretsProvider::from_env());
|
||||||
|
if let Err(e) = secrets.load().await {
|
||||||
|
tracing::warn!("secrets load: {e} — bucket credentials may be missing");
|
||||||
|
}
|
||||||
|
|
||||||
|
// Federation: bucket registry (primary + any [[storage.buckets]] entries + rescue)
|
||||||
|
let bucket_registry = std::sync::Arc::new(
|
||||||
|
storaged::registry::BucketRegistry::from_config(&config.storage, secrets.clone())
|
||||||
|
.await
|
||||||
|
.expect("bucket registry init failed")
|
||||||
|
);
|
||||||
|
let buckets = bucket_registry.list().await;
|
||||||
|
tracing::info!(
|
||||||
|
"bucket registry: {} buckets [{}], rescue={:?}",
|
||||||
|
buckets.len(),
|
||||||
|
buckets.iter().map(|b| format!("{}:{}", b.name, b.backend)).collect::<Vec<_>>().join(", "),
|
||||||
|
bucket_registry.rescue_name(),
|
||||||
|
);
|
||||||
|
|
||||||
|
// Back-compat: existing components that expect a single Arc<ObjectStore>
|
||||||
|
// get the primary bucket's store. As we migrate call sites to be bucket-
|
||||||
|
// aware, they'll call bucket_registry.get(name) directly.
|
||||||
|
let store = bucket_registry.default_store();
|
||||||
|
|
||||||
// Catalog
|
// Catalog
|
||||||
let registry = catalogd::registry::Registry::new(store.clone());
|
let registry = catalogd::registry::Registry::new(store.clone());
|
||||||
@@ -30,19 +55,126 @@ async fn main() {
|
|||||||
tracing::warn!("catalog rebuild failed (empty store?): {e}");
|
tracing::warn!("catalog rebuild failed (empty store?): {e}");
|
||||||
}
|
}
|
||||||
|
|
||||||
// Query engine
|
// Query engine with 16GB memory cache
|
||||||
let engine = queryd::context::QueryEngine::new(registry.clone(), store.clone());
|
let cache = queryd::cache::MemCache::new(16 * 1024 * 1024 * 1024);
|
||||||
|
let engine = queryd::context::QueryEngine::new(registry.clone(), bucket_registry.clone(), cache);
|
||||||
|
|
||||||
|
// Event journal — append-only mutation log (flush every 100 events)
|
||||||
|
let journal = journald::journal::Journal::new(store.clone(), 100);
|
||||||
|
|
||||||
|
// Access control
|
||||||
|
let access = access::AccessControl::new(config.auth.enabled);
|
||||||
|
access.register_defaults().await;
|
||||||
|
|
||||||
|
// Workspace manager for agent-specific overlays
|
||||||
|
let workspace_mgr = queryd::workspace::WorkspaceManager::new(store.clone());
|
||||||
|
if let Err(e) = workspace_mgr.rebuild().await {
|
||||||
|
tracing::warn!("workspace rebuild: {e}");
|
||||||
|
}
|
||||||
|
|
||||||
// AI sidecar client
|
// AI sidecar client
|
||||||
let ai_client = aibridge::client::AiClient::new(&config.sidecar.url);
|
let ai_client = aibridge::client::AiClient::new(&config.sidecar.url);
|
||||||
|
|
||||||
|
// Vector service components — built before the router because both the
|
||||||
|
// /vectors service AND ingestd need the agent handle to enqueue triggers.
|
||||||
|
let index_reg = vectord::index_registry::IndexRegistry::new(store.clone());
|
||||||
|
let _ = index_reg.rebuild().await;
|
||||||
|
let hnsw = vectord::hnsw::HnswStore::new();
|
||||||
|
let emb_cache = vectord::embedding_cache::EmbeddingCache::new(store.clone());
|
||||||
|
// Phase B (federation layer 2): trial journals + promotion files now
|
||||||
|
// live in each index's recorded bucket (IndexMeta.bucket), not always
|
||||||
|
// primary. Journal / registry resolve per-call via the bucket registry.
|
||||||
|
let tj = vectord::trial::TrialJournal::new(bucket_registry.clone(), index_reg.clone());
|
||||||
|
let pr = vectord::promotion::PromotionRegistry::new(bucket_registry.clone(), index_reg.clone());
|
||||||
|
|
||||||
|
// Phase 16.2: spawn the autotune agent. When config.agent.enabled=false
|
||||||
|
// this returns a handle that drops triggers silently — no surprise load.
|
||||||
|
let agent_cfg = vectord::agent::AgentConfig {
|
||||||
|
enabled: config.agent.enabled,
|
||||||
|
cycle_interval_secs: config.agent.cycle_interval_secs,
|
||||||
|
cooldown_between_trials_secs: config.agent.cooldown_between_trials_secs,
|
||||||
|
min_recall: config.agent.min_recall,
|
||||||
|
max_trials_per_hour: config.agent.max_trials_per_hour,
|
||||||
|
};
|
||||||
|
let agent_handle = vectord::agent::spawn(
|
||||||
|
agent_cfg,
|
||||||
|
vectord::agent::AgentDeps {
|
||||||
|
store: store.clone(),
|
||||||
|
ai_client: ai_client.clone(),
|
||||||
|
catalog: registry.clone(),
|
||||||
|
index_registry: index_reg.clone(),
|
||||||
|
hnsw_store: hnsw.clone(),
|
||||||
|
embedding_cache: emb_cache.clone(),
|
||||||
|
trial_journal: tj.clone(),
|
||||||
|
promotion_registry: pr.clone(),
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
// HTTP router
|
// HTTP router
|
||||||
let mut app = Router::new()
|
let mut app = Router::new()
|
||||||
.route("/health", get(health))
|
.route("/health", get(health))
|
||||||
.nest("/storage", storaged::service::router(store))
|
.nest("/storage", storaged::service::router(store.clone())
|
||||||
|
.merge(storaged::federation_service::router(bucket_registry.clone())))
|
||||||
.nest("/catalog", catalogd::service::router(registry.clone()))
|
.nest("/catalog", catalogd::service::router(registry.clone()))
|
||||||
.nest("/query", queryd::service::router(engine))
|
.nest("/query", queryd::service::router(engine.clone()))
|
||||||
.nest("/ai", aibridge::service::router(ai_client));
|
.nest("/ai", aibridge::service::router(ai_client.clone()))
|
||||||
|
;
|
||||||
|
|
||||||
|
// Phase E: scheduled ingest
|
||||||
|
let sched_store = ingestd::schedule::ScheduleStore::new(store.clone());
|
||||||
|
if let Ok(n) = sched_store.rebuild().await {
|
||||||
|
if n > 0 { tracing::info!("rebuilt {n} persisted schedule(s)"); }
|
||||||
|
}
|
||||||
|
let scheduler = ingestd::schedule::Scheduler {
|
||||||
|
store: sched_store.clone(),
|
||||||
|
ingest: ingestd::schedule::SchedulerIngestDeps {
|
||||||
|
store: store.clone(),
|
||||||
|
registry: registry.clone(),
|
||||||
|
buckets: bucket_registry.clone(),
|
||||||
|
agent_handle: Some(agent_handle.clone()),
|
||||||
|
index_registry: index_reg.clone(),
|
||||||
|
},
|
||||||
|
tick_secs: 10,
|
||||||
|
};
|
||||||
|
scheduler.spawn();
|
||||||
|
|
||||||
|
app = app
|
||||||
|
.nest("/ingest", ingestd::service::router(ingestd::service::IngestState {
|
||||||
|
store: store.clone(),
|
||||||
|
registry: registry.clone(),
|
||||||
|
buckets: bucket_registry.clone(),
|
||||||
|
agent_handle: agent_handle.clone(),
|
||||||
|
index_registry: index_reg.clone(),
|
||||||
|
schedules: sched_store,
|
||||||
|
}))
|
||||||
|
.nest("/vectors", vectord::service::router(vectord::service::VectorState {
|
||||||
|
store: store.clone(),
|
||||||
|
ai_client: ai_client.clone(),
|
||||||
|
job_tracker: vectord::jobs::JobTracker::new(),
|
||||||
|
index_registry: index_reg.clone(),
|
||||||
|
hnsw_store: hnsw,
|
||||||
|
embedding_cache: emb_cache,
|
||||||
|
trial_journal: tj,
|
||||||
|
catalog: registry.clone(),
|
||||||
|
promotion_registry: pr,
|
||||||
|
agent_handle,
|
||||||
|
bucket_registry: bucket_registry.clone(),
|
||||||
|
active_profile: std::sync::Arc::new(tokio::sync::RwLock::new(None)),
|
||||||
|
lance: vectord::lance_backend::LanceRegistry::new(
|
||||||
|
bucket_registry.clone(), index_reg.clone(),
|
||||||
|
),
|
||||||
|
}))
|
||||||
|
.nest("/workspaces", queryd::workspace_service::router(workspace_mgr))
|
||||||
|
.nest("/journal", journald::service::router(journal))
|
||||||
|
.nest("/access", access_service::router(access))
|
||||||
|
.nest("/tools", tools::service::router({
|
||||||
|
let tool_reg = tools::registry::ToolRegistry::new_with_defaults();
|
||||||
|
tool_reg.register_defaults().await;
|
||||||
|
tools::ToolState {
|
||||||
|
registry: tool_reg,
|
||||||
|
query_fn: tools::QueryExecutor::new(engine.clone()),
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
|
||||||
// Auth middleware (if enabled)
|
// Auth middleware (if enabled)
|
||||||
if config.auth.enabled {
|
if config.auth.enabled {
|
||||||
@@ -58,12 +190,20 @@ async fn main() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
app = app
|
app = app
|
||||||
|
.layer(DefaultBodyLimit::max(256 * 1024 * 1024)) // 256MB
|
||||||
.layer(CorsLayer::new()
|
.layer(CorsLayer::new()
|
||||||
.allow_origin(Any)
|
.allow_origin(Any)
|
||||||
.allow_methods(Any)
|
.allow_methods(Any)
|
||||||
.allow_headers(Any))
|
.allow_headers(Any))
|
||||||
.layer(TraceLayer::new_for_http());
|
.layer(TraceLayer::new_for_http());
|
||||||
|
|
||||||
|
// File watcher — auto-ingest from ./inbox
|
||||||
|
let _watcher = ingestd::watcher::spawn_watcher(
|
||||||
|
ingestd::watcher::WatchConfig::default(),
|
||||||
|
store.clone(),
|
||||||
|
registry.clone(),
|
||||||
|
);
|
||||||
|
|
||||||
// Start gRPC server on port+1
|
// Start gRPC server on port+1
|
||||||
let grpc_port = config.gateway.port + 1;
|
let grpc_port = config.gateway.port + 1;
|
||||||
let catalog_grpc = catalogd::grpc::CatalogGrpc::new(registry);
|
let catalog_grpc = catalogd::grpc::CatalogGrpc::new(registry);
|
||||||
|
|||||||
@@ -0,0 +1,47 @@
|
|||||||
|
pub mod registry;
|
||||||
|
pub mod service;
|
||||||
|
|
||||||
|
use queryd::context::QueryEngine;
|
||||||
|
use arrow::json::writer::{JsonArray, Writer as JsonWriter};
|
||||||
|
|
||||||
|
/// State for the tool system.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct ToolState {
|
||||||
|
pub registry: registry::ToolRegistry,
|
||||||
|
pub query_fn: QueryExecutor,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Wraps QueryEngine to provide a simple execute interface for tools.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct QueryExecutor {
|
||||||
|
engine: QueryEngine,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl QueryExecutor {
|
||||||
|
pub fn new(engine: QueryEngine) -> Self {
|
||||||
|
Self { engine }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute SQL and return (rows as JSON, row count).
|
||||||
|
pub async fn execute(&self, sql: &str) -> Result<(serde_json::Value, usize), String> {
|
||||||
|
let batches = self.engine.query(sql).await?;
|
||||||
|
|
||||||
|
if batches.is_empty() {
|
||||||
|
return Ok((serde_json::Value::Array(vec![]), 0));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
let mut writer = JsonWriter::<_, JsonArray>::new(&mut buf);
|
||||||
|
for batch in &batches {
|
||||||
|
writer.write(batch).map_err(|e| format!("JSON write: {e}"))?;
|
||||||
|
}
|
||||||
|
writer.finish().map_err(|e| format!("JSON finish: {e}"))?;
|
||||||
|
drop(writer);
|
||||||
|
|
||||||
|
let rows: serde_json::Value = serde_json::from_slice(&buf)
|
||||||
|
.map_err(|e| format!("JSON parse: {e}"))?;
|
||||||
|
let count = rows.as_array().map(|a| a.len()).unwrap_or(0);
|
||||||
|
|
||||||
|
Ok((rows, count))
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,274 @@
|
|||||||
|
/// Tool Registry: named, governed business actions for AI agents.
|
||||||
|
/// Instead of raw SQL, agents call validated tools with audit trails.
|
||||||
|
///
|
||||||
|
/// Each tool has:
|
||||||
|
/// - Name and description (for LLM tool-use)
|
||||||
|
/// - Parameter schema (validated before execution)
|
||||||
|
/// - Permission level (read / write / admin)
|
||||||
|
/// - Audit logging (every invocation recorded)
|
||||||
|
/// - Rate limiting (per agent)
|
||||||
|
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
/// Permission level for a tool.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||||
|
#[serde(rename_all = "lowercase")]
|
||||||
|
pub enum Permission {
|
||||||
|
Read, // auto-approved, no side effects
|
||||||
|
Write, // modifies data, logged
|
||||||
|
Admin, // destructive, requires confirmation
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parameter definition for a tool.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ParamDef {
|
||||||
|
pub name: String,
|
||||||
|
pub param_type: String, // "string", "integer", "boolean", "float"
|
||||||
|
pub required: bool,
|
||||||
|
pub description: String,
|
||||||
|
pub default: Option<serde_json::Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Tool definition — what agents see and can call.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ToolDef {
|
||||||
|
pub name: String,
|
||||||
|
pub description: String,
|
||||||
|
pub permission: Permission,
|
||||||
|
pub parameters: Vec<ParamDef>,
|
||||||
|
pub returns: String, // description of return value
|
||||||
|
pub sql_template: String, // SQL with {param} placeholders
|
||||||
|
pub category: String, // "candidates", "placements", "analytics"
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Audit log entry for a tool invocation.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ToolInvocation {
|
||||||
|
pub id: String,
|
||||||
|
pub tool_name: String,
|
||||||
|
pub agent: String,
|
||||||
|
pub params: serde_json::Value,
|
||||||
|
pub permission: Permission,
|
||||||
|
pub timestamp: DateTime<Utc>,
|
||||||
|
pub success: bool,
|
||||||
|
pub error: Option<String>,
|
||||||
|
pub rows_returned: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The registry — holds tool definitions and audit log.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct ToolRegistry {
|
||||||
|
tools: Arc<RwLock<HashMap<String, ToolDef>>>,
|
||||||
|
audit_log: Arc<RwLock<Vec<ToolInvocation>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ToolRegistry {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
let registry = Self {
|
||||||
|
tools: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
audit_log: Arc::new(RwLock::new(Vec::new())),
|
||||||
|
};
|
||||||
|
// Register built-in staffing tools
|
||||||
|
tokio::task::block_in_place(|| {
|
||||||
|
tokio::runtime::Handle::current().block_on(registry.register_defaults())
|
||||||
|
});
|
||||||
|
registry
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn new_with_defaults() -> Self {
|
||||||
|
let registry = Self {
|
||||||
|
tools: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
audit_log: Arc::new(RwLock::new(Vec::new())),
|
||||||
|
};
|
||||||
|
registry
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Register default staffing tools.
|
||||||
|
pub async fn register_defaults(&self) {
|
||||||
|
let tools = vec![
|
||||||
|
ToolDef {
|
||||||
|
name: "search_candidates".into(),
|
||||||
|
description: "Search for candidates by skills, city, state, availability, and experience. Returns matching candidates with contact info.".into(),
|
||||||
|
permission: Permission::Read,
|
||||||
|
parameters: vec![
|
||||||
|
ParamDef { name: "skills".into(), param_type: "string".into(), required: false, description: "Comma-separated skills to match (uses LIKE)".into(), default: None },
|
||||||
|
ParamDef { name: "city".into(), param_type: "string".into(), required: false, description: "City name".into(), default: None },
|
||||||
|
ParamDef { name: "state".into(), param_type: "string".into(), required: false, description: "State abbreviation".into(), default: None },
|
||||||
|
ParamDef { name: "min_years".into(), param_type: "integer".into(), required: false, description: "Minimum years of experience".into(), default: Some(serde_json::json!(0)) },
|
||||||
|
ParamDef { name: "status".into(), param_type: "string".into(), required: false, description: "Candidate status (active, inactive, placed)".into(), default: Some(serde_json::json!("active")) },
|
||||||
|
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Max results".into(), default: Some(serde_json::json!(20)) },
|
||||||
|
],
|
||||||
|
returns: "List of candidates with id, name, phone, email, skills, experience".into(),
|
||||||
|
sql_template: "SELECT candidate_id, first_name, last_name, phone, email, city, state, zip, vertical, skills, years_experience FROM candidates WHERE 1=1 {skills_filter} {city_filter} {state_filter} {years_filter} {status_filter} ORDER BY years_experience DESC LIMIT {limit}".into(),
|
||||||
|
category: "candidates".into(),
|
||||||
|
},
|
||||||
|
ToolDef {
|
||||||
|
name: "get_candidate".into(),
|
||||||
|
description: "Get full details for a specific candidate by ID.".into(),
|
||||||
|
permission: Permission::Read,
|
||||||
|
parameters: vec![
|
||||||
|
ParamDef { name: "candidate_id".into(), param_type: "string".into(), required: true, description: "Candidate ID (e.g. CAND-000001)".into(), default: None },
|
||||||
|
],
|
||||||
|
returns: "Full candidate record".into(),
|
||||||
|
sql_template: "SELECT * FROM candidates WHERE candidate_id = '{candidate_id}'".into(),
|
||||||
|
category: "candidates".into(),
|
||||||
|
},
|
||||||
|
ToolDef {
|
||||||
|
name: "revenue_by_client".into(),
|
||||||
|
description: "Show total billed revenue, pay costs, and gross profit by client. Filter by date range.".into(),
|
||||||
|
permission: Permission::Read,
|
||||||
|
parameters: vec![
|
||||||
|
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Top N clients".into(), default: Some(serde_json::json!(10)) },
|
||||||
|
],
|
||||||
|
returns: "Client name, total billed, total paid, gross profit, timesheet count".into(),
|
||||||
|
sql_template: "SELECT c.company_name, COUNT(*) as timesheets, ROUND(SUM(t.bill_total),2) as total_billed, ROUND(SUM(t.pay_total),2) as total_paid, ROUND(SUM(t.bill_total) - SUM(t.pay_total),2) as gross_profit FROM timesheets t JOIN clients c ON t.client_id = c.client_id WHERE t.approved = true GROUP BY c.company_name ORDER BY total_billed DESC LIMIT {limit}".into(),
|
||||||
|
category: "analytics".into(),
|
||||||
|
},
|
||||||
|
ToolDef {
|
||||||
|
name: "recruiter_performance".into(),
|
||||||
|
description: "Show recruiter performance: placements, unique candidates, and total revenue generated.".into(),
|
||||||
|
permission: Permission::Read,
|
||||||
|
parameters: vec![
|
||||||
|
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Top N recruiters".into(), default: Some(serde_json::json!(10)) },
|
||||||
|
],
|
||||||
|
returns: "Recruiter name, placement count, unique candidates, total revenue".into(),
|
||||||
|
sql_template: "SELECT p.recruiter, COUNT(DISTINCT p.placement_id) as placements, COUNT(DISTINCT p.candidate_id) as unique_candidates, ROUND(SUM(t.bill_total),2) as total_revenue FROM placements p JOIN timesheets t ON p.placement_id = t.placement_id GROUP BY p.recruiter ORDER BY total_revenue DESC LIMIT {limit}".into(),
|
||||||
|
category: "analytics".into(),
|
||||||
|
},
|
||||||
|
ToolDef {
|
||||||
|
name: "cold_leads".into(),
|
||||||
|
description: "Find candidates who were called multiple times but never placed — potential lost opportunities.".into(),
|
||||||
|
permission: Permission::Read,
|
||||||
|
parameters: vec![
|
||||||
|
ParamDef { name: "min_calls".into(), param_type: "integer".into(), required: false, description: "Minimum call count".into(), default: Some(serde_json::json!(5)) },
|
||||||
|
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Max results".into(), default: Some(serde_json::json!(20)) },
|
||||||
|
],
|
||||||
|
returns: "Candidates with call counts who were never placed".into(),
|
||||||
|
sql_template: "SELECT c.candidate_id, c.first_name, c.last_name, c.phone, c.vertical, cl.calls FROM candidates c JOIN (SELECT candidate_id, COUNT(*) as calls FROM call_log GROUP BY candidate_id HAVING COUNT(*) >= {min_calls}) cl ON c.candidate_id = cl.candidate_id WHERE c.candidate_id NOT IN (SELECT DISTINCT candidate_id FROM placements) ORDER BY cl.calls DESC LIMIT {limit}".into(),
|
||||||
|
category: "analytics".into(),
|
||||||
|
},
|
||||||
|
ToolDef {
|
||||||
|
name: "open_jobs".into(),
|
||||||
|
description: "List open job orders with client, title, rates, and location.".into(),
|
||||||
|
permission: Permission::Read,
|
||||||
|
parameters: vec![
|
||||||
|
ParamDef { name: "vertical".into(), param_type: "string".into(), required: false, description: "Filter by vertical (IT, Healthcare, etc)".into(), default: None },
|
||||||
|
ParamDef { name: "city".into(), param_type: "string".into(), required: false, description: "Filter by city".into(), default: None },
|
||||||
|
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Max results".into(), default: Some(serde_json::json!(20)) },
|
||||||
|
],
|
||||||
|
returns: "Open job orders with details".into(),
|
||||||
|
sql_template: "SELECT j.job_order_id, c.company_name, j.title, j.vertical, j.city, j.state, j.bill_rate, j.pay_rate, j.status FROM job_orders j JOIN clients c ON j.client_id = c.client_id WHERE j.status = 'open' {vertical_filter} {city_filter} ORDER BY j.bill_rate DESC LIMIT {limit}".into(),
|
||||||
|
category: "jobs".into(),
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
let mut reg = self.tools.write().await;
|
||||||
|
for tool in tools {
|
||||||
|
reg.insert(tool.name.clone(), tool);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get tool definition (for LLM tool-use schema).
|
||||||
|
pub async fn get_tool(&self, name: &str) -> Option<ToolDef> {
|
||||||
|
self.tools.read().await.get(name).cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all tools (for LLM tool discovery).
|
||||||
|
pub async fn list_tools(&self) -> Vec<ToolDef> {
|
||||||
|
self.tools.read().await.values().cloned().collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build SQL from tool parameters. Validates and sanitizes.
|
||||||
|
pub fn build_sql(tool: &ToolDef, params: &serde_json::Value) -> Result<String, String> {
|
||||||
|
let mut sql = tool.sql_template.clone();
|
||||||
|
let params = params.as_object().ok_or("params must be an object")?;
|
||||||
|
|
||||||
|
// Replace named parameters
|
||||||
|
for param_def in &tool.parameters {
|
||||||
|
let value = params.get(¶m_def.name)
|
||||||
|
.or(param_def.default.as_ref());
|
||||||
|
|
||||||
|
match param_def.name.as_str() {
|
||||||
|
// Handle filter parameters (conditional WHERE clauses)
|
||||||
|
"skills" => {
|
||||||
|
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
|
||||||
|
let filters: Vec<String> = v.split(',')
|
||||||
|
.map(|s| format!("skills LIKE '%{}%'", s.trim().replace('\'', "''")))
|
||||||
|
.collect();
|
||||||
|
sql = sql.replace("{skills_filter}", &format!("AND ({})", filters.join(" OR ")));
|
||||||
|
} else {
|
||||||
|
sql = sql.replace("{skills_filter}", "");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"city" => {
|
||||||
|
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
|
||||||
|
sql = sql.replace("{city_filter}", &format!("AND city = '{}'", v.replace('\'', "''")));
|
||||||
|
} else {
|
||||||
|
sql = sql.replace("{city_filter}", "");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"state" => {
|
||||||
|
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
|
||||||
|
sql = sql.replace("{state_filter}", &format!("AND state = '{}'", v.replace('\'', "''")));
|
||||||
|
} else {
|
||||||
|
sql = sql.replace("{state_filter}", "");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"vertical" => {
|
||||||
|
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
|
||||||
|
sql = sql.replace("{vertical_filter}", &format!("AND j.vertical = '{}'", v.replace('\'', "''")));
|
||||||
|
} else {
|
||||||
|
sql = sql.replace("{vertical_filter}", "");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"status" => {
|
||||||
|
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
|
||||||
|
sql = sql.replace("{status_filter}", &format!("AND status = '{}'", v.replace('\'', "''")));
|
||||||
|
} else {
|
||||||
|
sql = sql.replace("{status_filter}", "");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"min_years" => {
|
||||||
|
if let Some(v) = value.and_then(|v| v.as_i64()) {
|
||||||
|
sql = sql.replace("{years_filter}", &format!("AND years_experience >= {v}"));
|
||||||
|
} else {
|
||||||
|
sql = sql.replace("{years_filter}", "");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Direct substitution for simple params
|
||||||
|
_ => {
|
||||||
|
let placeholder = format!("{{{}}}", param_def.name);
|
||||||
|
if sql.contains(&placeholder) {
|
||||||
|
let val_str = match value {
|
||||||
|
Some(serde_json::Value::String(s)) => s.replace('\'', "''"),
|
||||||
|
Some(serde_json::Value::Number(n)) => n.to_string(),
|
||||||
|
Some(serde_json::Value::Bool(b)) => b.to_string(),
|
||||||
|
Some(v) => v.to_string(),
|
||||||
|
None if param_def.required => return Err(format!("missing required param: {}", param_def.name)),
|
||||||
|
None => continue,
|
||||||
|
};
|
||||||
|
sql = sql.replace(&placeholder, &val_str);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(sql)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Log a tool invocation.
|
||||||
|
pub async fn log_invocation(&self, inv: ToolInvocation) {
|
||||||
|
self.audit_log.write().await.push(inv);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get recent audit log.
|
||||||
|
pub async fn recent_audit(&self, limit: usize) -> Vec<ToolInvocation> {
|
||||||
|
let log = self.audit_log.read().await;
|
||||||
|
let start = log.len().saturating_sub(limit);
|
||||||
|
log[start..].iter().rev().cloned().collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,148 @@
|
|||||||
|
use axum::{
|
||||||
|
Json, Router,
|
||||||
|
extract::{Path, Query, State},
|
||||||
|
http::StatusCode,
|
||||||
|
response::IntoResponse,
|
||||||
|
routing::{get, post},
|
||||||
|
};
|
||||||
|
use serde::Deserialize;
|
||||||
|
|
||||||
|
use super::registry::{Permission, ToolInvocation, ToolRegistry};
|
||||||
|
use crate::tools::ToolState;
|
||||||
|
|
||||||
|
pub fn router(state: ToolState) -> Router {
|
||||||
|
Router::new()
|
||||||
|
.route("/", get(list_tools))
|
||||||
|
.route("/{name}", get(get_tool))
|
||||||
|
.route("/{name}/call", post(call_tool))
|
||||||
|
.route("/audit", get(audit_log))
|
||||||
|
.with_state(state)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all available tools (for LLM tool-use discovery).
|
||||||
|
async fn list_tools(State(state): State<ToolState>) -> impl IntoResponse {
|
||||||
|
let tools = state.registry.list_tools().await;
|
||||||
|
// Return in MCP-compatible format
|
||||||
|
let tool_list: Vec<serde_json::Value> = tools.iter().map(|t| {
|
||||||
|
serde_json::json!({
|
||||||
|
"name": t.name,
|
||||||
|
"description": t.description,
|
||||||
|
"permission": t.permission,
|
||||||
|
"category": t.category,
|
||||||
|
"parameters": t.parameters.iter().map(|p| {
|
||||||
|
serde_json::json!({
|
||||||
|
"name": p.name,
|
||||||
|
"type": p.param_type,
|
||||||
|
"required": p.required,
|
||||||
|
"description": p.description,
|
||||||
|
"default": p.default,
|
||||||
|
})
|
||||||
|
}).collect::<Vec<_>>(),
|
||||||
|
})
|
||||||
|
}).collect();
|
||||||
|
Json(tool_list)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get a specific tool definition.
|
||||||
|
async fn get_tool(
|
||||||
|
State(state): State<ToolState>,
|
||||||
|
Path(name): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match state.registry.get_tool(&name).await {
|
||||||
|
Some(tool) => Ok(Json(tool)),
|
||||||
|
None => Err((StatusCode::NOT_FOUND, format!("tool not found: {name}"))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Call a tool — validate params, execute SQL, log invocation.
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct CallRequest {
|
||||||
|
params: serde_json::Value,
|
||||||
|
agent: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn call_tool(
|
||||||
|
State(state): State<ToolState>,
|
||||||
|
Path(name): Path<String>,
|
||||||
|
Json(req): Json<CallRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let tool = match state.registry.get_tool(&name).await {
|
||||||
|
Some(t) => t,
|
||||||
|
None => return Err((StatusCode::NOT_FOUND, format!("tool not found: {name}"))),
|
||||||
|
};
|
||||||
|
|
||||||
|
tracing::info!("tool call: {} by agent '{}' with {:?}", name, req.agent, req.params);
|
||||||
|
|
||||||
|
// Build SQL from params
|
||||||
|
let sql = match ToolRegistry::build_sql(&tool, &req.params) {
|
||||||
|
Ok(s) => s,
|
||||||
|
Err(e) => {
|
||||||
|
state.registry.log_invocation(ToolInvocation {
|
||||||
|
id: format!("inv-{}", chrono::Utc::now().timestamp_millis()),
|
||||||
|
tool_name: name.clone(),
|
||||||
|
agent: req.agent.clone(),
|
||||||
|
params: req.params.clone(),
|
||||||
|
permission: tool.permission.clone(),
|
||||||
|
timestamp: chrono::Utc::now(),
|
||||||
|
success: false,
|
||||||
|
error: Some(e.clone()),
|
||||||
|
rows_returned: None,
|
||||||
|
}).await;
|
||||||
|
return Err((StatusCode::BAD_REQUEST, e));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Execute via query engine
|
||||||
|
let result = state.query_fn.execute(&sql).await;
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Ok((rows, row_count)) => {
|
||||||
|
state.registry.log_invocation(ToolInvocation {
|
||||||
|
id: format!("inv-{}", chrono::Utc::now().timestamp_millis()),
|
||||||
|
tool_name: name.clone(),
|
||||||
|
agent: req.agent,
|
||||||
|
params: req.params,
|
||||||
|
permission: tool.permission,
|
||||||
|
timestamp: chrono::Utc::now(),
|
||||||
|
success: true,
|
||||||
|
error: None,
|
||||||
|
rows_returned: Some(row_count),
|
||||||
|
}).await;
|
||||||
|
|
||||||
|
Ok(Json(serde_json::json!({
|
||||||
|
"tool": name,
|
||||||
|
"rows": rows,
|
||||||
|
"row_count": row_count,
|
||||||
|
"sql": sql,
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
state.registry.log_invocation(ToolInvocation {
|
||||||
|
id: format!("inv-{}", chrono::Utc::now().timestamp_millis()),
|
||||||
|
tool_name: name,
|
||||||
|
agent: req.agent,
|
||||||
|
params: req.params,
|
||||||
|
permission: tool.permission,
|
||||||
|
timestamp: chrono::Utc::now(),
|
||||||
|
success: false,
|
||||||
|
error: Some(e.clone()),
|
||||||
|
rows_returned: None,
|
||||||
|
}).await;
|
||||||
|
|
||||||
|
Err((StatusCode::INTERNAL_SERVER_ERROR, e))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct AuditQuery {
|
||||||
|
limit: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn audit_log(
|
||||||
|
State(state): State<ToolState>,
|
||||||
|
Query(q): Query<AuditQuery>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let log = state.registry.recent_audit(q.limit.unwrap_or(50)).await;
|
||||||
|
Json(log)
|
||||||
|
}
|
||||||
@@ -0,0 +1,26 @@
|
|||||||
|
[package]
|
||||||
|
name = "ingestd"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2024"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
shared = { path = "../shared" }
|
||||||
|
storaged = { path = "../storaged" }
|
||||||
|
catalogd = { path = "../catalogd" }
|
||||||
|
vectord = { path = "../vectord" }
|
||||||
|
tokio = { workspace = true }
|
||||||
|
axum = { workspace = true, features = ["multipart"] }
|
||||||
|
lopdf = { workspace = true }
|
||||||
|
serde = { workspace = true }
|
||||||
|
serde_json = { workspace = true }
|
||||||
|
tracing = { workspace = true }
|
||||||
|
arrow = { workspace = true }
|
||||||
|
parquet = { workspace = true }
|
||||||
|
bytes = { workspace = true }
|
||||||
|
sha2 = { workspace = true }
|
||||||
|
csv = { workspace = true }
|
||||||
|
chrono = { workspace = true }
|
||||||
|
object_store = { workspace = true }
|
||||||
|
tokio-postgres = { workspace = true }
|
||||||
|
mysql_async = { workspace = true }
|
||||||
|
uuid = { workspace = true }
|
||||||
@@ -0,0 +1,203 @@
|
|||||||
|
use arrow::array::{ArrayRef, Float64Array, Int64Array, StringArray, BooleanArray};
|
||||||
|
use arrow::datatypes::{DataType, Field, Schema};
|
||||||
|
use arrow::record_batch::RecordBatch;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// Inferred column type from sampling data.
|
||||||
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
|
enum InferredType {
|
||||||
|
Integer,
|
||||||
|
Float,
|
||||||
|
Boolean,
|
||||||
|
String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parse CSV bytes into Arrow RecordBatches with automatic schema detection.
|
||||||
|
/// Per ADR-010: ambiguous types default to String.
|
||||||
|
pub fn parse_csv(content: &[u8]) -> Result<(Arc<Schema>, Vec<RecordBatch>), String> {
|
||||||
|
let mut reader = csv::ReaderBuilder::new()
|
||||||
|
.flexible(true) // allow varying column counts
|
||||||
|
.trim(csv::Trim::All)
|
||||||
|
.from_reader(content);
|
||||||
|
|
||||||
|
let headers: Vec<String> = reader.headers()
|
||||||
|
.map_err(|e| format!("CSV header error: {e}"))?
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, h)| {
|
||||||
|
let h = h.trim().to_string();
|
||||||
|
if h.is_empty() { format!("column_{i}") } else { sanitize_column_name(&h) }
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let n_cols = headers.len();
|
||||||
|
if n_cols == 0 {
|
||||||
|
return Err("CSV has no columns".into());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Read all rows into string columns
|
||||||
|
let mut columns: Vec<Vec<String>> = vec![vec![]; n_cols];
|
||||||
|
let mut row_count = 0;
|
||||||
|
|
||||||
|
for result in reader.records() {
|
||||||
|
let record = result.map_err(|e| format!("CSV row error: {e}"))?;
|
||||||
|
for (i, field) in record.iter().enumerate() {
|
||||||
|
if i < n_cols {
|
||||||
|
columns[i].push(field.trim().to_string());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Pad short rows with empty strings
|
||||||
|
for col in columns.iter_mut().skip(record.len().min(n_cols)) {
|
||||||
|
col.push(String::new());
|
||||||
|
}
|
||||||
|
row_count += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
if row_count == 0 {
|
||||||
|
return Err("CSV has no data rows".into());
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!("parsed CSV: {row_count} rows × {n_cols} columns");
|
||||||
|
|
||||||
|
// Infer types by sampling (look at all values)
|
||||||
|
let types: Vec<InferredType> = columns.iter().map(|col| infer_column_type(col)).collect();
|
||||||
|
|
||||||
|
// Build Arrow schema
|
||||||
|
let fields: Vec<Field> = headers.iter().zip(types.iter()).map(|(name, typ)| {
|
||||||
|
let dt = match typ {
|
||||||
|
InferredType::Integer => DataType::Int64,
|
||||||
|
InferredType::Float => DataType::Float64,
|
||||||
|
InferredType::Boolean => DataType::Boolean,
|
||||||
|
InferredType::String => DataType::Utf8,
|
||||||
|
};
|
||||||
|
Field::new(name, dt, true) // all nullable
|
||||||
|
}).collect();
|
||||||
|
|
||||||
|
let schema = Arc::new(Schema::new(fields));
|
||||||
|
|
||||||
|
// Build arrays
|
||||||
|
let arrays: Vec<ArrayRef> = columns.iter().zip(types.iter()).map(|(col, typ)| {
|
||||||
|
match typ {
|
||||||
|
InferredType::Integer => {
|
||||||
|
let vals: Vec<Option<i64>> = col.iter().map(|v| {
|
||||||
|
if v.is_empty() { None } else { v.replace(',', "").parse().ok() }
|
||||||
|
}).collect();
|
||||||
|
Arc::new(Int64Array::from(vals)) as ArrayRef
|
||||||
|
}
|
||||||
|
InferredType::Float => {
|
||||||
|
let vals: Vec<Option<f64>> = col.iter().map(|v| {
|
||||||
|
if v.is_empty() { None }
|
||||||
|
else { v.replace(',', "").replace('$', "").replace('%', "").parse().ok() }
|
||||||
|
}).collect();
|
||||||
|
Arc::new(Float64Array::from(vals)) as ArrayRef
|
||||||
|
}
|
||||||
|
InferredType::Boolean => {
|
||||||
|
let vals: Vec<Option<bool>> = col.iter().map(|v| {
|
||||||
|
match v.to_lowercase().as_str() {
|
||||||
|
"true" | "yes" | "1" | "y" | "t" => Some(true),
|
||||||
|
"false" | "no" | "0" | "n" | "f" => Some(false),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
Arc::new(BooleanArray::from(vals)) as ArrayRef
|
||||||
|
}
|
||||||
|
InferredType::String => {
|
||||||
|
Arc::new(StringArray::from(col.clone())) as ArrayRef
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
|
||||||
|
let batch = RecordBatch::try_new(schema.clone(), arrays)
|
||||||
|
.map_err(|e| format!("RecordBatch error: {e}"))?;
|
||||||
|
|
||||||
|
Ok((schema, vec![batch]))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Infer column type from values. Conservative: defaults to String on ambiguity.
|
||||||
|
fn infer_column_type(values: &[String]) -> InferredType {
|
||||||
|
let non_empty: Vec<&str> = values.iter()
|
||||||
|
.map(|v| v.as_str())
|
||||||
|
.filter(|v| !v.is_empty() && *v != "NULL" && *v != "null" && *v != "N/A" && *v != "n/a")
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
if non_empty.is_empty() {
|
||||||
|
return InferredType::String;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check boolean
|
||||||
|
let all_bool = non_empty.iter().all(|v| {
|
||||||
|
matches!(v.to_lowercase().as_str(), "true" | "false" | "yes" | "no" | "1" | "0" | "y" | "n" | "t" | "f")
|
||||||
|
});
|
||||||
|
if all_bool && non_empty.len() >= 2 {
|
||||||
|
// Make sure it's not just all "1" and "0" which could be integers
|
||||||
|
let has_text_bool = non_empty.iter().any(|v| {
|
||||||
|
matches!(v.to_lowercase().as_str(), "true" | "false" | "yes" | "no" | "y" | "n" | "t" | "f")
|
||||||
|
});
|
||||||
|
if has_text_bool {
|
||||||
|
return InferredType::Boolean;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check integer (allow commas as thousands separator)
|
||||||
|
let int_rate = non_empty.iter()
|
||||||
|
.filter(|v| v.replace(',', "").parse::<i64>().is_ok())
|
||||||
|
.count() as f64 / non_empty.len() as f64;
|
||||||
|
if int_rate > 0.95 {
|
||||||
|
return InferredType::Integer;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check float (allow $, %, commas)
|
||||||
|
let float_rate = non_empty.iter()
|
||||||
|
.filter(|v| v.replace(',', "").replace('$', "").replace('%', "").parse::<f64>().is_ok())
|
||||||
|
.count() as f64 / non_empty.len() as f64;
|
||||||
|
if float_rate > 0.95 {
|
||||||
|
return InferredType::Float;
|
||||||
|
}
|
||||||
|
|
||||||
|
InferredType::String
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Sanitize column name for SQL compatibility.
|
||||||
|
fn sanitize_column_name(name: &str) -> String {
|
||||||
|
name.chars()
|
||||||
|
.map(|c| if c.is_alphanumeric() || c == '_' { c.to_ascii_lowercase() } else { '_' })
|
||||||
|
.collect::<String>()
|
||||||
|
.trim_matches('_')
|
||||||
|
.to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_simple_csv() {
|
||||||
|
let csv = b"Name,Age,Salary\nAlice,30,50000\nBob,25,45000\n";
|
||||||
|
let (schema, batches) = parse_csv(csv).unwrap();
|
||||||
|
assert_eq!(schema.fields().len(), 3);
|
||||||
|
assert_eq!(batches[0].num_rows(), 2);
|
||||||
|
assert_eq!(schema.field(1).data_type(), &DataType::Int64);
|
||||||
|
assert_eq!(schema.field(2).data_type(), &DataType::Int64);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_csv_with_mixed_types() {
|
||||||
|
let csv = b"id,value\n1,hello\n2,world\n3,N/A\n";
|
||||||
|
let (schema, _) = parse_csv(csv).unwrap();
|
||||||
|
assert_eq!(schema.field(0).data_type(), &DataType::Int64);
|
||||||
|
assert_eq!(schema.field(1).data_type(), &DataType::Utf8);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_csv_with_dollar_amounts() {
|
||||||
|
let csv = b"item,price\nWidget,$29.99\nGadget,$149.50\n";
|
||||||
|
let (schema, _) = parse_csv(csv).unwrap();
|
||||||
|
assert_eq!(schema.field(1).data_type(), &DataType::Float64);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn sanitize_names() {
|
||||||
|
assert_eq!(sanitize_column_name("First Name"), "first_name");
|
||||||
|
assert_eq!(sanitize_column_name("Bill Rate ($)"), "bill_rate");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,225 @@
|
|||||||
|
/// Database connector: pull tables from PostgreSQL directly into Parquet.
|
||||||
|
/// Connects, reads schema, streams rows into Arrow RecordBatches, writes Parquet.
|
||||||
|
/// No ORM, no schema definition — schema inferred from the database.
|
||||||
|
|
||||||
|
use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, StringArray};
|
||||||
|
use arrow::datatypes::{DataType, Field, Schema};
|
||||||
|
use arrow::record_batch::RecordBatch;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio_postgres::{Client, NoTls, types::Type as PgType};
|
||||||
|
|
||||||
|
/// Connection config for a database.
|
||||||
|
#[derive(Debug, Clone, serde::Deserialize)]
|
||||||
|
pub struct DbConfig {
|
||||||
|
pub host: String,
|
||||||
|
pub port: u16,
|
||||||
|
pub database: String,
|
||||||
|
pub user: String,
|
||||||
|
pub password: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl DbConfig {
|
||||||
|
pub fn connection_string(&self) -> String {
|
||||||
|
if self.password.is_empty() {
|
||||||
|
format!(
|
||||||
|
"host={} port={} dbname={} user={}",
|
||||||
|
self.host, self.port, self.database, self.user
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
format!(
|
||||||
|
"host={} port={} dbname={} user={} password={}",
|
||||||
|
self.host, self.port, self.database, self.user, self.password
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Result of a database table import.
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct DbIngestResult {
|
||||||
|
pub table_name: String,
|
||||||
|
pub rows: usize,
|
||||||
|
pub columns: usize,
|
||||||
|
pub schema_detected: Vec<String>, // column name: type pairs
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Connect to PostgreSQL and import a table as Arrow RecordBatches.
|
||||||
|
pub async fn import_postgres_table(
|
||||||
|
config: &DbConfig,
|
||||||
|
table_name: &str,
|
||||||
|
limit: Option<usize>,
|
||||||
|
) -> Result<(Arc<Schema>, Vec<RecordBatch>, DbIngestResult), String> {
|
||||||
|
// Connect
|
||||||
|
let (client, connection) = tokio_postgres::connect(&config.connection_string(), NoTls)
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("postgres connect error: {e}"))?;
|
||||||
|
|
||||||
|
// Spawn connection handler
|
||||||
|
tokio::spawn(async move {
|
||||||
|
if let Err(e) = connection.await {
|
||||||
|
tracing::error!("postgres connection error: {e}");
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
tracing::info!("connected to postgres {}:{}/{}", config.host, config.port, config.database);
|
||||||
|
|
||||||
|
// Get column info
|
||||||
|
let columns = get_table_columns(&client, table_name).await?;
|
||||||
|
if columns.is_empty() {
|
||||||
|
return Err(format!("table '{}' not found or has no columns", table_name));
|
||||||
|
}
|
||||||
|
|
||||||
|
let schema_detected: Vec<String> = columns.iter()
|
||||||
|
.map(|(name, pg_type)| format!("{}: {}", name, pg_type_to_string(pg_type)))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
tracing::info!("table '{}': {} columns", table_name, columns.len());
|
||||||
|
|
||||||
|
// Build Arrow schema
|
||||||
|
let arrow_fields: Vec<Field> = columns.iter()
|
||||||
|
.map(|(name, pg_type)| Field::new(name, pg_type_to_arrow(pg_type), true))
|
||||||
|
.collect();
|
||||||
|
let schema = Arc::new(Schema::new(arrow_fields));
|
||||||
|
|
||||||
|
// Query rows
|
||||||
|
let sql = match limit {
|
||||||
|
Some(n) => format!("SELECT * FROM \"{}\" LIMIT {}", table_name, n),
|
||||||
|
None => format!("SELECT * FROM \"{}\"", table_name),
|
||||||
|
};
|
||||||
|
|
||||||
|
let rows = client.query(&sql, &[])
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("query error: {e}"))?;
|
||||||
|
|
||||||
|
let row_count = rows.len();
|
||||||
|
tracing::info!("fetched {} rows from '{}'", row_count, table_name);
|
||||||
|
|
||||||
|
if row_count == 0 {
|
||||||
|
return Ok((schema.clone(), vec![], DbIngestResult {
|
||||||
|
table_name: table_name.to_string(),
|
||||||
|
rows: 0,
|
||||||
|
columns: columns.len(),
|
||||||
|
schema_detected,
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Convert to Arrow arrays
|
||||||
|
let mut arrays: Vec<ArrayRef> = Vec::new();
|
||||||
|
|
||||||
|
for (col_idx, (_, pg_type)) in columns.iter().enumerate() {
|
||||||
|
let array = rows_to_arrow_column(&rows, col_idx, pg_type)?;
|
||||||
|
arrays.push(array);
|
||||||
|
}
|
||||||
|
|
||||||
|
let batch = RecordBatch::try_new(schema.clone(), arrays)
|
||||||
|
.map_err(|e| format!("RecordBatch error: {e}"))?;
|
||||||
|
|
||||||
|
let result = DbIngestResult {
|
||||||
|
table_name: table_name.to_string(),
|
||||||
|
rows: row_count,
|
||||||
|
columns: columns.len(),
|
||||||
|
schema_detected,
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok((schema, vec![batch], result))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all tables in the database.
|
||||||
|
pub async fn list_postgres_tables(config: &DbConfig) -> Result<Vec<String>, String> {
|
||||||
|
let (client, connection) = tokio_postgres::connect(&config.connection_string(), NoTls)
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("connect error: {e}"))?;
|
||||||
|
|
||||||
|
tokio::spawn(async move { let _ = connection.await; });
|
||||||
|
|
||||||
|
let rows = client.query(
|
||||||
|
"SELECT table_name FROM information_schema.tables WHERE table_schema = 'public' ORDER BY table_name",
|
||||||
|
&[],
|
||||||
|
).await.map_err(|e| format!("query error: {e}"))?;
|
||||||
|
|
||||||
|
Ok(rows.iter().map(|r| r.get::<_, String>(0)).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get column names and types for a table.
|
||||||
|
async fn get_table_columns(client: &Client, table_name: &str) -> Result<Vec<(String, PgType)>, String> {
|
||||||
|
// Use a dummy query to get column types
|
||||||
|
let stmt = client.prepare(&format!("SELECT * FROM \"{}\" LIMIT 0", table_name))
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("prepare error for '{}': {e}", table_name))?;
|
||||||
|
|
||||||
|
Ok(stmt.columns().iter()
|
||||||
|
.map(|col| (col.name().to_string(), col.type_().clone()))
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Convert a Postgres type to Arrow DataType.
|
||||||
|
fn pg_type_to_arrow(pg_type: &PgType) -> DataType {
|
||||||
|
match *pg_type {
|
||||||
|
PgType::BOOL => DataType::Boolean,
|
||||||
|
PgType::INT2 | PgType::INT4 => DataType::Int32,
|
||||||
|
PgType::INT8 | PgType::OID => DataType::Int64,
|
||||||
|
PgType::FLOAT4 | PgType::FLOAT8 | PgType::NUMERIC => DataType::Float64,
|
||||||
|
_ => DataType::Utf8, // everything else → string (safe default per ADR-010)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn pg_type_to_string(pg_type: &PgType) -> String {
|
||||||
|
format!("{}", pg_type)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Convert a column of Postgres rows to an Arrow array.
|
||||||
|
fn rows_to_arrow_column(
|
||||||
|
rows: &[tokio_postgres::Row],
|
||||||
|
col_idx: usize,
|
||||||
|
pg_type: &PgType,
|
||||||
|
) -> Result<ArrayRef, String> {
|
||||||
|
match *pg_type {
|
||||||
|
PgType::BOOL => {
|
||||||
|
let vals: Vec<Option<bool>> = rows.iter()
|
||||||
|
.map(|r| r.try_get(col_idx).ok())
|
||||||
|
.collect();
|
||||||
|
Ok(Arc::new(BooleanArray::from(vals)))
|
||||||
|
}
|
||||||
|
PgType::INT2 => {
|
||||||
|
let vals: Vec<Option<i32>> = rows.iter()
|
||||||
|
.map(|r| r.try_get::<_, i16>(col_idx).ok().map(|v| v as i32))
|
||||||
|
.collect();
|
||||||
|
Ok(Arc::new(Int32Array::from(vals)))
|
||||||
|
}
|
||||||
|
PgType::INT4 => {
|
||||||
|
let vals: Vec<Option<i32>> = rows.iter()
|
||||||
|
.map(|r| r.try_get(col_idx).ok())
|
||||||
|
.collect();
|
||||||
|
Ok(Arc::new(Int32Array::from(vals)))
|
||||||
|
}
|
||||||
|
PgType::INT8 | PgType::OID => {
|
||||||
|
let vals: Vec<Option<i64>> = rows.iter()
|
||||||
|
.map(|r| r.try_get(col_idx).ok())
|
||||||
|
.collect();
|
||||||
|
Ok(Arc::new(Int64Array::from(vals)))
|
||||||
|
}
|
||||||
|
PgType::FLOAT4 => {
|
||||||
|
let vals: Vec<Option<f64>> = rows.iter()
|
||||||
|
.map(|r| r.try_get::<_, f32>(col_idx).ok().map(|v| v as f64))
|
||||||
|
.collect();
|
||||||
|
Ok(Arc::new(Float64Array::from(vals)))
|
||||||
|
}
|
||||||
|
PgType::FLOAT8 => {
|
||||||
|
let vals: Vec<Option<f64>> = rows.iter()
|
||||||
|
.map(|r| r.try_get(col_idx).ok())
|
||||||
|
.collect();
|
||||||
|
Ok(Arc::new(Float64Array::from(vals)))
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
// Default: try to get as string
|
||||||
|
let vals: Vec<Option<String>> = rows.iter()
|
||||||
|
.map(|r| {
|
||||||
|
r.try_get::<_, String>(col_idx).ok()
|
||||||
|
.or_else(|| r.try_get::<_, serde_json::Value>(col_idx).ok().map(|v| v.to_string()))
|
||||||
|
.or_else(|| Some("".to_string()))
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
Ok(Arc::new(StringArray::from(vals)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,103 @@
|
|||||||
|
use sha2::{Digest, Sha256};
|
||||||
|
|
||||||
|
/// Detected file type from content inspection.
|
||||||
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
|
pub enum FileType {
|
||||||
|
Csv,
|
||||||
|
Json,
|
||||||
|
NdJson, // newline-delimited JSON
|
||||||
|
Pdf,
|
||||||
|
Text, // plain text, SMS logs, etc.
|
||||||
|
Unknown,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Detect file type from filename extension and content sniffing.
|
||||||
|
pub fn detect_file_type(filename: &str, content: &[u8]) -> FileType {
|
||||||
|
// Extension-based first
|
||||||
|
let lower = filename.to_lowercase();
|
||||||
|
if lower.ends_with(".csv") || lower.ends_with(".tsv") {
|
||||||
|
return FileType::Csv;
|
||||||
|
}
|
||||||
|
if lower.ends_with(".json") {
|
||||||
|
// Check if it's newline-delimited JSON
|
||||||
|
if content.iter().take(4096).filter(|&&b| b == b'\n').count() > 2 {
|
||||||
|
let first_line = content.split(|&b| b == b'\n').next().unwrap_or(b"");
|
||||||
|
if first_line.starts_with(b"{") {
|
||||||
|
return FileType::NdJson;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return FileType::Json;
|
||||||
|
}
|
||||||
|
if lower.ends_with(".ndjson") || lower.ends_with(".jsonl") {
|
||||||
|
return FileType::NdJson;
|
||||||
|
}
|
||||||
|
if lower.ends_with(".pdf") {
|
||||||
|
return FileType::Pdf;
|
||||||
|
}
|
||||||
|
if lower.ends_with(".txt") || lower.ends_with(".log") || lower.ends_with(".sms") {
|
||||||
|
return FileType::Text;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Content sniffing fallback
|
||||||
|
if content.starts_with(b"%PDF") {
|
||||||
|
return FileType::Pdf;
|
||||||
|
}
|
||||||
|
if content.starts_with(b"[") || content.starts_with(b"{") {
|
||||||
|
return FileType::Json;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if it looks like CSV (has commas and newlines in first chunk)
|
||||||
|
let sample = &content[..content.len().min(4096)];
|
||||||
|
let comma_count = sample.iter().filter(|&&b| b == b',').count();
|
||||||
|
let newline_count = sample.iter().filter(|&&b| b == b'\n').count();
|
||||||
|
if comma_count > 3 && newline_count > 1 {
|
||||||
|
return FileType::Csv;
|
||||||
|
}
|
||||||
|
|
||||||
|
// If it's valid UTF-8, treat as text
|
||||||
|
if std::str::from_utf8(sample).is_ok() {
|
||||||
|
return FileType::Text;
|
||||||
|
}
|
||||||
|
|
||||||
|
FileType::Unknown
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Compute SHA-256 hash of content for deduplication.
|
||||||
|
pub fn content_hash(content: &[u8]) -> String {
|
||||||
|
let mut hasher = Sha256::new();
|
||||||
|
hasher.update(content);
|
||||||
|
format!("{:x}", hasher.finalize())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn detect_csv_by_extension() {
|
||||||
|
assert_eq!(detect_file_type("data.csv", b"a,b,c\n1,2,3"), FileType::Csv);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn detect_json_by_extension() {
|
||||||
|
assert_eq!(detect_file_type("data.json", b"[{\"a\":1}]"), FileType::Json);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn detect_pdf_by_magic() {
|
||||||
|
assert_eq!(detect_file_type("unknown", b"%PDF-1.4 blah"), FileType::Pdf);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn detect_csv_by_content() {
|
||||||
|
let csv = b"name,age,city\nAlice,30,NYC\nBob,25,LA\n";
|
||||||
|
assert_eq!(detect_file_type("unknown.dat", csv), FileType::Csv);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn content_hash_deterministic() {
|
||||||
|
let h1 = content_hash(b"hello world");
|
||||||
|
let h2 = content_hash(b"hello world");
|
||||||
|
assert_eq!(h1, h2);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,167 @@
|
|||||||
|
use arrow::array::{ArrayRef, Float64Array, Int64Array, StringArray, BooleanArray};
|
||||||
|
use arrow::datatypes::{DataType, Field, Schema};
|
||||||
|
use arrow::record_batch::RecordBatch;
|
||||||
|
use std::collections::BTreeMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// Parse JSON (array of objects or newline-delimited) into Arrow RecordBatches.
|
||||||
|
/// Nested objects are flattened: {"a": {"b": 1}} → column "a_b".
|
||||||
|
pub fn parse_json(content: &[u8]) -> Result<(Arc<Schema>, Vec<RecordBatch>), String> {
|
||||||
|
let text = std::str::from_utf8(content).map_err(|e| format!("invalid UTF-8: {e}"))?;
|
||||||
|
let text = text.trim();
|
||||||
|
|
||||||
|
// Parse into Vec<Map>
|
||||||
|
let rows: Vec<serde_json::Map<String, serde_json::Value>> = if text.starts_with('[') {
|
||||||
|
// JSON array
|
||||||
|
let arr: Vec<serde_json::Value> = serde_json::from_str(text)
|
||||||
|
.map_err(|e| format!("JSON parse error: {e}"))?;
|
||||||
|
arr.into_iter()
|
||||||
|
.filter_map(|v| v.as_object().cloned())
|
||||||
|
.collect()
|
||||||
|
} else {
|
||||||
|
// Newline-delimited JSON
|
||||||
|
text.lines()
|
||||||
|
.filter(|l| !l.trim().is_empty())
|
||||||
|
.map(|l| {
|
||||||
|
let v: serde_json::Value = serde_json::from_str(l)
|
||||||
|
.map_err(|e| format!("NDJSON line parse error: {e}"))?;
|
||||||
|
v.as_object().cloned().ok_or_else(|| "NDJSON line is not an object".into())
|
||||||
|
})
|
||||||
|
.collect::<Result<Vec<_>, String>>()?
|
||||||
|
};
|
||||||
|
|
||||||
|
if rows.is_empty() {
|
||||||
|
return Err("JSON has no records".into());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Flatten and collect all columns
|
||||||
|
let mut columns: BTreeMap<String, Vec<serde_json::Value>> = BTreeMap::new();
|
||||||
|
let n_rows = rows.len();
|
||||||
|
|
||||||
|
for (row_idx, row) in rows.iter().enumerate() {
|
||||||
|
let flat = flatten_object(row, "");
|
||||||
|
// Pad existing columns that aren't in this row
|
||||||
|
for key in columns.keys().cloned().collect::<Vec<_>>() {
|
||||||
|
if !flat.contains_key(&key) {
|
||||||
|
columns.get_mut(&key).unwrap().push(serde_json::Value::Null);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Add new columns or append values
|
||||||
|
for (key, val) in flat {
|
||||||
|
let col = columns.entry(key).or_insert_with(Vec::new);
|
||||||
|
// Pad with nulls if this column is new and we've seen prior rows
|
||||||
|
while col.len() < row_idx {
|
||||||
|
col.push(serde_json::Value::Null);
|
||||||
|
}
|
||||||
|
col.push(val);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Pad short columns
|
||||||
|
for col in columns.values_mut() {
|
||||||
|
while col.len() < n_rows {
|
||||||
|
col.push(serde_json::Value::Null);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!("parsed JSON: {n_rows} rows × {} columns", columns.len());
|
||||||
|
|
||||||
|
// Infer types and build schema
|
||||||
|
let mut fields = Vec::new();
|
||||||
|
let mut arrays: Vec<ArrayRef> = Vec::new();
|
||||||
|
|
||||||
|
for (name, values) in &columns {
|
||||||
|
let (dt, array) = json_column_to_arrow(values);
|
||||||
|
fields.push(Field::new(name, dt, true));
|
||||||
|
arrays.push(array);
|
||||||
|
}
|
||||||
|
|
||||||
|
let schema = Arc::new(Schema::new(fields));
|
||||||
|
let batch = RecordBatch::try_new(schema.clone(), arrays)
|
||||||
|
.map_err(|e| format!("RecordBatch error: {e}"))?;
|
||||||
|
|
||||||
|
Ok((schema, vec![batch]))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Flatten nested JSON object: {"a": {"b": 1}} → {"a_b": 1}
|
||||||
|
fn flatten_object(obj: &serde_json::Map<String, serde_json::Value>, prefix: &str) -> BTreeMap<String, serde_json::Value> {
|
||||||
|
let mut result = BTreeMap::new();
|
||||||
|
for (key, val) in obj {
|
||||||
|
let full_key = if prefix.is_empty() { key.clone() } else { format!("{prefix}_{key}") };
|
||||||
|
match val {
|
||||||
|
serde_json::Value::Object(inner) => {
|
||||||
|
result.extend(flatten_object(inner, &full_key));
|
||||||
|
}
|
||||||
|
other => {
|
||||||
|
result.insert(full_key, other.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Convert a column of JSON values to an Arrow array.
|
||||||
|
fn json_column_to_arrow(values: &[serde_json::Value]) -> (DataType, ArrayRef) {
|
||||||
|
// Check if all non-null values are the same type
|
||||||
|
let non_null: Vec<&serde_json::Value> = values.iter()
|
||||||
|
.filter(|v| !v.is_null())
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
if non_null.is_empty() {
|
||||||
|
return (DataType::Utf8, Arc::new(StringArray::from(vec![None::<&str>; values.len()])));
|
||||||
|
}
|
||||||
|
|
||||||
|
let all_bool = non_null.iter().all(|v| v.is_boolean());
|
||||||
|
let all_i64 = non_null.iter().all(|v| v.is_i64() || v.is_u64());
|
||||||
|
let all_f64 = non_null.iter().all(|v| v.is_number());
|
||||||
|
|
||||||
|
if all_bool {
|
||||||
|
let arr: Vec<Option<bool>> = values.iter().map(|v| v.as_bool()).collect();
|
||||||
|
(DataType::Boolean, Arc::new(BooleanArray::from(arr)))
|
||||||
|
} else if all_i64 {
|
||||||
|
let arr: Vec<Option<i64>> = values.iter().map(|v| v.as_i64()).collect();
|
||||||
|
(DataType::Int64, Arc::new(Int64Array::from(arr)))
|
||||||
|
} else if all_f64 {
|
||||||
|
let arr: Vec<Option<f64>> = values.iter().map(|v| v.as_f64()).collect();
|
||||||
|
(DataType::Float64, Arc::new(Float64Array::from(arr)))
|
||||||
|
} else {
|
||||||
|
// Default to string
|
||||||
|
let arr: Vec<Option<String>> = values.iter().map(|v| {
|
||||||
|
match v {
|
||||||
|
serde_json::Value::Null => None,
|
||||||
|
serde_json::Value::String(s) => Some(s.clone()),
|
||||||
|
other => Some(other.to_string()),
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
(DataType::Utf8, Arc::new(StringArray::from(arr)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_json_array() {
|
||||||
|
let json = br#"[{"name":"Alice","age":30},{"name":"Bob","age":25}]"#;
|
||||||
|
let (schema, batches) = parse_json(json).unwrap();
|
||||||
|
assert_eq!(batches[0].num_rows(), 2);
|
||||||
|
assert!(schema.field_with_name("age").unwrap().data_type() == &DataType::Int64);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_ndjson() {
|
||||||
|
let json = b"{\"x\":1}\n{\"x\":2}\n{\"x\":3}\n";
|
||||||
|
let (_, batches) = parse_json(json).unwrap();
|
||||||
|
assert_eq!(batches[0].num_rows(), 3);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn flatten_nested() {
|
||||||
|
let json = br#"[{"user":{"name":"Alice","address":{"city":"NYC"}},"score":9.5}]"#;
|
||||||
|
let (schema, _) = parse_json(json).unwrap();
|
||||||
|
let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
|
||||||
|
assert!(names.contains(&"user_name"));
|
||||||
|
assert!(names.contains(&"user_address_city"));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,12 @@
|
|||||||
|
pub mod db_ingest;
|
||||||
|
pub mod my_stream;
|
||||||
|
pub mod pg_stream;
|
||||||
|
pub mod schedule;
|
||||||
|
pub mod detect;
|
||||||
|
pub mod csv_ingest;
|
||||||
|
pub mod json_ingest;
|
||||||
|
pub mod pdf_ingest;
|
||||||
|
pub mod pipeline;
|
||||||
|
pub mod schema_evolution;
|
||||||
|
pub mod service;
|
||||||
|
pub mod watcher;
|
||||||
@@ -0,0 +1,410 @@
|
|||||||
|
/// Streaming MySQL ingest.
|
||||||
|
///
|
||||||
|
/// Mirrors `pg_stream` for MySQL sources. Same OFFSET-paginated strategy,
|
||||||
|
/// same Parquet-streaming shape. Uses `mysql_async` (pure-rust) so we
|
||||||
|
/// don't need a C client library at build time.
|
||||||
|
///
|
||||||
|
/// Type mapping follows ADR-010 (default to string on ambiguity):
|
||||||
|
/// - Booleans (TINYINT(1)) and integer types map to Arrow Int32/Int64
|
||||||
|
/// - Floating point and decimals → Float64
|
||||||
|
/// - Everything else (text, varchar, json, date, timestamp) → Utf8
|
||||||
|
///
|
||||||
|
/// What's deliberately not supported (yet):
|
||||||
|
/// - TLS connections — `minimal` feature is plain TCP only. Upgrade
|
||||||
|
/// when a tenant actually needs it.
|
||||||
|
/// - Keyset pagination — OFFSET scans are O(N²) at multi-million-row
|
||||||
|
/// scale; same caveat as `pg_stream`.
|
||||||
|
/// - BINARY/BLOB columns — currently rendered as base64 or empty string
|
||||||
|
/// via Display-fallback.
|
||||||
|
|
||||||
|
use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray};
|
||||||
|
use arrow::datatypes::{DataType, Field, Schema};
|
||||||
|
use mysql_async::prelude::*;
|
||||||
|
use mysql_async::{Conn, Opts, Row, Value};
|
||||||
|
use parquet::arrow::ArrowWriter;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// Request shape for MySQL streaming ingest.
|
||||||
|
#[derive(Debug, Clone, serde::Deserialize)]
|
||||||
|
pub struct MyStreamRequest {
|
||||||
|
/// mysql://user:pass@host:port/db
|
||||||
|
pub dsn: String,
|
||||||
|
pub table: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub dataset_name: Option<String>,
|
||||||
|
/// Rows per fetch. Default 10_000.
|
||||||
|
#[serde(default)]
|
||||||
|
pub batch_size: Option<usize>,
|
||||||
|
/// Column to ORDER BY for stable pagination. If omitted, the first
|
||||||
|
/// column returned by the schema probe is used.
|
||||||
|
#[serde(default)]
|
||||||
|
pub order_by: Option<String>,
|
||||||
|
/// Hard cap on total rows (for sampling / previews).
|
||||||
|
#[serde(default)]
|
||||||
|
pub limit: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct MyStreamResult {
|
||||||
|
pub table: String,
|
||||||
|
pub rows: usize,
|
||||||
|
pub batches: usize,
|
||||||
|
pub columns: usize,
|
||||||
|
pub schema: Vec<String>,
|
||||||
|
pub parquet_bytes: u64,
|
||||||
|
pub duration_secs: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parsed DSN pieces. Kept local (rather than reusing pg's DbConfig) so
|
||||||
|
/// the MySQL connector doesn't depend on the PG path.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct MyConfig {
|
||||||
|
pub host: String,
|
||||||
|
pub port: u16,
|
||||||
|
pub user: String,
|
||||||
|
pub password: String,
|
||||||
|
pub database: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl MyConfig {
|
||||||
|
/// Build a mysql_async `Opts` from the parsed config.
|
||||||
|
pub fn to_opts(&self) -> Opts {
|
||||||
|
let url = if self.password.is_empty() {
|
||||||
|
format!(
|
||||||
|
"mysql://{}@{}:{}/{}",
|
||||||
|
percent(&self.user), percent(&self.host), self.port, percent(&self.database),
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
format!(
|
||||||
|
"mysql://{}:{}@{}:{}/{}",
|
||||||
|
percent(&self.user), percent(&self.password),
|
||||||
|
percent(&self.host), self.port, percent(&self.database),
|
||||||
|
)
|
||||||
|
};
|
||||||
|
Opts::from_url(&url).expect("MyConfig.to_opts: rebuilt URL should parse")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Minimal URL-encoder for the few characters that commonly appear in
|
||||||
|
/// MySQL passwords. mysql_async's URL parser expects valid URL-encoded
|
||||||
|
/// components.
|
||||||
|
fn percent(s: &str) -> String {
|
||||||
|
let mut out = String::with_capacity(s.len());
|
||||||
|
for c in s.chars() {
|
||||||
|
match c {
|
||||||
|
'@' | ':' | '/' | '?' | '#' | '[' | ']' | ' ' | '%' => {
|
||||||
|
out.push_str(&format!("%{:02X}", c as u32));
|
||||||
|
}
|
||||||
|
_ => out.push(c),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parse a mysql:// DSN.
|
||||||
|
/// Supports: mysql://[user[:password]@]host[:port][/db]
|
||||||
|
pub fn parse_dsn(dsn: &str) -> Result<MyConfig, String> {
|
||||||
|
let rest = dsn
|
||||||
|
.strip_prefix("mysql://")
|
||||||
|
.ok_or_else(|| "DSN must start with mysql://".to_string())?;
|
||||||
|
|
||||||
|
let (auth_host, database) = match rest.split_once('/') {
|
||||||
|
Some((ah, db)) => (ah, db.split('?').next().unwrap_or(db).to_string()),
|
||||||
|
None => (rest, String::new()),
|
||||||
|
};
|
||||||
|
|
||||||
|
let (userpass, hostport) = match auth_host.rsplit_once('@') {
|
||||||
|
Some((up, hp)) => (Some(up), hp),
|
||||||
|
None => (None, auth_host),
|
||||||
|
};
|
||||||
|
|
||||||
|
let (user, password) = match userpass {
|
||||||
|
Some(up) => match up.split_once(':') {
|
||||||
|
Some((u, p)) => (u.to_string(), p.to_string()),
|
||||||
|
None => (up.to_string(), String::new()),
|
||||||
|
},
|
||||||
|
None => ("root".to_string(), String::new()),
|
||||||
|
};
|
||||||
|
|
||||||
|
let (host, port) = match hostport.rsplit_once(':') {
|
||||||
|
Some((h, p)) => (
|
||||||
|
h.to_string(),
|
||||||
|
p.parse::<u16>().map_err(|_| format!("invalid port in DSN: {p}"))?,
|
||||||
|
),
|
||||||
|
None => (hostport.to_string(), 3306),
|
||||||
|
};
|
||||||
|
|
||||||
|
if host.is_empty() {
|
||||||
|
return Err("DSN has no host".into());
|
||||||
|
}
|
||||||
|
if database.is_empty() {
|
||||||
|
return Err("DSN has no database path (mysql://... /db)".into());
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(MyConfig { host, port, user, password, database })
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Stream a MySQL table as Parquet bytes.
|
||||||
|
pub async fn stream_table_to_parquet(
|
||||||
|
req: &MyStreamRequest,
|
||||||
|
) -> Result<(bytes::Bytes, MyStreamResult), String> {
|
||||||
|
let t0 = std::time::Instant::now();
|
||||||
|
let config = parse_dsn(&req.dsn)?;
|
||||||
|
let batch_size = req.batch_size.unwrap_or(10_000).max(1);
|
||||||
|
|
||||||
|
let mut conn = Conn::new(config.to_opts()).await
|
||||||
|
.map_err(|e| format!("mysql connect: {e}"))?;
|
||||||
|
|
||||||
|
// Probe columns via information_schema — gives real type names that
|
||||||
|
// we can map to Arrow dtypes without needing to fetch a row first.
|
||||||
|
let columns = probe_columns(&mut conn, &config.database, &req.table).await?;
|
||||||
|
if columns.is_empty() {
|
||||||
|
return Err(format!("table '{}' not found or has no columns", req.table));
|
||||||
|
}
|
||||||
|
|
||||||
|
let arrow_fields: Vec<Field> = columns
|
||||||
|
.iter()
|
||||||
|
.map(|(name, ty)| Field::new(name, mysql_type_to_arrow(ty), true))
|
||||||
|
.collect();
|
||||||
|
let schema = Arc::new(Schema::new(arrow_fields));
|
||||||
|
let schema_report: Vec<String> = columns
|
||||||
|
.iter()
|
||||||
|
.map(|(n, t)| format!("{}:{}", n, t))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let order_col = req.order_by.clone().unwrap_or_else(|| columns[0].0.clone());
|
||||||
|
|
||||||
|
let mut buf: Vec<u8> = Vec::with_capacity(1024 * 1024);
|
||||||
|
let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)
|
||||||
|
.map_err(|e| format!("ArrowWriter init: {e}"))?;
|
||||||
|
|
||||||
|
let mut total_rows: usize = 0;
|
||||||
|
let mut batch_count: usize = 0;
|
||||||
|
let row_cap = req.limit.unwrap_or(usize::MAX);
|
||||||
|
|
||||||
|
loop {
|
||||||
|
let remaining = row_cap.saturating_sub(total_rows);
|
||||||
|
if remaining == 0 { break; }
|
||||||
|
let fetch = remaining.min(batch_size);
|
||||||
|
|
||||||
|
// Backticks are MySQL's identifier quote. Forbid backticks in
|
||||||
|
// table/column names to prevent injection — neither pg nor us
|
||||||
|
// should be accepting those anyway.
|
||||||
|
if req.table.contains('`') || order_col.contains('`') {
|
||||||
|
return Err("table or order_by column contains backticks — refused".into());
|
||||||
|
}
|
||||||
|
let sql = format!(
|
||||||
|
"SELECT * FROM `{}` ORDER BY `{}` LIMIT {} OFFSET {}",
|
||||||
|
req.table, order_col, fetch, total_rows,
|
||||||
|
);
|
||||||
|
|
||||||
|
let rows: Vec<Row> = conn.query(&sql).await
|
||||||
|
.map_err(|e| format!("fetch batch at offset {total_rows}: {e}"))?;
|
||||||
|
|
||||||
|
if rows.is_empty() { break; }
|
||||||
|
let n = rows.len();
|
||||||
|
|
||||||
|
let arrays: Vec<ArrayRef> = columns
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(idx, (_, ty))| rows_to_column(&rows, idx, ty))
|
||||||
|
.collect::<Result<_, _>>()?;
|
||||||
|
let batch = RecordBatch::try_new(schema.clone(), arrays)
|
||||||
|
.map_err(|e| format!("RecordBatch: {e}"))?;
|
||||||
|
writer.write(&batch).map_err(|e| format!("ArrowWriter::write: {e}"))?;
|
||||||
|
|
||||||
|
total_rows += n;
|
||||||
|
batch_count += 1;
|
||||||
|
tracing::info!(
|
||||||
|
"mysql stream '{}': fetched batch {} ({} rows, total {})",
|
||||||
|
req.table, batch_count, n, total_rows,
|
||||||
|
);
|
||||||
|
|
||||||
|
if n < fetch { break; }
|
||||||
|
}
|
||||||
|
|
||||||
|
writer.close().map_err(|e| format!("ArrowWriter::close: {e}"))?;
|
||||||
|
conn.disconnect().await.ok();
|
||||||
|
|
||||||
|
let result = MyStreamResult {
|
||||||
|
table: req.table.clone(),
|
||||||
|
rows: total_rows,
|
||||||
|
batches: batch_count,
|
||||||
|
columns: columns.len(),
|
||||||
|
schema: schema_report,
|
||||||
|
parquet_bytes: buf.len() as u64,
|
||||||
|
duration_secs: t0.elapsed().as_secs_f32(),
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok((bytes::Bytes::from(buf), result))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn probe_columns(
|
||||||
|
conn: &mut Conn,
|
||||||
|
schema: &str,
|
||||||
|
table: &str,
|
||||||
|
) -> Result<Vec<(String, String)>, String> {
|
||||||
|
let sql = format!(
|
||||||
|
"SELECT COLUMN_NAME, DATA_TYPE FROM information_schema.columns \
|
||||||
|
WHERE table_schema = '{}' AND table_name = '{}' \
|
||||||
|
ORDER BY ORDINAL_POSITION",
|
||||||
|
schema.replace('\'', "''"),
|
||||||
|
table.replace('\'', "''"),
|
||||||
|
);
|
||||||
|
let rows: Vec<(String, String)> = conn.query(&sql).await
|
||||||
|
.map_err(|e| format!("probe columns: {e}"))?;
|
||||||
|
Ok(rows)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// MySQL data type string → Arrow DataType. Conservative: anything we
|
||||||
|
/// don't recognize becomes Utf8 (ADR-010).
|
||||||
|
fn mysql_type_to_arrow(ty: &str) -> DataType {
|
||||||
|
match ty.to_lowercase().as_str() {
|
||||||
|
"tinyint" | "smallint" | "mediumint" | "int" | "integer" => DataType::Int32,
|
||||||
|
"bigint" => DataType::Int64,
|
||||||
|
"float" | "double" | "decimal" | "numeric" | "real" => DataType::Float64,
|
||||||
|
"bit" | "bool" | "boolean" => DataType::Boolean,
|
||||||
|
_ => DataType::Utf8,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Convert a single column slice of MySQL rows into an Arrow array.
|
||||||
|
fn rows_to_column(
|
||||||
|
rows: &[Row],
|
||||||
|
idx: usize,
|
||||||
|
ty: &str,
|
||||||
|
) -> Result<ArrayRef, String> {
|
||||||
|
let arrow_ty = mysql_type_to_arrow(ty);
|
||||||
|
match arrow_ty {
|
||||||
|
DataType::Boolean => {
|
||||||
|
let v: Vec<Option<bool>> = rows.iter().map(|r| cell_as_bool(r, idx)).collect();
|
||||||
|
Ok(Arc::new(BooleanArray::from(v)))
|
||||||
|
}
|
||||||
|
DataType::Int32 => {
|
||||||
|
let v: Vec<Option<i32>> = rows.iter().map(|r| cell_as_i64(r, idx).map(|n| n as i32)).collect();
|
||||||
|
Ok(Arc::new(Int32Array::from(v)))
|
||||||
|
}
|
||||||
|
DataType::Int64 => {
|
||||||
|
let v: Vec<Option<i64>> = rows.iter().map(|r| cell_as_i64(r, idx)).collect();
|
||||||
|
Ok(Arc::new(Int64Array::from(v)))
|
||||||
|
}
|
||||||
|
DataType::Float64 => {
|
||||||
|
let v: Vec<Option<f64>> = rows.iter().map(|r| cell_as_f64(r, idx)).collect();
|
||||||
|
Ok(Arc::new(Float64Array::from(v)))
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
let v: Vec<Option<String>> = rows.iter().map(|r| cell_as_string(r, idx)).collect();
|
||||||
|
Ok(Arc::new(StringArray::from(v)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cell(r: &Row, idx: usize) -> &Value { r.as_ref(idx).unwrap_or(&Value::NULL) }
|
||||||
|
|
||||||
|
fn cell_as_bool(r: &Row, idx: usize) -> Option<bool> {
|
||||||
|
match cell(r, idx) {
|
||||||
|
Value::NULL => None,
|
||||||
|
Value::Int(n) => Some(*n != 0),
|
||||||
|
Value::UInt(n) => Some(*n != 0),
|
||||||
|
Value::Bytes(b) => std::str::from_utf8(b).ok().and_then(|s| {
|
||||||
|
match s.to_ascii_lowercase().as_str() {
|
||||||
|
"true" | "1" | "y" | "yes" => Some(true),
|
||||||
|
"false" | "0" | "n" | "no" => Some(false),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cell_as_i64(r: &Row, idx: usize) -> Option<i64> {
|
||||||
|
match cell(r, idx) {
|
||||||
|
Value::NULL => None,
|
||||||
|
Value::Int(n) => Some(*n),
|
||||||
|
Value::UInt(n) => i64::try_from(*n).ok(),
|
||||||
|
Value::Float(f) => Some(*f as i64),
|
||||||
|
Value::Double(f) => Some(*f as i64),
|
||||||
|
Value::Bytes(b) => std::str::from_utf8(b).ok().and_then(|s| s.parse().ok()),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cell_as_f64(r: &Row, idx: usize) -> Option<f64> {
|
||||||
|
match cell(r, idx) {
|
||||||
|
Value::NULL => None,
|
||||||
|
Value::Int(n) => Some(*n as f64),
|
||||||
|
Value::UInt(n) => Some(*n as f64),
|
||||||
|
Value::Float(f) => Some(*f as f64),
|
||||||
|
Value::Double(f) => Some(*f),
|
||||||
|
Value::Bytes(b) => std::str::from_utf8(b).ok().and_then(|s| s.parse().ok()),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cell_as_string(r: &Row, idx: usize) -> Option<String> {
|
||||||
|
match cell(r, idx) {
|
||||||
|
Value::NULL => None,
|
||||||
|
Value::Bytes(b) => Some(String::from_utf8_lossy(b).into_owned()),
|
||||||
|
Value::Int(n) => Some(n.to_string()),
|
||||||
|
Value::UInt(n) => Some(n.to_string()),
|
||||||
|
Value::Float(f) => Some(f.to_string()),
|
||||||
|
Value::Double(f) => Some(f.to_string()),
|
||||||
|
Value::Date(y, mo, d, h, mi, s, _us) => {
|
||||||
|
Some(format!("{:04}-{:02}-{:02} {:02}:{:02}:{:02}", y, mo, d, h, mi, s))
|
||||||
|
}
|
||||||
|
Value::Time(neg, days, h, mi, s, _us) => {
|
||||||
|
let sign = if *neg { "-" } else { "" };
|
||||||
|
Some(format!("{}{}d {:02}:{:02}:{:02}", sign, days, h, mi, s))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_dsn_full() {
|
||||||
|
let c = parse_dsn("mysql://daisy:secret@db.example.com:3307/my_db").unwrap();
|
||||||
|
assert_eq!(c.host, "db.example.com");
|
||||||
|
assert_eq!(c.port, 3307);
|
||||||
|
assert_eq!(c.user, "daisy");
|
||||||
|
assert_eq!(c.password, "secret");
|
||||||
|
assert_eq!(c.database, "my_db");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_dsn_default_port() {
|
||||||
|
let c = parse_dsn("mysql://root@localhost/shop").unwrap();
|
||||||
|
assert_eq!(c.port, 3306);
|
||||||
|
assert_eq!(c.user, "root");
|
||||||
|
assert_eq!(c.password, "");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_dsn_no_auth() {
|
||||||
|
let c = parse_dsn("mysql://127.0.0.1:3306/analytics").unwrap();
|
||||||
|
assert_eq!(c.user, "root");
|
||||||
|
assert_eq!(c.host, "127.0.0.1");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_dsn_rejects_non_mysql() {
|
||||||
|
assert!(parse_dsn("postgresql://host/db").is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_dsn_requires_db() {
|
||||||
|
assert!(parse_dsn("mysql://localhost:3306").is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn type_map_basics() {
|
||||||
|
assert_eq!(mysql_type_to_arrow("int"), DataType::Int32);
|
||||||
|
assert_eq!(mysql_type_to_arrow("BIGINT"), DataType::Int64);
|
||||||
|
assert_eq!(mysql_type_to_arrow("decimal"), DataType::Float64);
|
||||||
|
assert_eq!(mysql_type_to_arrow("varchar"), DataType::Utf8);
|
||||||
|
assert_eq!(mysql_type_to_arrow("json"), DataType::Utf8);
|
||||||
|
assert_eq!(mysql_type_to_arrow("bool"), DataType::Boolean);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,279 @@
|
|||||||
|
use arrow::array::{ArrayRef, Int32Array, StringArray};
|
||||||
|
use arrow::datatypes::{DataType, Field, Schema};
|
||||||
|
use arrow::record_batch::RecordBatch;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// Extract text from a PDF file.
|
||||||
|
///
|
||||||
|
/// Two-tier approach:
|
||||||
|
/// 1. **Text extraction (lopdf)** — fast, works on digital/native PDFs.
|
||||||
|
/// Returns one row per page: (source_file, page_number, text_content).
|
||||||
|
/// 2. **OCR fallback (Tesseract)** — if text extraction yields zero
|
||||||
|
/// pages (common with scanned documents, image-only PDFs, or PDFs
|
||||||
|
/// exported from scanners without an OCR layer), we extract embedded
|
||||||
|
/// images from the PDF and run Tesseract on each. Falls back
|
||||||
|
/// gracefully if Tesseract isn't installed (returns an error
|
||||||
|
/// directing the operator to install it).
|
||||||
|
///
|
||||||
|
/// Schema is identical for both paths so downstream consumers (chunker,
|
||||||
|
/// vectord, queryd) don't need to know how the text was produced.
|
||||||
|
pub fn parse_pdf(content: &[u8], source_filename: &str) -> Result<(Arc<Schema>, Vec<RecordBatch>), String> {
|
||||||
|
let doc = lopdf::Document::load_mem(content)
|
||||||
|
.map_err(|e| format!("PDF load error: {e}"))?;
|
||||||
|
|
||||||
|
let pages = doc.get_pages();
|
||||||
|
let mut page_numbers: Vec<i32> = Vec::new();
|
||||||
|
let mut page_texts: Vec<String> = Vec::new();
|
||||||
|
let mut sources: Vec<String> = Vec::new();
|
||||||
|
|
||||||
|
for (&page_num, _) in pages.iter() {
|
||||||
|
let text = doc.extract_text(&[page_num]).unwrap_or_default();
|
||||||
|
let text = text.trim().to_string();
|
||||||
|
|
||||||
|
if !text.is_empty() {
|
||||||
|
page_numbers.push(page_num as i32);
|
||||||
|
page_texts.push(text);
|
||||||
|
sources.push(source_filename.to_string());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Tier 2: OCR fallback for scanned pages.
|
||||||
|
if page_numbers.is_empty() {
|
||||||
|
tracing::info!(
|
||||||
|
"no extractable text in '{}' — attempting OCR on {} pages",
|
||||||
|
source_filename, pages.len(),
|
||||||
|
);
|
||||||
|
let ocr_results = ocr_pdf_images(&doc, source_filename)?;
|
||||||
|
if ocr_results.is_empty() {
|
||||||
|
return Err(
|
||||||
|
"PDF contains no extractable text and OCR produced no output. \
|
||||||
|
Ensure tesseract is installed (`apt install tesseract-ocr tesseract-ocr-eng`)."
|
||||||
|
.into(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
for (pg, text) in ocr_results {
|
||||||
|
page_numbers.push(pg);
|
||||||
|
page_texts.push(text);
|
||||||
|
sources.push(source_filename.to_string());
|
||||||
|
}
|
||||||
|
tracing::info!(
|
||||||
|
"OCR recovered {} pages from '{}'", page_numbers.len(), source_filename,
|
||||||
|
);
|
||||||
|
} else {
|
||||||
|
tracing::info!(
|
||||||
|
"extracted {} pages with text from PDF '{}'",
|
||||||
|
page_numbers.len(), source_filename,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let schema = Arc::new(Schema::new(vec![
|
||||||
|
Field::new("source_file", DataType::Utf8, false),
|
||||||
|
Field::new("page_number", DataType::Int32, false),
|
||||||
|
Field::new("text_content", DataType::Utf8, false),
|
||||||
|
]));
|
||||||
|
|
||||||
|
let arrays: Vec<ArrayRef> = vec![
|
||||||
|
Arc::new(StringArray::from(sources)),
|
||||||
|
Arc::new(Int32Array::from(page_numbers)),
|
||||||
|
Arc::new(StringArray::from(page_texts)),
|
||||||
|
];
|
||||||
|
|
||||||
|
let batch = RecordBatch::try_new(schema.clone(), arrays)
|
||||||
|
.map_err(|e| format!("RecordBatch error: {e}"))?;
|
||||||
|
|
||||||
|
Ok((schema, vec![batch]))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extract embedded images from the PDF, write each to a temp file,
|
||||||
|
/// run Tesseract, collect the OCR'd text. Returns (page_number, text)
|
||||||
|
/// pairs for pages that yielded non-empty output.
|
||||||
|
///
|
||||||
|
/// Image extraction strategy: walk every page's resources looking for
|
||||||
|
/// XObject streams with /Subtype /Image. Decode the raw stream bytes
|
||||||
|
/// (lopdf handles FlateDecode/DCTDecode) and write them as-is to a
|
||||||
|
/// temp file with the appropriate extension. Tesseract handles JPEG,
|
||||||
|
/// PNG, TIFF, and BMP natively.
|
||||||
|
fn ocr_pdf_images(
|
||||||
|
doc: &lopdf::Document,
|
||||||
|
source_filename: &str,
|
||||||
|
) -> Result<Vec<(i32, String)>, String> {
|
||||||
|
if !tesseract_available() {
|
||||||
|
return Err(
|
||||||
|
"tesseract not found on PATH — install with `apt install tesseract-ocr tesseract-ocr-eng`"
|
||||||
|
.into(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let pages = doc.get_pages();
|
||||||
|
let mut results: Vec<(i32, String)> = Vec::new();
|
||||||
|
|
||||||
|
for (&page_num, &page_id) in &pages {
|
||||||
|
let images = extract_page_images(doc, page_id);
|
||||||
|
if images.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// OCR each image on this page and concatenate the text.
|
||||||
|
let mut page_text = String::new();
|
||||||
|
for (idx, img_bytes) in images.iter().enumerate() {
|
||||||
|
match ocr_image_bytes(img_bytes, page_num, idx) {
|
||||||
|
Ok(text) if !text.trim().is_empty() => {
|
||||||
|
if !page_text.is_empty() { page_text.push('\n'); }
|
||||||
|
page_text.push_str(text.trim());
|
||||||
|
}
|
||||||
|
Ok(_) => {} // empty OCR — skip
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!(
|
||||||
|
"OCR failed on '{}' page {} image {}: {e}",
|
||||||
|
source_filename, page_num, idx,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !page_text.is_empty() {
|
||||||
|
results.push((page_num as i32, page_text));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(results)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extract raw image bytes from a PDF page's XObject resources.
|
||||||
|
/// Handles the common case: page /Resources → /XObject dict → streams
|
||||||
|
/// with /Subtype /Image.
|
||||||
|
fn extract_page_images(
|
||||||
|
doc: &lopdf::Document,
|
||||||
|
page_id: lopdf::ObjectId,
|
||||||
|
) -> Vec<Vec<u8>> {
|
||||||
|
let mut images = Vec::new();
|
||||||
|
|
||||||
|
let page_obj = match doc.get_object(page_id) {
|
||||||
|
Ok(o) => o,
|
||||||
|
Err(_) => return images,
|
||||||
|
};
|
||||||
|
let page_dict = match page_obj.as_dict() {
|
||||||
|
Ok(d) => d,
|
||||||
|
Err(_) => return images,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Resolve /Resources → /XObject
|
||||||
|
let resources = page_dict
|
||||||
|
.get(b"Resources")
|
||||||
|
.ok()
|
||||||
|
.and_then(|r| doc.dereference(r).ok())
|
||||||
|
.map(|(_, o)| o.clone());
|
||||||
|
let xobjects = resources
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|r| r.as_dict().ok())
|
||||||
|
.and_then(|d| d.get(b"XObject").ok())
|
||||||
|
.and_then(|x| doc.dereference(x).ok())
|
||||||
|
.map(|(_, o)| o.clone());
|
||||||
|
|
||||||
|
let xobj_dict = match xobjects.as_ref().and_then(|x| x.as_dict().ok()) {
|
||||||
|
Some(d) => d,
|
||||||
|
None => return images,
|
||||||
|
};
|
||||||
|
|
||||||
|
for (_name, obj_ref) in xobj_dict.iter() {
|
||||||
|
let obj = match doc.dereference(obj_ref) {
|
||||||
|
Ok((_, o)) => o.clone(),
|
||||||
|
Err(_) => continue,
|
||||||
|
};
|
||||||
|
let stream = match obj.as_stream() {
|
||||||
|
Ok(s) => s,
|
||||||
|
Err(_) => continue,
|
||||||
|
};
|
||||||
|
// Check /Subtype == /Image
|
||||||
|
let is_image = stream.dict.get(b"Subtype")
|
||||||
|
.ok()
|
||||||
|
.and_then(|s| s.as_name().ok())
|
||||||
|
.map(|n| n == b"Image")
|
||||||
|
.unwrap_or(false);
|
||||||
|
if !is_image {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
// Get decoded content — lopdf handles FlateDecode/DCTDecode.
|
||||||
|
let content = match stream.decompressed_content() {
|
||||||
|
Ok(c) => c,
|
||||||
|
Err(_) => {
|
||||||
|
// Try raw content as fallback (some DCT streams are stored raw).
|
||||||
|
stream.content.clone()
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if !content.is_empty() {
|
||||||
|
images.push(content);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
images
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run Tesseract on raw image bytes. Writes to a temp file, shells out,
|
||||||
|
/// reads the output. Returns the OCR'd text.
|
||||||
|
fn ocr_image_bytes(
|
||||||
|
img_bytes: &[u8],
|
||||||
|
page_num: u32,
|
||||||
|
img_idx: usize,
|
||||||
|
) -> Result<String, String> {
|
||||||
|
use std::io::Write;
|
||||||
|
use std::process::Command;
|
||||||
|
|
||||||
|
let tmp_dir = std::env::temp_dir();
|
||||||
|
// Detect image format from magic bytes to pick the right extension.
|
||||||
|
let ext = if img_bytes.starts_with(&[0xFF, 0xD8]) {
|
||||||
|
"jpg"
|
||||||
|
} else if img_bytes.starts_with(&[0x89, 0x50, 0x4E, 0x47]) {
|
||||||
|
"png"
|
||||||
|
} else if img_bytes.starts_with(b"II") || img_bytes.starts_with(b"MM") {
|
||||||
|
"tiff"
|
||||||
|
} else {
|
||||||
|
// Tesseract can often handle raw pixel data with the right options,
|
||||||
|
// but writing as PNG is safer. For truly unknown formats, try PNG
|
||||||
|
// wrapping; if Tesseract can't read it, it'll error cleanly.
|
||||||
|
"png"
|
||||||
|
};
|
||||||
|
|
||||||
|
let img_path = tmp_dir.join(format!("lh_ocr_p{page_num}_i{img_idx}.{ext}"));
|
||||||
|
let out_base = tmp_dir.join(format!("lh_ocr_p{page_num}_i{img_idx}_out"));
|
||||||
|
let out_path = tmp_dir.join(format!("lh_ocr_p{page_num}_i{img_idx}_out.txt"));
|
||||||
|
|
||||||
|
// Write image bytes to temp file.
|
||||||
|
let mut f = std::fs::File::create(&img_path)
|
||||||
|
.map_err(|e| format!("create temp image: {e}"))?;
|
||||||
|
f.write_all(img_bytes)
|
||||||
|
.map_err(|e| format!("write temp image: {e}"))?;
|
||||||
|
drop(f);
|
||||||
|
|
||||||
|
// Run Tesseract.
|
||||||
|
let status = Command::new("tesseract")
|
||||||
|
.arg(&img_path)
|
||||||
|
.arg(&out_base)
|
||||||
|
.arg("--oem").arg("3") // LSTM + legacy combined
|
||||||
|
.arg("--psm").arg("6") // assume uniform block of text
|
||||||
|
.arg("-l").arg("eng")
|
||||||
|
.stdout(std::process::Stdio::null())
|
||||||
|
.stderr(std::process::Stdio::null())
|
||||||
|
.status()
|
||||||
|
.map_err(|e| format!("tesseract exec: {e}"))?;
|
||||||
|
|
||||||
|
// Cleanup input regardless of outcome.
|
||||||
|
let _ = std::fs::remove_file(&img_path);
|
||||||
|
|
||||||
|
if !status.success() {
|
||||||
|
let _ = std::fs::remove_file(&out_path);
|
||||||
|
return Err(format!("tesseract exited with status {status}"));
|
||||||
|
}
|
||||||
|
|
||||||
|
let text = std::fs::read_to_string(&out_path).unwrap_or_default();
|
||||||
|
let _ = std::fs::remove_file(&out_path);
|
||||||
|
Ok(text)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Quick check: is `tesseract` on PATH?
|
||||||
|
fn tesseract_available() -> bool {
|
||||||
|
std::process::Command::new("tesseract")
|
||||||
|
.arg("--version")
|
||||||
|
.stdout(std::process::Stdio::null())
|
||||||
|
.stderr(std::process::Stdio::null())
|
||||||
|
.status()
|
||||||
|
.map(|s| s.success())
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
@@ -0,0 +1,308 @@
|
|||||||
|
/// Streaming PostgreSQL ingest.
|
||||||
|
///
|
||||||
|
/// The original `db_ingest::import_postgres_table` loads every row into
|
||||||
|
/// memory before emitting Parquet — fine for small tables, blows up on
|
||||||
|
/// 1M+ rows. This module paginates via `ORDER BY <pk> LIMIT N OFFSET M`
|
||||||
|
/// and streams batches into an `ArrowWriter`, closing once exhausted.
|
||||||
|
///
|
||||||
|
/// Pagination is OFFSET-based (not keyset) — simpler, works on any table
|
||||||
|
/// without needing to know the PK type. OFFSET N scans N rows per fetch,
|
||||||
|
/// which becomes quadratic for very large tables. Upgrade to keyset when
|
||||||
|
/// a user actually ingests something that hurts.
|
||||||
|
|
||||||
|
use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray};
|
||||||
|
use arrow::datatypes::{DataType, Field, Schema};
|
||||||
|
use parquet::arrow::ArrowWriter;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio_postgres::{Client, NoTls, types::Type as PgType};
|
||||||
|
|
||||||
|
use crate::db_ingest::DbConfig;
|
||||||
|
|
||||||
|
/// Request shape for streaming ingest — takes a DSN string and optional
|
||||||
|
/// tuning knobs.
|
||||||
|
#[derive(Debug, Clone, serde::Deserialize)]
|
||||||
|
pub struct PgStreamRequest {
|
||||||
|
/// postgresql://user:pass@host:port/db — alternative to DbConfig.
|
||||||
|
pub dsn: String,
|
||||||
|
pub table: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub dataset_name: Option<String>,
|
||||||
|
/// Rows per fetch. Default 10_000.
|
||||||
|
#[serde(default)]
|
||||||
|
pub batch_size: Option<usize>,
|
||||||
|
/// Column to ORDER BY for stable pagination. If omitted, the first
|
||||||
|
/// column returned by the schema probe is used.
|
||||||
|
#[serde(default)]
|
||||||
|
pub order_by: Option<String>,
|
||||||
|
/// Hard cap on total rows (for sampling / previews).
|
||||||
|
#[serde(default)]
|
||||||
|
pub limit: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct PgStreamResult {
|
||||||
|
pub table: String,
|
||||||
|
pub rows: usize,
|
||||||
|
pub batches: usize,
|
||||||
|
pub columns: usize,
|
||||||
|
pub schema: Vec<String>,
|
||||||
|
pub parquet_bytes: u64,
|
||||||
|
pub duration_secs: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parse a postgresql:// DSN into a DbConfig.
|
||||||
|
/// Supports: postgresql://[user[:password]@]host[:port][/db]
|
||||||
|
/// Does NOT support: query parameters (sslmode etc) — add if needed.
|
||||||
|
pub fn parse_dsn(dsn: &str) -> Result<DbConfig, String> {
|
||||||
|
let rest = dsn
|
||||||
|
.strip_prefix("postgresql://")
|
||||||
|
.or_else(|| dsn.strip_prefix("postgres://"))
|
||||||
|
.ok_or_else(|| "DSN must start with postgresql:// or postgres://".to_string())?;
|
||||||
|
|
||||||
|
// Split off path (database) first.
|
||||||
|
let (auth_host, database) = match rest.split_once('/') {
|
||||||
|
Some((ah, db)) => (ah, db.split('?').next().unwrap_or(db).to_string()),
|
||||||
|
None => (rest, "postgres".to_string()),
|
||||||
|
};
|
||||||
|
|
||||||
|
// Split user[:password] from host[:port].
|
||||||
|
let (userpass, hostport) = match auth_host.rsplit_once('@') {
|
||||||
|
Some((up, hp)) => (Some(up), hp),
|
||||||
|
None => (None, auth_host),
|
||||||
|
};
|
||||||
|
|
||||||
|
let (user, password) = match userpass {
|
||||||
|
Some(up) => match up.split_once(':') {
|
||||||
|
Some((u, p)) => (u.to_string(), p.to_string()),
|
||||||
|
None => (up.to_string(), String::new()),
|
||||||
|
},
|
||||||
|
None => ("postgres".to_string(), String::new()),
|
||||||
|
};
|
||||||
|
|
||||||
|
let (host, port) = match hostport.rsplit_once(':') {
|
||||||
|
Some((h, p)) => (
|
||||||
|
h.to_string(),
|
||||||
|
p.parse::<u16>().map_err(|_| format!("invalid port in DSN: {p}"))?,
|
||||||
|
),
|
||||||
|
None => (hostport.to_string(), 5432),
|
||||||
|
};
|
||||||
|
|
||||||
|
if host.is_empty() {
|
||||||
|
return Err("DSN has no host".into());
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(DbConfig { host, port, database, user, password })
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Stream a Postgres table as Parquet bytes.
|
||||||
|
///
|
||||||
|
/// Returns the full Parquet payload + summary stats. The payload is
|
||||||
|
/// written to an in-memory buffer via `ArrowWriter` so each batch gets
|
||||||
|
/// columnar compression as it arrives — memory footprint stays roughly
|
||||||
|
/// at one batch plus Parquet footer state.
|
||||||
|
pub async fn stream_table_to_parquet(
|
||||||
|
req: &PgStreamRequest,
|
||||||
|
) -> Result<(bytes::Bytes, PgStreamResult), String> {
|
||||||
|
let t0 = std::time::Instant::now();
|
||||||
|
let config = parse_dsn(&req.dsn)?;
|
||||||
|
let batch_size = req.batch_size.unwrap_or(10_000).max(1);
|
||||||
|
|
||||||
|
let (client, connection) = tokio_postgres::connect(&config.connection_string(), NoTls)
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("postgres connect: {e}"))?;
|
||||||
|
tokio::spawn(async move {
|
||||||
|
if let Err(e) = connection.await {
|
||||||
|
tracing::error!("pg connection: {e}");
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// Probe schema.
|
||||||
|
let columns = probe_columns(&client, &req.table).await?;
|
||||||
|
if columns.is_empty() {
|
||||||
|
return Err(format!("table '{}' not found or has no columns", req.table));
|
||||||
|
}
|
||||||
|
|
||||||
|
let arrow_fields: Vec<Field> = columns
|
||||||
|
.iter()
|
||||||
|
.map(|(name, pg)| Field::new(name, pg_type_to_arrow(pg), true))
|
||||||
|
.collect();
|
||||||
|
let schema = Arc::new(Schema::new(arrow_fields));
|
||||||
|
let schema_report: Vec<String> = columns
|
||||||
|
.iter()
|
||||||
|
.map(|(n, t)| format!("{}:{}", n, t))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// Pagination key: user-specified or first column.
|
||||||
|
let order_col = req.order_by.clone().unwrap_or_else(|| columns[0].0.clone());
|
||||||
|
|
||||||
|
// Stream batches into ArrowWriter.
|
||||||
|
let mut buf: Vec<u8> = Vec::with_capacity(1024 * 1024);
|
||||||
|
let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)
|
||||||
|
.map_err(|e| format!("ArrowWriter init: {e}"))?;
|
||||||
|
|
||||||
|
let mut total_rows: usize = 0;
|
||||||
|
let mut batch_count: usize = 0;
|
||||||
|
let row_cap = req.limit.unwrap_or(usize::MAX);
|
||||||
|
|
||||||
|
loop {
|
||||||
|
let remaining = row_cap.saturating_sub(total_rows);
|
||||||
|
if remaining == 0 { break; }
|
||||||
|
let fetch = remaining.min(batch_size);
|
||||||
|
|
||||||
|
let sql = format!(
|
||||||
|
"SELECT * FROM \"{}\" ORDER BY \"{}\" LIMIT {} OFFSET {}",
|
||||||
|
req.table, order_col, fetch, total_rows,
|
||||||
|
);
|
||||||
|
|
||||||
|
let rows = client
|
||||||
|
.query(&sql, &[])
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("fetch batch at offset {total_rows}: {e}"))?;
|
||||||
|
|
||||||
|
if rows.is_empty() { break; }
|
||||||
|
|
||||||
|
let n = rows.len();
|
||||||
|
let arrays: Vec<ArrayRef> = columns
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(idx, (_, pg))| rows_to_column(&rows, idx, pg))
|
||||||
|
.collect::<Result<_, _>>()?;
|
||||||
|
let batch = RecordBatch::try_new(schema.clone(), arrays)
|
||||||
|
.map_err(|e| format!("RecordBatch: {e}"))?;
|
||||||
|
writer.write(&batch).map_err(|e| format!("ArrowWriter::write: {e}"))?;
|
||||||
|
|
||||||
|
total_rows += n;
|
||||||
|
batch_count += 1;
|
||||||
|
tracing::info!(
|
||||||
|
"pg stream '{}': fetched batch {} ({} rows, total {})",
|
||||||
|
req.table, batch_count, n, total_rows,
|
||||||
|
);
|
||||||
|
|
||||||
|
if n < fetch { break; } // short read = end of table
|
||||||
|
}
|
||||||
|
|
||||||
|
writer.close().map_err(|e| format!("ArrowWriter::close: {e}"))?;
|
||||||
|
let parquet_bytes = buf.len() as u64;
|
||||||
|
let duration = t0.elapsed().as_secs_f32();
|
||||||
|
|
||||||
|
let result = PgStreamResult {
|
||||||
|
table: req.table.clone(),
|
||||||
|
rows: total_rows,
|
||||||
|
batches: batch_count,
|
||||||
|
columns: columns.len(),
|
||||||
|
schema: schema_report,
|
||||||
|
parquet_bytes,
|
||||||
|
duration_secs: duration,
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok((bytes::Bytes::from(buf), result))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn probe_columns(client: &Client, table: &str) -> Result<Vec<(String, PgType)>, String> {
|
||||||
|
let stmt = client
|
||||||
|
.prepare(&format!("SELECT * FROM \"{}\" LIMIT 0", table))
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("prepare schema probe: {e}"))?;
|
||||||
|
Ok(stmt
|
||||||
|
.columns()
|
||||||
|
.iter()
|
||||||
|
.map(|c| (c.name().to_string(), c.type_().clone()))
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn pg_type_to_arrow(pg: &PgType) -> DataType {
|
||||||
|
match *pg {
|
||||||
|
PgType::BOOL => DataType::Boolean,
|
||||||
|
PgType::INT2 | PgType::INT4 => DataType::Int32,
|
||||||
|
PgType::INT8 | PgType::OID => DataType::Int64,
|
||||||
|
PgType::FLOAT4 | PgType::FLOAT8 | PgType::NUMERIC => DataType::Float64,
|
||||||
|
_ => DataType::Utf8,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn rows_to_column(
|
||||||
|
rows: &[tokio_postgres::Row],
|
||||||
|
idx: usize,
|
||||||
|
pg: &PgType,
|
||||||
|
) -> Result<ArrayRef, String> {
|
||||||
|
match *pg {
|
||||||
|
PgType::BOOL => {
|
||||||
|
let v: Vec<Option<bool>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
|
||||||
|
Ok(Arc::new(BooleanArray::from(v)))
|
||||||
|
}
|
||||||
|
PgType::INT2 => {
|
||||||
|
let v: Vec<Option<i32>> = rows.iter()
|
||||||
|
.map(|r| r.try_get::<_, i16>(idx).ok().map(|x| x as i32)).collect();
|
||||||
|
Ok(Arc::new(Int32Array::from(v)))
|
||||||
|
}
|
||||||
|
PgType::INT4 => {
|
||||||
|
let v: Vec<Option<i32>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
|
||||||
|
Ok(Arc::new(Int32Array::from(v)))
|
||||||
|
}
|
||||||
|
PgType::INT8 | PgType::OID => {
|
||||||
|
let v: Vec<Option<i64>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
|
||||||
|
Ok(Arc::new(Int64Array::from(v)))
|
||||||
|
}
|
||||||
|
PgType::FLOAT4 => {
|
||||||
|
let v: Vec<Option<f64>> = rows.iter()
|
||||||
|
.map(|r| r.try_get::<_, f32>(idx).ok().map(|x| x as f64)).collect();
|
||||||
|
Ok(Arc::new(Float64Array::from(v)))
|
||||||
|
}
|
||||||
|
PgType::FLOAT8 => {
|
||||||
|
let v: Vec<Option<f64>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
|
||||||
|
Ok(Arc::new(Float64Array::from(v)))
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
// Safe-default per ADR-010: everything else becomes string.
|
||||||
|
let v: Vec<Option<String>> = rows.iter().map(|r| {
|
||||||
|
r.try_get::<_, String>(idx).ok()
|
||||||
|
.or_else(|| r.try_get::<_, serde_json::Value>(idx).ok().map(|j| j.to_string()))
|
||||||
|
.or_else(|| {
|
||||||
|
// Timestamps + uuid, rendered via Display.
|
||||||
|
r.try_get::<_, chrono::DateTime<chrono::Utc>>(idx).ok()
|
||||||
|
.map(|t| t.to_rfc3339())
|
||||||
|
})
|
||||||
|
.or_else(|| r.try_get::<_, uuid::Uuid>(idx).ok().map(|u| u.to_string()))
|
||||||
|
.or(Some(String::new()))
|
||||||
|
}).collect();
|
||||||
|
Ok(Arc::new(StringArray::from(v)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_dsn_full() {
|
||||||
|
let c = parse_dsn("postgresql://daisy:secret@db.example.com:6543/my_db").unwrap();
|
||||||
|
assert_eq!(c.host, "db.example.com");
|
||||||
|
assert_eq!(c.port, 6543);
|
||||||
|
assert_eq!(c.user, "daisy");
|
||||||
|
assert_eq!(c.password, "secret");
|
||||||
|
assert_eq!(c.database, "my_db");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_dsn_minimal() {
|
||||||
|
let c = parse_dsn("postgresql://localhost/knowledge_base").unwrap();
|
||||||
|
assert_eq!(c.host, "localhost");
|
||||||
|
assert_eq!(c.port, 5432);
|
||||||
|
assert_eq!(c.user, "postgres");
|
||||||
|
assert_eq!(c.password, "");
|
||||||
|
assert_eq!(c.database, "knowledge_base");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_dsn_no_password() {
|
||||||
|
let c = parse_dsn("postgres://postgres@127.0.0.1:5432/mydb").unwrap();
|
||||||
|
assert_eq!(c.user, "postgres");
|
||||||
|
assert_eq!(c.password, "");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_dsn_rejects_non_pg() {
|
||||||
|
assert!(parse_dsn("http://host/db").is_err());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,219 @@
|
|||||||
|
use arrow::datatypes::SchemaRef;
|
||||||
|
use arrow::record_batch::RecordBatch;
|
||||||
|
use bytes::Bytes;
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use catalogd::registry::Registry;
|
||||||
|
use shared::arrow_helpers::{fingerprint_schema, record_batch_to_parquet};
|
||||||
|
use shared::types::{ObjectRef, SchemaFingerprint};
|
||||||
|
use storaged::ops;
|
||||||
|
|
||||||
|
use crate::detect::{FileType, content_hash, detect_file_type};
|
||||||
|
use crate::csv_ingest;
|
||||||
|
use crate::json_ingest;
|
||||||
|
use crate::pdf_ingest;
|
||||||
|
|
||||||
|
/// Result of an ingest operation.
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct IngestResult {
|
||||||
|
pub dataset_name: String,
|
||||||
|
pub file_type: String,
|
||||||
|
pub rows: usize,
|
||||||
|
pub columns: usize,
|
||||||
|
pub storage_key: String,
|
||||||
|
pub content_hash: String,
|
||||||
|
pub schema_fingerprint: String,
|
||||||
|
pub deduplicated: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Full ingest pipeline: detect → parse → dedup → store → register.
|
||||||
|
///
|
||||||
|
/// Bucket name (federation layer 2) is recorded on the ObjectRef so the
|
||||||
|
/// catalog knows which bucket holds the data — defaults to "primary" via
|
||||||
|
/// the ingest service's resolve_bucket helper.
|
||||||
|
pub async fn ingest_file(
|
||||||
|
filename: &str,
|
||||||
|
content: &[u8],
|
||||||
|
dataset_name: Option<&str>,
|
||||||
|
store: &Arc<dyn ObjectStore>,
|
||||||
|
registry: &Registry,
|
||||||
|
) -> Result<IngestResult, String> {
|
||||||
|
ingest_file_to_bucket(filename, content, dataset_name, "primary", store, registry).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Same as `ingest_file` but with explicit target bucket name on the
|
||||||
|
/// ObjectRef. Header-aware ingest endpoints call this directly.
|
||||||
|
pub async fn ingest_file_to_bucket(
|
||||||
|
filename: &str,
|
||||||
|
content: &[u8],
|
||||||
|
dataset_name: Option<&str>,
|
||||||
|
bucket: &str,
|
||||||
|
store: &Arc<dyn ObjectStore>,
|
||||||
|
registry: &Registry,
|
||||||
|
) -> Result<IngestResult, String> {
|
||||||
|
// 1. Detect file type
|
||||||
|
let file_type = detect_file_type(filename, content);
|
||||||
|
tracing::info!("ingesting '{}' as {:?} ({} bytes)", filename, file_type, content.len());
|
||||||
|
|
||||||
|
// 2. Parse into Arrow
|
||||||
|
let (schema, batches) = match file_type {
|
||||||
|
FileType::Csv => csv_ingest::parse_csv(content)?,
|
||||||
|
FileType::Json | FileType::NdJson => json_ingest::parse_json(content)?,
|
||||||
|
FileType::Pdf => pdf_ingest::parse_pdf(content, filename)?,
|
||||||
|
FileType::Text => parse_text_file(content, filename)?,
|
||||||
|
FileType::Unknown => return Err(format!("unknown file type for '{filename}'")),
|
||||||
|
};
|
||||||
|
|
||||||
|
let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
|
||||||
|
let total_cols = schema.fields().len();
|
||||||
|
|
||||||
|
if total_rows == 0 {
|
||||||
|
return Err("file contains no data".into());
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. Content hash for dedup
|
||||||
|
let hash = content_hash(content);
|
||||||
|
|
||||||
|
// 4. Check if already ingested (same content hash)
|
||||||
|
let existing = registry.list().await;
|
||||||
|
if existing.iter().any(|d| d.schema_fingerprint.0 == hash) {
|
||||||
|
tracing::info!("file already ingested (hash: {}), skipping", &hash[..12]);
|
||||||
|
return Ok(IngestResult {
|
||||||
|
dataset_name: dataset_name.unwrap_or(filename).to_string(),
|
||||||
|
file_type: format!("{:?}", file_type),
|
||||||
|
rows: total_rows,
|
||||||
|
columns: total_cols,
|
||||||
|
storage_key: String::new(),
|
||||||
|
content_hash: hash,
|
||||||
|
schema_fingerprint: String::new(),
|
||||||
|
deduplicated: true,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// 5. Convert to Parquet
|
||||||
|
let mut all_parquet = Vec::new();
|
||||||
|
for batch in &batches {
|
||||||
|
let pq = record_batch_to_parquet(batch)?;
|
||||||
|
all_parquet.extend_from_slice(&pq);
|
||||||
|
}
|
||||||
|
let parquet_bytes = Bytes::from(all_parquet);
|
||||||
|
let parquet_size = parquet_bytes.len() as u64;
|
||||||
|
|
||||||
|
// 6. Store in object storage
|
||||||
|
let name = dataset_name.unwrap_or_else(|| {
|
||||||
|
filename.rsplit('/').next().unwrap_or(filename)
|
||||||
|
.rsplit('.').last().unwrap_or(filename)
|
||||||
|
});
|
||||||
|
let safe_name = sanitize_dataset_name(name);
|
||||||
|
let storage_key = format!("datasets/{}.parquet", safe_name);
|
||||||
|
|
||||||
|
ops::put(store, &storage_key, parquet_bytes).await?;
|
||||||
|
tracing::info!("stored {} as {} ({} bytes)", filename, storage_key, parquet_size);
|
||||||
|
|
||||||
|
// 7. Register in catalog with rich metadata
|
||||||
|
let schema_fp = fingerprint_schema(&schema);
|
||||||
|
let now = chrono::Utc::now();
|
||||||
|
let obj_ref = ObjectRef {
|
||||||
|
bucket: bucket.to_string(),
|
||||||
|
key: storage_key.clone(),
|
||||||
|
size_bytes: parquet_size,
|
||||||
|
created_at: now,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Register base dataset
|
||||||
|
registry.register(
|
||||||
|
safe_name.clone(),
|
||||||
|
SchemaFingerprint(hash.clone()),
|
||||||
|
vec![obj_ref],
|
||||||
|
).await?;
|
||||||
|
|
||||||
|
// Auto-populate metadata: PII detection, column info, lineage, row count
|
||||||
|
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
|
||||||
|
let dataset_sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
|
||||||
|
|
||||||
|
let column_meta: Vec<shared::types::ColumnMeta> = schema.fields().iter().map(|f| {
|
||||||
|
let sens = shared::pii::detect_sensitivity(f.name());
|
||||||
|
shared::types::ColumnMeta {
|
||||||
|
name: f.name().clone(),
|
||||||
|
data_type: f.data_type().to_string(),
|
||||||
|
sensitivity: sens.clone(),
|
||||||
|
description: String::new(),
|
||||||
|
is_pii: matches!(sens, Some(shared::types::Sensitivity::Pii)),
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
|
||||||
|
let lineage = shared::types::Lineage {
|
||||||
|
source_system: format!("{:?}", file_type).to_lowercase(),
|
||||||
|
source_file: filename.to_string(),
|
||||||
|
ingest_job: format!("ingest-{}", now.timestamp_millis()),
|
||||||
|
ingest_timestamp: now,
|
||||||
|
parent_datasets: vec![],
|
||||||
|
};
|
||||||
|
|
||||||
|
let pii_count = column_meta.iter().filter(|c| c.is_pii).count();
|
||||||
|
if pii_count > 0 {
|
||||||
|
tracing::info!("auto-detected {} PII columns in '{}'", pii_count, safe_name);
|
||||||
|
}
|
||||||
|
|
||||||
|
let _ = registry.update_metadata(&safe_name, catalogd::registry::MetadataUpdate {
|
||||||
|
sensitivity: dataset_sensitivity,
|
||||||
|
columns: Some(column_meta),
|
||||||
|
lineage: Some(lineage),
|
||||||
|
row_count: Some(total_rows as u64),
|
||||||
|
..Default::default()
|
||||||
|
}).await;
|
||||||
|
|
||||||
|
// Phase C: if this dataset already had embeddings, they're now stale.
|
||||||
|
// mark_embeddings_stale is a no-op for never-embedded datasets.
|
||||||
|
let _ = registry.mark_embeddings_stale(&safe_name).await;
|
||||||
|
|
||||||
|
Ok(IngestResult {
|
||||||
|
dataset_name: safe_name,
|
||||||
|
file_type: format!("{:?}", file_type),
|
||||||
|
rows: total_rows,
|
||||||
|
columns: total_cols,
|
||||||
|
storage_key,
|
||||||
|
content_hash: hash,
|
||||||
|
schema_fingerprint: schema_fp.0,
|
||||||
|
deduplicated: false,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parse plain text / SMS logs into a simple table.
|
||||||
|
fn parse_text_file(content: &[u8], filename: &str) -> Result<(SchemaRef, Vec<RecordBatch>), String> {
|
||||||
|
let text = String::from_utf8_lossy(content);
|
||||||
|
let lines: Vec<&str> = text.lines().filter(|l| !l.trim().is_empty()).collect();
|
||||||
|
|
||||||
|
if lines.is_empty() {
|
||||||
|
return Err("text file is empty".into());
|
||||||
|
}
|
||||||
|
|
||||||
|
let schema = Arc::new(arrow::datatypes::Schema::new(vec![
|
||||||
|
arrow::datatypes::Field::new("source_file", arrow::datatypes::DataType::Utf8, false),
|
||||||
|
arrow::datatypes::Field::new("line_number", arrow::datatypes::DataType::Int32, false),
|
||||||
|
arrow::datatypes::Field::new("text", arrow::datatypes::DataType::Utf8, false),
|
||||||
|
]));
|
||||||
|
|
||||||
|
let sources: Vec<&str> = vec![filename; lines.len()];
|
||||||
|
let line_nums: Vec<i32> = (1..=lines.len() as i32).collect();
|
||||||
|
|
||||||
|
let arrays: Vec<arrow::array::ArrayRef> = vec![
|
||||||
|
Arc::new(arrow::array::StringArray::from(sources)),
|
||||||
|
Arc::new(arrow::array::Int32Array::from(line_nums)),
|
||||||
|
Arc::new(arrow::array::StringArray::from(lines)),
|
||||||
|
];
|
||||||
|
|
||||||
|
let batch = RecordBatch::try_new(schema.clone(), arrays)
|
||||||
|
.map_err(|e| format!("RecordBatch error: {e}"))?;
|
||||||
|
|
||||||
|
Ok((schema, vec![batch]))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn sanitize_dataset_name(name: &str) -> String {
|
||||||
|
let clean: String = name.chars()
|
||||||
|
.map(|c| if c.is_alphanumeric() || c == '_' { c.to_ascii_lowercase() } else { '_' })
|
||||||
|
.collect();
|
||||||
|
let trimmed = clean.trim_matches('_').to_string();
|
||||||
|
if trimmed.is_empty() { "unnamed_dataset".to_string() } else { trimmed }
|
||||||
|
}
|
||||||
@@ -0,0 +1,626 @@
|
|||||||
|
//! Scheduled ingest — the production-substrate piece that turns the
|
||||||
|
//! lakehouse from "manual API toolkit" into a system that runs itself.
|
||||||
|
//!
|
||||||
|
//! A `ScheduleDef` declares "what to ingest" + "when to fire." A
|
||||||
|
//! background `Scheduler` task wakes periodically, asks each enabled
|
||||||
|
//! schedule whether it's due, and runs the matching ingest path. Outcomes
|
||||||
|
//! land back on the schedule so operators can see what happened without
|
||||||
|
//! `journalctl` archaeology.
|
||||||
|
//!
|
||||||
|
//! Storage shape: one JSON file per schedule at
|
||||||
|
//! `primary://_schedules/{id}.json`. Serializable via serde, rebuilt on
|
||||||
|
//! startup. Same write-once shape as the rest of the catalog — no DB.
|
||||||
|
//!
|
||||||
|
//! Concurrency model: single-flight per schedule_id. The scheduler holds
|
||||||
|
//! a `HashSet<String>` of currently-running ids; a tick that finds a
|
||||||
|
//! schedule already in flight skips it (no queueing — a slow ingest
|
||||||
|
//! shouldn't pile up runs). Different schedules run concurrently.
|
||||||
|
//!
|
||||||
|
//! What's deliberately not in scope here:
|
||||||
|
//! - Cron expressions (the trigger enum has the variant but parsing is
|
||||||
|
//! stubbed). Intervals cover 90% of operational scheduling; cron is
|
||||||
|
//! easy to bolt on later.
|
||||||
|
//! - Backoff / retry policies. A failed run records the failure and
|
||||||
|
//! schedules `next_run_at` as if it succeeded — no exponential backoff.
|
||||||
|
//! - Distributed coordination. Single-machine scheduler; no leader
|
||||||
|
//! election. If you run two gateway instances they'll both fire.
|
||||||
|
|
||||||
|
use chrono::{DateTime, Duration, Utc};
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::{HashMap, HashSet};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
use storaged::ops;
|
||||||
|
|
||||||
|
const PREFIX: &str = "_schedules";
|
||||||
|
|
||||||
|
// =================== Public types ===================
|
||||||
|
|
||||||
|
/// What kind of ingest the schedule fires.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
#[serde(tag = "type", rename_all = "snake_case")]
|
||||||
|
pub enum ScheduleKind {
|
||||||
|
/// Pull a MySQL/MariaDB table via the my_stream path.
|
||||||
|
Mysql {
|
||||||
|
/// `mysql://user:pass@host:port/db` — DSN with full credentials.
|
||||||
|
/// Stored encrypted-at-rest is a future concern; for now it lives
|
||||||
|
/// in plain JSON in the catalog. Don't put production root creds
|
||||||
|
/// here without disk-level encryption.
|
||||||
|
dsn: String,
|
||||||
|
table: String,
|
||||||
|
#[serde(default)]
|
||||||
|
dataset_name: Option<String>,
|
||||||
|
#[serde(default)]
|
||||||
|
batch_size: Option<usize>,
|
||||||
|
#[serde(default)]
|
||||||
|
order_by: Option<String>,
|
||||||
|
},
|
||||||
|
/// Pull a Postgres table via the pg_stream path.
|
||||||
|
Postgres {
|
||||||
|
dsn: String,
|
||||||
|
table: String,
|
||||||
|
#[serde(default)]
|
||||||
|
dataset_name: Option<String>,
|
||||||
|
#[serde(default)]
|
||||||
|
batch_size: Option<usize>,
|
||||||
|
#[serde(default)]
|
||||||
|
order_by: Option<String>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ScheduleKind {
|
||||||
|
pub fn label(&self) -> String {
|
||||||
|
match self {
|
||||||
|
ScheduleKind::Mysql { table, .. } => format!("mysql:{table}"),
|
||||||
|
ScheduleKind::Postgres { table, .. } => format!("postgres:{table}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// When the schedule fires.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
#[serde(tag = "type", rename_all = "snake_case")]
|
||||||
|
pub enum ScheduleTrigger {
|
||||||
|
/// Run every N seconds. Time anchor = compute_next_run_at decision.
|
||||||
|
Interval { secs: u64 },
|
||||||
|
/// Cron expression — parsing not implemented yet. Defining the
|
||||||
|
/// variant now so the JSON shape is forward-compatible.
|
||||||
|
Cron { expr: String },
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Outcome of a single run, kept on the schedule for observability.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct RunOutcome {
|
||||||
|
pub at: DateTime<Utc>,
|
||||||
|
pub success: bool,
|
||||||
|
pub message: String,
|
||||||
|
pub rows: Option<usize>,
|
||||||
|
pub duration_secs: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One scheduled ingest definition.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ScheduleDef {
|
||||||
|
pub id: String,
|
||||||
|
pub kind: ScheduleKind,
|
||||||
|
pub trigger: ScheduleTrigger,
|
||||||
|
pub enabled: bool,
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub created_by: String,
|
||||||
|
/// When the next run is due. Recomputed after each completion via
|
||||||
|
/// `compute_next_run_at`. On creation: set to "now" so the schedule
|
||||||
|
/// fires on the next scheduler tick.
|
||||||
|
pub next_run_at: DateTime<Utc>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub last_run_at: Option<DateTime<Utc>>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub last_outcome: Option<RunOutcome>,
|
||||||
|
/// Total successful runs and total failures since creation. Useful
|
||||||
|
/// for at-a-glance "is this schedule healthy?" checks.
|
||||||
|
#[serde(default)]
|
||||||
|
pub run_count: u64,
|
||||||
|
#[serde(default)]
|
||||||
|
pub failure_count: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ScheduleDef {
|
||||||
|
/// True if this schedule should fire at `now`. Honors `enabled`.
|
||||||
|
pub fn is_due(&self, now: DateTime<Utc>) -> bool {
|
||||||
|
self.enabled && now >= self.next_run_at
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// =================== Trigger semantics ===================
|
||||||
|
//
|
||||||
|
// J: THIS IS YOURS TO IMPLEMENT
|
||||||
|
// =====================================================
|
||||||
|
//
|
||||||
|
// `compute_next_run_at` decides when a schedule fires next. The current
|
||||||
|
// state is the "since last completion" semantics — simple but has real
|
||||||
|
// trade-offs vs other strategies:
|
||||||
|
//
|
||||||
|
// 1. SINCE LAST COMPLETION (current default below):
|
||||||
|
// next = completed_at + interval
|
||||||
|
// Easy to reason about. A long-running ingest delays the next fire
|
||||||
|
// by however long it took. Drift is bounded — the schedule never
|
||||||
|
// "falls behind" trying to catch up.
|
||||||
|
//
|
||||||
|
// 2. ANCHORED INTERVAL:
|
||||||
|
// next = previous_target + interval, regardless of when run finished
|
||||||
|
// Keeps wall-clock alignment ("every 30min on the :00 and :30").
|
||||||
|
// Risk: if ingest takes longer than the interval, the next run is
|
||||||
|
// "due immediately" and you can fire back-to-back forever.
|
||||||
|
//
|
||||||
|
// 3. ANCHORED + SKIP:
|
||||||
|
// Anchored, but skip a tick if the previous one is still running.
|
||||||
|
// Best of both, slightly more code.
|
||||||
|
//
|
||||||
|
// 4. CRON:
|
||||||
|
// Parse the expression, find the next match after `now`. Most
|
||||||
|
// expressive but pulls in a cron parser dep.
|
||||||
|
//
|
||||||
|
// Cron is currently stubbed — falling back to "now + 1 hour" so a
|
||||||
|
// mistakenly-configured Cron schedule doesn't melt the GPU.
|
||||||
|
pub fn compute_next_run_at(
|
||||||
|
trigger: &ScheduleTrigger,
|
||||||
|
completed_at: DateTime<Utc>,
|
||||||
|
_previous_next: DateTime<Utc>,
|
||||||
|
) -> DateTime<Utc> {
|
||||||
|
// TODO(J): pick a strategy. Starter = since-last-completion.
|
||||||
|
match trigger {
|
||||||
|
ScheduleTrigger::Interval { secs } => {
|
||||||
|
completed_at + Duration::seconds(*secs as i64)
|
||||||
|
}
|
||||||
|
ScheduleTrigger::Cron { expr: _ } => {
|
||||||
|
// Cron parsing not implemented — fall back to a safe 1h
|
||||||
|
// window so a bad config can't fire-loop the system.
|
||||||
|
completed_at + Duration::hours(1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// =================== Persistence ===================
|
||||||
|
|
||||||
|
/// Storage path for a schedule's JSON record.
|
||||||
|
fn schedule_key(id: &str) -> String {
|
||||||
|
let safe: String = id
|
||||||
|
.chars()
|
||||||
|
.map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' })
|
||||||
|
.collect();
|
||||||
|
format!("{PREFIX}/{safe}.json")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// In-memory + on-disk schedule registry. The Scheduler reads from here;
|
||||||
|
/// HTTP CRUD writes through here.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct ScheduleStore {
|
||||||
|
/// Bucket where schedule JSON lives. Always primary today — a
|
||||||
|
/// future enhancement could move to per-profile schedule buckets.
|
||||||
|
store: Arc<dyn ObjectStore>,
|
||||||
|
schedules: Arc<RwLock<HashMap<String, ScheduleDef>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ScheduleStore {
|
||||||
|
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||||
|
Self {
|
||||||
|
store,
|
||||||
|
schedules: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Load every persisted schedule into memory. Call on startup.
|
||||||
|
pub async fn rebuild(&self) -> Result<usize, String> {
|
||||||
|
let keys = ops::list(&self.store, Some(&format!("{PREFIX}/")))
|
||||||
|
.await
|
||||||
|
.unwrap_or_default();
|
||||||
|
let mut count = 0usize;
|
||||||
|
for key in keys {
|
||||||
|
if !key.ends_with(".json") { continue; }
|
||||||
|
match ops::get(&self.store, &key).await {
|
||||||
|
Ok(bytes) => match serde_json::from_slice::<ScheduleDef>(&bytes) {
|
||||||
|
Ok(def) => {
|
||||||
|
self.schedules.write().await.insert(def.id.clone(), def);
|
||||||
|
count += 1;
|
||||||
|
}
|
||||||
|
Err(e) => tracing::warn!("schedule {key}: malformed: {e}"),
|
||||||
|
},
|
||||||
|
Err(e) => tracing::warn!("schedule {key}: read failed: {e}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list(&self) -> Vec<ScheduleDef> {
|
||||||
|
let mut out: Vec<ScheduleDef> = self.schedules.read().await.values().cloned().collect();
|
||||||
|
out.sort_by(|a, b| a.id.cmp(&b.id));
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get(&self, id: &str) -> Option<ScheduleDef> {
|
||||||
|
self.schedules.read().await.get(id).cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn put(&self, def: ScheduleDef) -> Result<ScheduleDef, String> {
|
||||||
|
let json = serde_json::to_vec_pretty(&def).map_err(|e| e.to_string())?;
|
||||||
|
ops::put(&self.store, &schedule_key(&def.id), json.into()).await?;
|
||||||
|
self.schedules.write().await.insert(def.id.clone(), def.clone());
|
||||||
|
Ok(def)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn delete(&self, id: &str) -> Result<bool, String> {
|
||||||
|
// Object_store's delete is idempotent — silently OKs missing keys.
|
||||||
|
let _ = ops::delete(&self.store, &schedule_key(id)).await;
|
||||||
|
Ok(self.schedules.write().await.remove(id).is_some())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Update completion metadata + persist. Used by the scheduler
|
||||||
|
/// after a run finishes (success or failure).
|
||||||
|
pub async fn record_run(
|
||||||
|
&self,
|
||||||
|
id: &str,
|
||||||
|
outcome: RunOutcome,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let mut guard = self.schedules.write().await;
|
||||||
|
let Some(def) = guard.get_mut(id) else {
|
||||||
|
return Err(format!("schedule '{id}' not registered"));
|
||||||
|
};
|
||||||
|
def.last_run_at = Some(outcome.at);
|
||||||
|
if outcome.success {
|
||||||
|
def.run_count += 1;
|
||||||
|
} else {
|
||||||
|
def.failure_count += 1;
|
||||||
|
}
|
||||||
|
def.next_run_at = compute_next_run_at(&def.trigger, outcome.at, def.next_run_at);
|
||||||
|
def.last_outcome = Some(outcome);
|
||||||
|
|
||||||
|
let json = serde_json::to_vec_pretty(def).map_err(|e| e.to_string())?;
|
||||||
|
let key = schedule_key(id);
|
||||||
|
let store = self.store.clone();
|
||||||
|
// Drop the write lock before async I/O.
|
||||||
|
drop(guard);
|
||||||
|
ops::put(&store, &key, json.into()).await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// =================== Scheduler task ===================
|
||||||
|
|
||||||
|
/// Long-running task that fires due schedules. Spawned at gateway
|
||||||
|
/// startup. Stop semantics: drop the handle (the loop exits when its
|
||||||
|
/// store is the only reference left, or sooner if you wire a stop signal).
|
||||||
|
pub struct Scheduler {
|
||||||
|
pub store: ScheduleStore,
|
||||||
|
pub ingest: SchedulerIngestDeps,
|
||||||
|
/// Polling cadence for the loop itself. The schedules' own intervals
|
||||||
|
/// are independent — this just controls how granularly we check
|
||||||
|
/// "is anything due?" Default 10s.
|
||||||
|
pub tick_secs: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// What the scheduler needs to actually run an ingest. We pass in only
|
||||||
|
/// what's needed — keeps the dep set narrow.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct SchedulerIngestDeps {
|
||||||
|
pub store: Arc<dyn ObjectStore>,
|
||||||
|
pub registry: catalogd::registry::Registry,
|
||||||
|
pub buckets: Arc<storaged::registry::BucketRegistry>,
|
||||||
|
/// Optional Phase 16.5 hook — push a DatasetAppended event if any
|
||||||
|
/// HNSW index is bound to the dataset that just got refreshed.
|
||||||
|
pub agent_handle: Option<vectord::agent::AgentHandle>,
|
||||||
|
pub index_registry: vectord::index_registry::IndexRegistry,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Scheduler {
|
||||||
|
/// Spawn the loop in the background. Returns immediately.
|
||||||
|
pub fn spawn(self) {
|
||||||
|
tokio::spawn(async move {
|
||||||
|
self.run().await;
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run(self) {
|
||||||
|
tracing::info!("scheduler started — tick={}s", self.tick_secs);
|
||||||
|
let in_flight: Arc<RwLock<HashSet<String>>> = Arc::new(RwLock::new(HashSet::new()));
|
||||||
|
|
||||||
|
loop {
|
||||||
|
tokio::time::sleep(std::time::Duration::from_secs(self.tick_secs)).await;
|
||||||
|
let now = Utc::now();
|
||||||
|
let due: Vec<ScheduleDef> = self.store.list().await
|
||||||
|
.into_iter()
|
||||||
|
.filter(|s| s.is_due(now))
|
||||||
|
.collect();
|
||||||
|
if due.is_empty() { continue; }
|
||||||
|
|
||||||
|
for def in due {
|
||||||
|
{
|
||||||
|
let in_flight_guard = in_flight.read().await;
|
||||||
|
if in_flight_guard.contains(&def.id) {
|
||||||
|
// Previous run still going — skip this tick.
|
||||||
|
// We don't reschedule; next tick re-evaluates.
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let store = self.store.clone();
|
||||||
|
let ingest = self.ingest.clone();
|
||||||
|
let in_flight = in_flight.clone();
|
||||||
|
let id = def.id.clone();
|
||||||
|
|
||||||
|
in_flight.write().await.insert(id.clone());
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let t0 = std::time::Instant::now();
|
||||||
|
let outcome = run_schedule(&def, &ingest).await;
|
||||||
|
let elapsed = t0.elapsed().as_secs_f32();
|
||||||
|
let outcome = RunOutcome {
|
||||||
|
at: Utc::now(),
|
||||||
|
success: outcome.is_ok(),
|
||||||
|
message: match &outcome {
|
||||||
|
Ok(s) => s.clone(),
|
||||||
|
Err(e) => e.clone(),
|
||||||
|
},
|
||||||
|
rows: outcome.as_ref().ok().and_then(parse_rows_from_message),
|
||||||
|
duration_secs: elapsed,
|
||||||
|
};
|
||||||
|
if let Err(e) = store.record_run(&id, outcome).await {
|
||||||
|
tracing::warn!("scheduler: record_run({id}): {e}");
|
||||||
|
}
|
||||||
|
in_flight.write().await.remove(&id);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pluck a "rows=N" hint out of a success message — best-effort, no
|
||||||
|
/// hard contract. Lets us surface row counts in the schedule status
|
||||||
|
/// endpoint without forcing every ingest path to return structured
|
||||||
|
/// stats here.
|
||||||
|
fn parse_rows_from_message(msg: &String) -> Option<usize> {
|
||||||
|
msg.split_whitespace()
|
||||||
|
.find_map(|tok| tok.strip_prefix("rows=").and_then(|n| n.parse().ok()))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fire one schedule's ingest. Returns Ok(success_message) or Err(reason).
|
||||||
|
pub async fn run_schedule(
|
||||||
|
def: &ScheduleDef,
|
||||||
|
deps: &SchedulerIngestDeps,
|
||||||
|
) -> Result<String, String> {
|
||||||
|
tracing::info!("scheduler: firing '{}' ({})", def.id, def.kind.label());
|
||||||
|
match &def.kind {
|
||||||
|
ScheduleKind::Mysql { dsn, table, dataset_name, batch_size, order_by } => {
|
||||||
|
let req = crate::my_stream::MyStreamRequest {
|
||||||
|
dsn: dsn.clone(),
|
||||||
|
table: table.clone(),
|
||||||
|
dataset_name: dataset_name.clone(),
|
||||||
|
batch_size: *batch_size,
|
||||||
|
order_by: order_by.clone(),
|
||||||
|
limit: None,
|
||||||
|
};
|
||||||
|
run_mysql(req, deps).await
|
||||||
|
}
|
||||||
|
ScheduleKind::Postgres { dsn, table, dataset_name, batch_size, order_by } => {
|
||||||
|
let req = crate::pg_stream::PgStreamRequest {
|
||||||
|
dsn: dsn.clone(),
|
||||||
|
table: table.clone(),
|
||||||
|
dataset_name: dataset_name.clone(),
|
||||||
|
batch_size: *batch_size,
|
||||||
|
order_by: order_by.clone(),
|
||||||
|
limit: None,
|
||||||
|
};
|
||||||
|
run_postgres(req, deps).await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run_mysql(
|
||||||
|
req: crate::my_stream::MyStreamRequest,
|
||||||
|
deps: &SchedulerIngestDeps,
|
||||||
|
) -> Result<String, String> {
|
||||||
|
let (parquet, result) = crate::my_stream::stream_table_to_parquet(&req).await?;
|
||||||
|
if result.rows == 0 {
|
||||||
|
return Ok(format!("rows=0 — table empty"));
|
||||||
|
}
|
||||||
|
persist_and_register(
|
||||||
|
&parquet, &result.table, req.dataset_name.as_deref(),
|
||||||
|
result.rows, "mysql", &req.dsn, deps,
|
||||||
|
).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run_postgres(
|
||||||
|
req: crate::pg_stream::PgStreamRequest,
|
||||||
|
deps: &SchedulerIngestDeps,
|
||||||
|
) -> Result<String, String> {
|
||||||
|
let (parquet, result) = crate::pg_stream::stream_table_to_parquet(&req).await?;
|
||||||
|
if result.rows == 0 {
|
||||||
|
return Ok(format!("rows=0 — table empty"));
|
||||||
|
}
|
||||||
|
persist_and_register(
|
||||||
|
&parquet, &result.table, req.dataset_name.as_deref(),
|
||||||
|
result.rows, "postgresql", &req.dsn, deps,
|
||||||
|
).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Shared write+register path. Mirrors the HTTP ingest handlers in
|
||||||
|
/// service.rs so the scheduled path produces identical catalog output:
|
||||||
|
/// PII detection, lineage with redacted DSN, mark-stale + agent trigger.
|
||||||
|
async fn persist_and_register(
|
||||||
|
parquet: &bytes::Bytes,
|
||||||
|
table: &str,
|
||||||
|
dataset_name_override: Option<&str>,
|
||||||
|
rows: usize,
|
||||||
|
source_system: &str,
|
||||||
|
dsn: &str,
|
||||||
|
deps: &SchedulerIngestDeps,
|
||||||
|
) -> Result<String, String> {
|
||||||
|
use shared::types::{ColumnMeta, Lineage, ObjectRef, SchemaFingerprint, Sensitivity};
|
||||||
|
|
||||||
|
let (schema, _) = shared::arrow_helpers::parquet_to_record_batches(parquet)
|
||||||
|
.map_err(|e| format!("reparse parquet: {e}"))?;
|
||||||
|
|
||||||
|
let dataset_name = dataset_name_override.unwrap_or(table).to_string();
|
||||||
|
let storage_key = format!("datasets/{}.parquet", dataset_name);
|
||||||
|
let parquet_size = parquet.len() as u64;
|
||||||
|
let bucket = "primary".to_string();
|
||||||
|
|
||||||
|
let target_store = deps.buckets.get(&bucket)?;
|
||||||
|
ops::put(&target_store, &storage_key, parquet.clone()).await?;
|
||||||
|
|
||||||
|
let schema_fp = shared::arrow_helpers::fingerprint_schema(&schema);
|
||||||
|
let now = Utc::now();
|
||||||
|
deps.registry.register(
|
||||||
|
dataset_name.clone(),
|
||||||
|
SchemaFingerprint(schema_fp.0),
|
||||||
|
vec![ObjectRef {
|
||||||
|
bucket: bucket.clone(),
|
||||||
|
key: storage_key.clone(),
|
||||||
|
size_bytes: parquet_size,
|
||||||
|
created_at: now,
|
||||||
|
}],
|
||||||
|
).await?;
|
||||||
|
|
||||||
|
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
|
||||||
|
let sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
|
||||||
|
let columns: Vec<ColumnMeta> = schema.fields().iter().map(|f| {
|
||||||
|
let sens = shared::pii::detect_sensitivity(f.name());
|
||||||
|
ColumnMeta {
|
||||||
|
name: f.name().clone(),
|
||||||
|
data_type: f.data_type().to_string(),
|
||||||
|
sensitivity: sens.clone(),
|
||||||
|
description: String::new(),
|
||||||
|
is_pii: matches!(sens, Some(Sensitivity::Pii)),
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
|
||||||
|
let lineage = Lineage {
|
||||||
|
source_system: source_system.to_string(),
|
||||||
|
source_file: format!("dsn: {}", redact_dsn(dsn)),
|
||||||
|
ingest_job: format!("scheduled-{}-{}", source_system, now.timestamp_millis()),
|
||||||
|
ingest_timestamp: now,
|
||||||
|
parent_datasets: vec![],
|
||||||
|
};
|
||||||
|
|
||||||
|
let _ = deps.registry.update_metadata(&dataset_name, catalogd::registry::MetadataUpdate {
|
||||||
|
sensitivity,
|
||||||
|
columns: Some(columns),
|
||||||
|
lineage: Some(lineage),
|
||||||
|
row_count: Some(rows as u64),
|
||||||
|
..Default::default()
|
||||||
|
}).await;
|
||||||
|
|
||||||
|
let _ = deps.registry.mark_embeddings_stale(&dataset_name).await;
|
||||||
|
|
||||||
|
if let Some(agent) = &deps.agent_handle {
|
||||||
|
let bound = deps.index_registry.list(Some(&dataset_name), None).await;
|
||||||
|
for meta in bound {
|
||||||
|
let event = vectord::agent::TriggerEvent::dataset_appended(
|
||||||
|
meta.index_name.clone(), &dataset_name,
|
||||||
|
);
|
||||||
|
let _ = agent.enqueue(event).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(format!("rows={rows} dataset={dataset_name} bytes={parquet_size}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Same generic redactor used in service.rs — duplicated here to avoid
|
||||||
|
/// pulling service::redact_dsn behind a `pub` visibility just for this.
|
||||||
|
fn redact_dsn(dsn: &str) -> String {
|
||||||
|
let scheme_end = match dsn.find("://") {
|
||||||
|
Some(i) => i + 3,
|
||||||
|
None => return dsn.to_string(),
|
||||||
|
};
|
||||||
|
let at_idx = match dsn.rfind('@') {
|
||||||
|
Some(i) if i > scheme_end => i,
|
||||||
|
_ => return dsn.to_string(),
|
||||||
|
};
|
||||||
|
let userpass = &dsn[scheme_end..at_idx];
|
||||||
|
let colon_offset = match userpass.find(':') {
|
||||||
|
Some(i) => i,
|
||||||
|
None => return dsn.to_string(),
|
||||||
|
};
|
||||||
|
let colon_idx = scheme_end + colon_offset;
|
||||||
|
format!("{}:***{}", &dsn[..colon_idx], &dsn[at_idx..])
|
||||||
|
}
|
||||||
|
|
||||||
|
// =================== Tests ===================
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn mk(secs: u64) -> ScheduleDef {
|
||||||
|
ScheduleDef {
|
||||||
|
id: "test".into(),
|
||||||
|
kind: ScheduleKind::Mysql {
|
||||||
|
dsn: "mysql://x@host/db".into(),
|
||||||
|
table: "t".into(),
|
||||||
|
dataset_name: None,
|
||||||
|
batch_size: None,
|
||||||
|
order_by: None,
|
||||||
|
},
|
||||||
|
trigger: ScheduleTrigger::Interval { secs },
|
||||||
|
enabled: true,
|
||||||
|
created_at: Utc::now(),
|
||||||
|
created_by: "tester".into(),
|
||||||
|
next_run_at: Utc::now(),
|
||||||
|
last_run_at: None,
|
||||||
|
last_outcome: None,
|
||||||
|
run_count: 0,
|
||||||
|
failure_count: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn is_due_respects_enabled() {
|
||||||
|
let mut s = mk(60);
|
||||||
|
s.next_run_at = Utc::now() - Duration::seconds(5);
|
||||||
|
assert!(s.is_due(Utc::now()));
|
||||||
|
s.enabled = false;
|
||||||
|
assert!(!s.is_due(Utc::now()));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn is_due_respects_future_time() {
|
||||||
|
let mut s = mk(60);
|
||||||
|
s.next_run_at = Utc::now() + Duration::seconds(60);
|
||||||
|
assert!(!s.is_due(Utc::now()));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn next_run_since_last_completion_advances() {
|
||||||
|
let s = mk(120);
|
||||||
|
let completed = Utc::now();
|
||||||
|
let next = compute_next_run_at(&s.trigger, completed, completed);
|
||||||
|
let delta = next - completed;
|
||||||
|
assert_eq!(delta.num_seconds(), 120);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cron_falls_back_to_one_hour() {
|
||||||
|
let trig = ScheduleTrigger::Cron { expr: "* * * * *".into() };
|
||||||
|
let now = Utc::now();
|
||||||
|
let next = compute_next_run_at(&trig, now, now);
|
||||||
|
let delta = next - now;
|
||||||
|
assert_eq!(delta.num_seconds(), 3600);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn schedule_kind_label() {
|
||||||
|
let k = ScheduleKind::Mysql {
|
||||||
|
dsn: "x".into(), table: "customers".into(),
|
||||||
|
dataset_name: None, batch_size: None, order_by: None,
|
||||||
|
};
|
||||||
|
assert_eq!(k.label(), "mysql:customers");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn redact_dsn_handles_mysql_and_postgres() {
|
||||||
|
assert_eq!(redact_dsn("mysql://u:secret@h/db"), "mysql://u:***@h/db");
|
||||||
|
assert_eq!(redact_dsn("postgresql://u:secret@h:5432/db"), "postgresql://u:***@h:5432/db");
|
||||||
|
assert_eq!(redact_dsn("mysql://u@h/db"), "mysql://u@h/db"); // no password
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,279 @@
|
|||||||
|
/// Schema evolution: detect changes between ingests and suggest migration rules.
|
||||||
|
/// When a source changes format (columns renamed, added, removed, type changed),
|
||||||
|
/// the system detects the diff and can auto-map using AI or heuristic matching.
|
||||||
|
|
||||||
|
use arrow::datatypes::{DataType, Schema, SchemaRef};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// A detected change between two schema versions.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct SchemaChange {
|
||||||
|
pub change_type: ChangeType,
|
||||||
|
pub column: String,
|
||||||
|
pub detail: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
#[serde(rename_all = "snake_case")]
|
||||||
|
pub enum ChangeType {
|
||||||
|
Added, // new column in incoming schema
|
||||||
|
Removed, // column in old schema not in new
|
||||||
|
TypeChanged, // same name, different type
|
||||||
|
Renamed, // likely renamed (fuzzy match)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A migration rule — how to map old schema to new.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct MigrationRule {
|
||||||
|
pub source_column: String,
|
||||||
|
pub target_column: String,
|
||||||
|
pub transform: String, // "direct", "cast", "concat", "split", "drop"
|
||||||
|
pub confidence: f32, // 0.0 - 1.0
|
||||||
|
pub auto_generated: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Compare two schemas and detect changes.
|
||||||
|
pub fn diff_schemas(old: &SchemaRef, new: &SchemaRef) -> Vec<SchemaChange> {
|
||||||
|
let mut changes = Vec::new();
|
||||||
|
|
||||||
|
let old_fields: HashMap<&str, &DataType> = old.fields().iter()
|
||||||
|
.map(|f| (f.name().as_str(), f.data_type()))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let new_fields: HashMap<&str, &DataType> = new.fields().iter()
|
||||||
|
.map(|f| (f.name().as_str(), f.data_type()))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// Check for removed columns
|
||||||
|
for (name, _dt) in &old_fields {
|
||||||
|
if !new_fields.contains_key(name) {
|
||||||
|
// Check if it was renamed (fuzzy match)
|
||||||
|
let possible_rename = find_similar_column(name, &new_fields, &old_fields);
|
||||||
|
if let Some((new_name, confidence)) = possible_rename {
|
||||||
|
changes.push(SchemaChange {
|
||||||
|
change_type: ChangeType::Renamed,
|
||||||
|
column: name.to_string(),
|
||||||
|
detail: format!("likely renamed to '{}' (confidence: {:.0}%)", new_name, confidence * 100.0),
|
||||||
|
});
|
||||||
|
} else {
|
||||||
|
changes.push(SchemaChange {
|
||||||
|
change_type: ChangeType::Removed,
|
||||||
|
column: name.to_string(),
|
||||||
|
detail: "column no longer present in incoming data".into(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check for added or type-changed columns
|
||||||
|
for (name, new_dt) in &new_fields {
|
||||||
|
match old_fields.get(name) {
|
||||||
|
None => {
|
||||||
|
// Skip if we already flagged this as a rename target
|
||||||
|
let is_rename_target = changes.iter().any(|c| {
|
||||||
|
c.change_type == ChangeType::Renamed && c.detail.contains(name)
|
||||||
|
});
|
||||||
|
if !is_rename_target {
|
||||||
|
changes.push(SchemaChange {
|
||||||
|
change_type: ChangeType::Added,
|
||||||
|
column: name.to_string(),
|
||||||
|
detail: format!("new column (type: {})", new_dt),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(old_dt) => {
|
||||||
|
if old_dt != new_dt {
|
||||||
|
changes.push(SchemaChange {
|
||||||
|
change_type: ChangeType::TypeChanged,
|
||||||
|
column: name.to_string(),
|
||||||
|
detail: format!("{} → {}", old_dt, new_dt),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
changes
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Generate migration rules from detected changes.
|
||||||
|
pub fn generate_rules(changes: &[SchemaChange]) -> Vec<MigrationRule> {
|
||||||
|
changes.iter().filter_map(|change| {
|
||||||
|
match change.change_type {
|
||||||
|
ChangeType::Renamed => {
|
||||||
|
// Extract the new name from the detail string
|
||||||
|
let new_name = change.detail
|
||||||
|
.split('\'')
|
||||||
|
.nth(1)
|
||||||
|
.unwrap_or(&change.column);
|
||||||
|
Some(MigrationRule {
|
||||||
|
source_column: change.column.clone(),
|
||||||
|
target_column: new_name.to_string(),
|
||||||
|
transform: "direct".into(),
|
||||||
|
confidence: 0.8,
|
||||||
|
auto_generated: true,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
ChangeType::TypeChanged => {
|
||||||
|
Some(MigrationRule {
|
||||||
|
source_column: change.column.clone(),
|
||||||
|
target_column: change.column.clone(),
|
||||||
|
transform: "cast".into(),
|
||||||
|
confidence: 0.9,
|
||||||
|
auto_generated: true,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
ChangeType::Added => {
|
||||||
|
Some(MigrationRule {
|
||||||
|
source_column: String::new(),
|
||||||
|
target_column: change.column.clone(),
|
||||||
|
transform: "new_column".into(),
|
||||||
|
confidence: 1.0,
|
||||||
|
auto_generated: true,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
ChangeType::Removed => {
|
||||||
|
Some(MigrationRule {
|
||||||
|
source_column: change.column.clone(),
|
||||||
|
target_column: String::new(),
|
||||||
|
transform: "drop".into(),
|
||||||
|
confidence: 0.7,
|
||||||
|
auto_generated: true,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build an AI prompt to suggest migration rules for complex cases.
|
||||||
|
pub fn build_migration_prompt(
|
||||||
|
old_schema: &SchemaRef,
|
||||||
|
new_schema: &SchemaRef,
|
||||||
|
changes: &[SchemaChange],
|
||||||
|
) -> String {
|
||||||
|
let old_cols: Vec<String> = old_schema.fields().iter()
|
||||||
|
.map(|f| format!(" {} ({})", f.name(), f.data_type()))
|
||||||
|
.collect();
|
||||||
|
let new_cols: Vec<String> = new_schema.fields().iter()
|
||||||
|
.map(|f| format!(" {} ({})", f.name(), f.data_type()))
|
||||||
|
.collect();
|
||||||
|
let change_list: Vec<String> = changes.iter()
|
||||||
|
.map(|c| format!(" [{:?}] {} — {}", c.change_type, c.column, c.detail))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
format!(
|
||||||
|
"You are a data migration expert. A data source has changed its schema.\n\n\
|
||||||
|
OLD SCHEMA:\n{}\n\n\
|
||||||
|
NEW SCHEMA:\n{}\n\n\
|
||||||
|
DETECTED CHANGES:\n{}\n\n\
|
||||||
|
For each change, suggest a migration rule as JSON:\n\
|
||||||
|
{{\"source_column\": \"old_name\", \"target_column\": \"new_name\", \"transform\": \"direct|cast|concat|split|drop\", \"confidence\": 0.0-1.0}}\n\n\
|
||||||
|
Output ONLY a JSON array of rules, nothing else.",
|
||||||
|
old_cols.join("\n"),
|
||||||
|
new_cols.join("\n"),
|
||||||
|
change_list.join("\n"),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fuzzy match: find a column in new schema that's similar to a removed old column.
|
||||||
|
fn find_similar_column<'a>(
|
||||||
|
old_name: &str,
|
||||||
|
new_fields: &HashMap<&'a str, &DataType>,
|
||||||
|
old_fields: &HashMap<&str, &DataType>,
|
||||||
|
) -> Option<(String, f32)> {
|
||||||
|
let old_lower = old_name.to_lowercase();
|
||||||
|
let old_parts: Vec<&str> = old_lower.split('_').collect();
|
||||||
|
|
||||||
|
let mut best_match: Option<(String, f32)> = None;
|
||||||
|
|
||||||
|
for (new_name, _) in new_fields {
|
||||||
|
// Skip columns that already exist in old schema (not renamed)
|
||||||
|
if old_fields.contains_key(new_name) { continue; }
|
||||||
|
|
||||||
|
let new_lower = new_name.to_lowercase();
|
||||||
|
let new_parts: Vec<&str> = new_lower.split('_').collect();
|
||||||
|
|
||||||
|
// Exact substring match
|
||||||
|
if old_lower.contains(&new_lower) || new_lower.contains(&old_lower) {
|
||||||
|
let confidence = 0.85;
|
||||||
|
if best_match.as_ref().map_or(true, |(_, c)| confidence > *c) {
|
||||||
|
best_match = Some((new_name.to_string(), confidence));
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Shared word parts (e.g., "first_name" → "fname", "full_name")
|
||||||
|
let shared: usize = old_parts.iter()
|
||||||
|
.filter(|p| new_parts.iter().any(|np| np.contains(*p) || p.contains(np)))
|
||||||
|
.count();
|
||||||
|
|
||||||
|
if shared > 0 {
|
||||||
|
let confidence = shared as f32 / old_parts.len().max(new_parts.len()) as f32 * 0.75;
|
||||||
|
if confidence >= 0.3 && best_match.as_ref().map_or(true, |(_, c)| confidence > *c) {
|
||||||
|
best_match = Some((new_name.to_string(), confidence));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
best_match
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use arrow::datatypes::Field;
|
||||||
|
|
||||||
|
fn schema(fields: Vec<(&str, DataType)>) -> SchemaRef {
|
||||||
|
Arc::new(Schema::new(
|
||||||
|
fields.into_iter().map(|(n, t)| Field::new(n, t, true)).collect::<Vec<_>>()
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn detect_added_column() {
|
||||||
|
let old = schema(vec![("id", DataType::Int64), ("name", DataType::Utf8)]);
|
||||||
|
let new = schema(vec![("id", DataType::Int64), ("name", DataType::Utf8), ("email", DataType::Utf8)]);
|
||||||
|
let changes = diff_schemas(&old, &new);
|
||||||
|
assert_eq!(changes.len(), 1);
|
||||||
|
assert_eq!(changes[0].change_type, ChangeType::Added);
|
||||||
|
assert_eq!(changes[0].column, "email");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn detect_removed_column() {
|
||||||
|
let old = schema(vec![("id", DataType::Int64), ("name", DataType::Utf8), ("age", DataType::Int64)]);
|
||||||
|
let new = schema(vec![("id", DataType::Int64), ("name", DataType::Utf8)]);
|
||||||
|
let changes = diff_schemas(&old, &new);
|
||||||
|
assert!(changes.iter().any(|c| c.column == "age" && c.change_type == ChangeType::Removed));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn detect_type_change() {
|
||||||
|
let old = schema(vec![("id", DataType::Int64), ("score", DataType::Int64)]);
|
||||||
|
let new = schema(vec![("id", DataType::Int64), ("score", DataType::Float64)]);
|
||||||
|
let changes = diff_schemas(&old, &new);
|
||||||
|
assert_eq!(changes.len(), 1);
|
||||||
|
assert_eq!(changes[0].change_type, ChangeType::TypeChanged);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn detect_rename() {
|
||||||
|
let old = schema(vec![("first_name", DataType::Utf8), ("last_name", DataType::Utf8)]);
|
||||||
|
let new = schema(vec![("full_name", DataType::Utf8)]);
|
||||||
|
let changes = diff_schemas(&old, &new);
|
||||||
|
// Should detect first_name → full_name as a potential rename
|
||||||
|
assert!(changes.iter().any(|c| c.change_type == ChangeType::Renamed));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn generate_rules_from_changes() {
|
||||||
|
let old = schema(vec![("id", DataType::Int64), ("name", DataType::Utf8)]);
|
||||||
|
let new = schema(vec![("id", DataType::Int64), ("name", DataType::Utf8), ("email", DataType::Utf8)]);
|
||||||
|
let changes = diff_schemas(&old, &new);
|
||||||
|
let rules = generate_rules(&changes);
|
||||||
|
assert_eq!(rules.len(), 1);
|
||||||
|
assert_eq!(rules[0].transform, "new_column");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,624 @@
|
|||||||
|
use axum::{
|
||||||
|
Json, Router,
|
||||||
|
extract::{Multipart, Path, Query, State},
|
||||||
|
http::{HeaderMap, StatusCode},
|
||||||
|
response::IntoResponse,
|
||||||
|
routing::{delete, get, patch, post},
|
||||||
|
};
|
||||||
|
use bytes::Bytes;
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use serde::Deserialize;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use catalogd::registry::Registry;
|
||||||
|
use crate::{db_ingest, my_stream, pg_stream, pipeline, schedule};
|
||||||
|
use shared::arrow_helpers::record_batch_to_parquet;
|
||||||
|
use shared::types::{ObjectRef, SchemaFingerprint};
|
||||||
|
use storaged::ops;
|
||||||
|
use storaged::registry::BucketRegistry;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct IngestState {
|
||||||
|
pub store: Arc<dyn ObjectStore>,
|
||||||
|
pub registry: Registry,
|
||||||
|
/// Federation layer 2: lookup target bucket from request headers.
|
||||||
|
pub buckets: Arc<BucketRegistry>,
|
||||||
|
/// Phase 16.5: when ingest marks a dataset's embeddings stale, push
|
||||||
|
/// a `DatasetAppended` trigger so the autotune agent can schedule
|
||||||
|
/// re-trials. Holds the agent handle (no-op when agent disabled).
|
||||||
|
pub agent_handle: vectord::agent::AgentHandle,
|
||||||
|
/// Used to look up which HNSW indexes are attached to the
|
||||||
|
/// just-ingested dataset. Each matching index gets one trigger.
|
||||||
|
pub index_registry: vectord::index_registry::IndexRegistry,
|
||||||
|
/// Scheduled-ingest registry. The scheduler task runs against this
|
||||||
|
/// store; HTTP CRUD endpoints write through it.
|
||||||
|
pub schedules: schedule::ScheduleStore,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Push `DatasetAppended` triggers for every HNSW index bound to this
|
||||||
|
/// dataset. Called after a successful ingest. Logs failures but never
|
||||||
|
/// fails the ingest — the agent is best-effort, not load-bearing.
|
||||||
|
async fn notify_agent_on_append(state: &IngestState, dataset_name: &str) {
|
||||||
|
let indexes = state.index_registry.list(Some(dataset_name), None).await;
|
||||||
|
if indexes.is_empty() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
for meta in indexes {
|
||||||
|
let event = vectord::agent::TriggerEvent::dataset_appended(
|
||||||
|
meta.index_name.clone(), dataset_name,
|
||||||
|
);
|
||||||
|
if let Err(e) = state.agent_handle.enqueue(event).await {
|
||||||
|
tracing::warn!("agent enqueue failed for '{}': {}", meta.index_name, e);
|
||||||
|
} else {
|
||||||
|
tracing::info!(
|
||||||
|
"agent: enqueued DatasetAppended({}) for index '{}'",
|
||||||
|
dataset_name, meta.index_name,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolve the target bucket from `X-Lakehouse-Bucket` header.
|
||||||
|
/// Returns `(bucket_name, store_for_writes)`. Falls back to "primary"
|
||||||
|
/// when the header is absent. Returns Err with the canonical bucket
|
||||||
|
/// list when the header names an unknown bucket.
|
||||||
|
fn resolve_bucket(
|
||||||
|
headers: &HeaderMap,
|
||||||
|
buckets: &BucketRegistry,
|
||||||
|
) -> Result<(String, Arc<dyn ObjectStore>), (StatusCode, String)> {
|
||||||
|
let target = headers
|
||||||
|
.get("x-lakehouse-bucket")
|
||||||
|
.and_then(|h| h.to_str().ok())
|
||||||
|
.map(|s| s.to_string())
|
||||||
|
.unwrap_or_else(|| buckets.default_name().to_string());
|
||||||
|
|
||||||
|
match buckets.get(&target) {
|
||||||
|
Ok(store) => Ok((target, store)),
|
||||||
|
Err(_) => Err((
|
||||||
|
StatusCode::NOT_FOUND,
|
||||||
|
format!("unknown bucket '{}' — use GET /storage/buckets to list", target),
|
||||||
|
)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn router(state: IngestState) -> Router {
|
||||||
|
Router::new()
|
||||||
|
.route("/health", get(health))
|
||||||
|
.route("/file", post(ingest_file))
|
||||||
|
.route("/postgres/tables", post(list_pg_tables))
|
||||||
|
.route("/postgres/import", post(import_pg_table))
|
||||||
|
.route("/db", post(ingest_db_stream))
|
||||||
|
.route("/mysql", post(ingest_mysql_stream))
|
||||||
|
// Phase E: scheduled ingest
|
||||||
|
.route("/schedules", get(list_schedules).post(create_schedule))
|
||||||
|
.route("/schedules/{id}", get(get_schedule)
|
||||||
|
.patch(patch_schedule)
|
||||||
|
.delete(delete_schedule))
|
||||||
|
.route("/schedules/{id}/run-now", post(run_schedule_now))
|
||||||
|
.with_state(state)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn health() -> &'static str {
|
||||||
|
"ingestd ok"
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct IngestQuery {
|
||||||
|
name: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn ingest_file(
|
||||||
|
State(state): State<IngestState>,
|
||||||
|
Query(query): Query<IngestQuery>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
mut multipart: Multipart,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let (bucket, store) = resolve_bucket(&headers, &state.buckets)?;
|
||||||
|
|
||||||
|
let field = match multipart.next_field().await {
|
||||||
|
Ok(Some(f)) => f,
|
||||||
|
Ok(None) => return Err((StatusCode::BAD_REQUEST, "no file uploaded".to_string())),
|
||||||
|
Err(e) => return Err((StatusCode::BAD_REQUEST, format!("multipart error: {e}"))),
|
||||||
|
};
|
||||||
|
|
||||||
|
let filename = field.file_name().unwrap_or("unknown").to_string();
|
||||||
|
let content = field.bytes().await
|
||||||
|
.map_err(|e| (StatusCode::BAD_REQUEST, format!("read error: {e}")))?;
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
"ingest '{}' ({} bytes) -> bucket={}",
|
||||||
|
filename, content.len(), bucket,
|
||||||
|
);
|
||||||
|
|
||||||
|
let dataset_name = query.name.as_deref();
|
||||||
|
|
||||||
|
match pipeline::ingest_file_to_bucket(&filename, &content, dataset_name, &bucket, &store, &state.registry).await {
|
||||||
|
Ok(result) => {
|
||||||
|
if !result.deduplicated {
|
||||||
|
notify_agent_on_append(&state, &result.dataset_name).await;
|
||||||
|
}
|
||||||
|
if result.deduplicated {
|
||||||
|
Ok((StatusCode::OK, Json(result)))
|
||||||
|
} else {
|
||||||
|
Ok((StatusCode::CREATED, Json(result)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- PostgreSQL Import ---
|
||||||
|
|
||||||
|
/// List tables in a PostgreSQL database.
|
||||||
|
async fn list_pg_tables(
|
||||||
|
Json(config): Json<db_ingest::DbConfig>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match db_ingest::list_postgres_tables(&config).await {
|
||||||
|
Ok(tables) => Ok(Json(tables)),
|
||||||
|
Err(e) => Err((StatusCode::BAD_GATEWAY, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct PgImportRequest {
|
||||||
|
#[serde(flatten)]
|
||||||
|
config: db_ingest::DbConfig,
|
||||||
|
table: String,
|
||||||
|
/// Override dataset name (defaults to table name)
|
||||||
|
dataset_name: Option<String>,
|
||||||
|
/// Max rows to import (None = all)
|
||||||
|
limit: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Import a PostgreSQL table into the lakehouse.
|
||||||
|
async fn import_pg_table(
|
||||||
|
State(state): State<IngestState>,
|
||||||
|
Json(req): Json<PgImportRequest>,
|
||||||
|
) -> Result<(StatusCode, Json<serde_json::Value>), (StatusCode, String)> {
|
||||||
|
tracing::info!("importing postgres table '{}' from {}:{}/{}",
|
||||||
|
req.table, req.config.host, req.config.port, req.config.database);
|
||||||
|
|
||||||
|
// Import from Postgres
|
||||||
|
let (schema, batches, db_result) = db_ingest::import_postgres_table(
|
||||||
|
&req.config, &req.table, req.limit,
|
||||||
|
).await.map_err(|e| (StatusCode::BAD_GATEWAY, e))?;
|
||||||
|
|
||||||
|
if batches.is_empty() || db_result.rows == 0 {
|
||||||
|
return Ok((StatusCode::OK, Json(serde_json::json!({
|
||||||
|
"table": req.table,
|
||||||
|
"rows": 0,
|
||||||
|
"message": "table is empty",
|
||||||
|
}))));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Convert to Parquet
|
||||||
|
let mut all_parquet = Vec::new();
|
||||||
|
for batch in &batches {
|
||||||
|
let pq = record_batch_to_parquet(batch)
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||||
|
all_parquet.extend_from_slice(&pq);
|
||||||
|
}
|
||||||
|
|
||||||
|
let dataset_name = req.dataset_name.unwrap_or_else(|| req.table.clone());
|
||||||
|
let storage_key = format!("datasets/{}.parquet", dataset_name);
|
||||||
|
let parquet_size = all_parquet.len() as u64;
|
||||||
|
|
||||||
|
// Store
|
||||||
|
ops::put(&state.store, &storage_key, Bytes::from(all_parquet))
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||||
|
|
||||||
|
// Register
|
||||||
|
let schema_fp = shared::arrow_helpers::fingerprint_schema(&schema);
|
||||||
|
let now = chrono::Utc::now();
|
||||||
|
state.registry.register(
|
||||||
|
dataset_name.clone(),
|
||||||
|
SchemaFingerprint(schema_fp.0),
|
||||||
|
vec![ObjectRef {
|
||||||
|
bucket: "data".to_string(),
|
||||||
|
key: storage_key.clone(),
|
||||||
|
size_bytes: parquet_size,
|
||||||
|
created_at: now,
|
||||||
|
}],
|
||||||
|
).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||||
|
|
||||||
|
// Auto-populate metadata (PII detection, lineage)
|
||||||
|
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
|
||||||
|
let sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
|
||||||
|
let columns: Vec<shared::types::ColumnMeta> = schema.fields().iter().map(|f| {
|
||||||
|
let sens = shared::pii::detect_sensitivity(f.name());
|
||||||
|
shared::types::ColumnMeta {
|
||||||
|
name: f.name().clone(),
|
||||||
|
data_type: f.data_type().to_string(),
|
||||||
|
sensitivity: sens.clone(),
|
||||||
|
description: String::new(),
|
||||||
|
is_pii: matches!(sens, Some(shared::types::Sensitivity::Pii)),
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
|
||||||
|
let lineage = shared::types::Lineage {
|
||||||
|
source_system: "postgresql".to_string(),
|
||||||
|
source_file: format!("{}:{}/{}.{}", req.config.host, req.config.port, req.config.database, req.table),
|
||||||
|
ingest_job: format!("pg-import-{}", now.timestamp_millis()),
|
||||||
|
ingest_timestamp: now,
|
||||||
|
parent_datasets: vec![],
|
||||||
|
};
|
||||||
|
|
||||||
|
let _ = state.registry.update_metadata(&dataset_name, catalogd::registry::MetadataUpdate {
|
||||||
|
sensitivity,
|
||||||
|
columns: Some(columns),
|
||||||
|
lineage: Some(lineage),
|
||||||
|
row_count: Some(db_result.rows as u64),
|
||||||
|
..Default::default()
|
||||||
|
}).await;
|
||||||
|
|
||||||
|
tracing::info!("imported '{}' from postgres: {} rows → {}", dataset_name, db_result.rows, storage_key);
|
||||||
|
|
||||||
|
Ok((StatusCode::CREATED, Json(serde_json::json!({
|
||||||
|
"dataset_name": dataset_name,
|
||||||
|
"table": req.table,
|
||||||
|
"rows": db_result.rows,
|
||||||
|
"columns": db_result.columns,
|
||||||
|
"schema": db_result.schema_detected,
|
||||||
|
"storage_key": storage_key,
|
||||||
|
"size_bytes": parquet_size,
|
||||||
|
}))))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Streaming DSN-based ingest. Paginates via ORDER BY + LIMIT/OFFSET so
|
||||||
|
/// large tables don't blow up memory. This is the path the task spec calls
|
||||||
|
/// `POST /ingest/db` with `{dsn, table, dataset_name}`.
|
||||||
|
async fn ingest_db_stream(
|
||||||
|
State(state): State<IngestState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(req): Json<pg_stream::PgStreamRequest>,
|
||||||
|
) -> Result<(StatusCode, Json<serde_json::Value>), (StatusCode, String)> {
|
||||||
|
let (bucket, store) = resolve_bucket(&headers, &state.buckets)?;
|
||||||
|
tracing::info!(
|
||||||
|
"pg stream ingest: table='{}' dataset='{:?}' bucket='{}' batch_size={:?}",
|
||||||
|
req.table, req.dataset_name, bucket, req.batch_size,
|
||||||
|
);
|
||||||
|
|
||||||
|
// Stream from Postgres into Parquet bytes.
|
||||||
|
let (parquet, stream_result) = pg_stream::stream_table_to_parquet(&req)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::BAD_GATEWAY, e))?;
|
||||||
|
|
||||||
|
if stream_result.rows == 0 {
|
||||||
|
return Ok((StatusCode::OK, Json(serde_json::json!({
|
||||||
|
"table": req.table,
|
||||||
|
"rows": 0,
|
||||||
|
"message": "table is empty",
|
||||||
|
}))));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Recover schema from the parquet footer — keeps a single source of truth
|
||||||
|
// for the schema fingerprint + column metadata.
|
||||||
|
let (schema, _) = shared::arrow_helpers::parquet_to_record_batches(&parquet)
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("reparse parquet: {e}")))?;
|
||||||
|
|
||||||
|
let dataset_name = req.dataset_name.clone().unwrap_or_else(|| req.table.clone());
|
||||||
|
let storage_key = format!("datasets/{}.parquet", dataset_name);
|
||||||
|
let size_bytes = parquet.len() as u64;
|
||||||
|
|
||||||
|
ops::put(&store, &storage_key, parquet)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||||
|
|
||||||
|
let schema_fp = shared::arrow_helpers::fingerprint_schema(&schema);
|
||||||
|
let now = chrono::Utc::now();
|
||||||
|
state.registry.register(
|
||||||
|
dataset_name.clone(),
|
||||||
|
SchemaFingerprint(schema_fp.0),
|
||||||
|
vec![ObjectRef {
|
||||||
|
bucket: bucket.clone(),
|
||||||
|
key: storage_key.clone(),
|
||||||
|
size_bytes,
|
||||||
|
created_at: now,
|
||||||
|
}],
|
||||||
|
).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||||
|
|
||||||
|
// Rich metadata: PII auto-detect, lineage, row_count, columns.
|
||||||
|
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
|
||||||
|
let sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
|
||||||
|
let columns: Vec<shared::types::ColumnMeta> = schema.fields().iter().map(|f| {
|
||||||
|
let sens = shared::pii::detect_sensitivity(f.name());
|
||||||
|
shared::types::ColumnMeta {
|
||||||
|
name: f.name().clone(),
|
||||||
|
data_type: f.data_type().to_string(),
|
||||||
|
sensitivity: sens.clone(),
|
||||||
|
description: String::new(),
|
||||||
|
is_pii: matches!(sens, Some(shared::types::Sensitivity::Pii)),
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
|
||||||
|
let lineage = shared::types::Lineage {
|
||||||
|
source_system: "postgresql".to_string(),
|
||||||
|
source_file: format!("dsn: {}", redact_dsn(&req.dsn)),
|
||||||
|
ingest_job: format!("pg-stream-{}", now.timestamp_millis()),
|
||||||
|
ingest_timestamp: now,
|
||||||
|
parent_datasets: vec![],
|
||||||
|
};
|
||||||
|
|
||||||
|
let _ = state.registry.update_metadata(&dataset_name, catalogd::registry::MetadataUpdate {
|
||||||
|
sensitivity,
|
||||||
|
columns: Some(columns),
|
||||||
|
lineage: Some(lineage),
|
||||||
|
row_count: Some(stream_result.rows as u64),
|
||||||
|
..Default::default()
|
||||||
|
}).await;
|
||||||
|
|
||||||
|
// Phase C: mark embeddings stale if the dataset already had a vector
|
||||||
|
// index. No-op for newly-created datasets.
|
||||||
|
let _ = state.registry.mark_embeddings_stale(&dataset_name).await;
|
||||||
|
|
||||||
|
// Phase 16.5: notify autotune agent. No-op if no indexes attached.
|
||||||
|
notify_agent_on_append(&state, &dataset_name).await;
|
||||||
|
|
||||||
|
Ok((StatusCode::CREATED, Json(serde_json::json!({
|
||||||
|
"dataset_name": dataset_name,
|
||||||
|
"table": stream_result.table,
|
||||||
|
"rows": stream_result.rows,
|
||||||
|
"batches": stream_result.batches,
|
||||||
|
"columns": stream_result.columns,
|
||||||
|
"schema": stream_result.schema,
|
||||||
|
"storage_key": storage_key,
|
||||||
|
"size_bytes": size_bytes,
|
||||||
|
"duration_secs": stream_result.duration_secs,
|
||||||
|
}))))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Streaming MySQL ingest. Mirrors `ingest_db_stream` — same pagination +
|
||||||
|
/// Parquet-footer schema recovery + PII/lineage metadata. Different driver,
|
||||||
|
/// otherwise identical.
|
||||||
|
async fn ingest_mysql_stream(
|
||||||
|
State(state): State<IngestState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(req): Json<my_stream::MyStreamRequest>,
|
||||||
|
) -> Result<(StatusCode, Json<serde_json::Value>), (StatusCode, String)> {
|
||||||
|
let (bucket, store) = resolve_bucket(&headers, &state.buckets)?;
|
||||||
|
tracing::info!(
|
||||||
|
"mysql stream ingest: table='{}' dataset='{:?}' bucket='{}' batch_size={:?}",
|
||||||
|
req.table, req.dataset_name, bucket, req.batch_size,
|
||||||
|
);
|
||||||
|
|
||||||
|
let (parquet, stream_result) = my_stream::stream_table_to_parquet(&req)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::BAD_GATEWAY, e))?;
|
||||||
|
|
||||||
|
if stream_result.rows == 0 {
|
||||||
|
return Ok((StatusCode::OK, Json(serde_json::json!({
|
||||||
|
"table": req.table,
|
||||||
|
"rows": 0,
|
||||||
|
"message": "table is empty",
|
||||||
|
}))));
|
||||||
|
}
|
||||||
|
|
||||||
|
let (schema, _) = shared::arrow_helpers::parquet_to_record_batches(&parquet)
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("reparse parquet: {e}")))?;
|
||||||
|
|
||||||
|
let dataset_name = req.dataset_name.clone().unwrap_or_else(|| req.table.clone());
|
||||||
|
let storage_key = format!("datasets/{}.parquet", dataset_name);
|
||||||
|
let size_bytes = parquet.len() as u64;
|
||||||
|
|
||||||
|
ops::put(&store, &storage_key, parquet)
|
||||||
|
.await
|
||||||
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||||
|
|
||||||
|
let schema_fp = shared::arrow_helpers::fingerprint_schema(&schema);
|
||||||
|
let now = chrono::Utc::now();
|
||||||
|
state.registry.register(
|
||||||
|
dataset_name.clone(),
|
||||||
|
SchemaFingerprint(schema_fp.0),
|
||||||
|
vec![ObjectRef {
|
||||||
|
bucket: bucket.clone(),
|
||||||
|
key: storage_key.clone(),
|
||||||
|
size_bytes,
|
||||||
|
created_at: now,
|
||||||
|
}],
|
||||||
|
).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||||
|
|
||||||
|
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
|
||||||
|
let sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
|
||||||
|
let columns: Vec<shared::types::ColumnMeta> = schema.fields().iter().map(|f| {
|
||||||
|
let sens = shared::pii::detect_sensitivity(f.name());
|
||||||
|
shared::types::ColumnMeta {
|
||||||
|
name: f.name().clone(),
|
||||||
|
data_type: f.data_type().to_string(),
|
||||||
|
sensitivity: sens.clone(),
|
||||||
|
description: String::new(),
|
||||||
|
is_pii: matches!(sens, Some(shared::types::Sensitivity::Pii)),
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
|
||||||
|
let lineage = shared::types::Lineage {
|
||||||
|
source_system: "mysql".to_string(),
|
||||||
|
source_file: format!("dsn: {}", redact_dsn(&req.dsn)),
|
||||||
|
ingest_job: format!("mysql-stream-{}", now.timestamp_millis()),
|
||||||
|
ingest_timestamp: now,
|
||||||
|
parent_datasets: vec![],
|
||||||
|
};
|
||||||
|
|
||||||
|
let _ = state.registry.update_metadata(&dataset_name, catalogd::registry::MetadataUpdate {
|
||||||
|
sensitivity,
|
||||||
|
columns: Some(columns),
|
||||||
|
lineage: Some(lineage),
|
||||||
|
row_count: Some(stream_result.rows as u64),
|
||||||
|
..Default::default()
|
||||||
|
}).await;
|
||||||
|
|
||||||
|
// Phase C: mark stale + Phase 16.5: notify agent.
|
||||||
|
let _ = state.registry.mark_embeddings_stale(&dataset_name).await;
|
||||||
|
notify_agent_on_append(&state, &dataset_name).await;
|
||||||
|
|
||||||
|
Ok((StatusCode::CREATED, Json(serde_json::json!({
|
||||||
|
"dataset_name": dataset_name,
|
||||||
|
"table": stream_result.table,
|
||||||
|
"rows": stream_result.rows,
|
||||||
|
"batches": stream_result.batches,
|
||||||
|
"columns": stream_result.columns,
|
||||||
|
"schema": stream_result.schema,
|
||||||
|
"storage_key": storage_key,
|
||||||
|
"size_bytes": size_bytes,
|
||||||
|
"duration_secs": stream_result.duration_secs,
|
||||||
|
}))))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Redact the password in any `scheme://user:pass@host/...` DSN. Works
|
||||||
|
/// for postgres, mysql, and any other scheme using the same shape.
|
||||||
|
fn redact_dsn(dsn: &str) -> String {
|
||||||
|
let scheme_end = match dsn.find("://") {
|
||||||
|
Some(i) => i + 3,
|
||||||
|
None => return dsn.to_string(),
|
||||||
|
};
|
||||||
|
let at_idx = match dsn.rfind('@') {
|
||||||
|
Some(i) if i > scheme_end => i,
|
||||||
|
_ => return dsn.to_string(),
|
||||||
|
};
|
||||||
|
// Find the password separator (first `:` after scheme, before `@`).
|
||||||
|
let userpass = &dsn[scheme_end..at_idx];
|
||||||
|
let colon_offset = match userpass.find(':') {
|
||||||
|
Some(i) => i,
|
||||||
|
None => return dsn.to_string(), // no password in DSN
|
||||||
|
};
|
||||||
|
let colon_idx = scheme_end + colon_offset;
|
||||||
|
format!("{}:***{}", &dsn[..colon_idx], &dsn[at_idx..])
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Phase E: Scheduled ingest CRUD endpoints ---
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct CreateScheduleRequest {
|
||||||
|
/// Stable id used in URLs and on disk. Required so callers can
|
||||||
|
/// PATCH/DELETE deterministically. Sanitized on storage but the
|
||||||
|
/// caller's id is what gets persisted as the canonical key.
|
||||||
|
id: String,
|
||||||
|
kind: schedule::ScheduleKind,
|
||||||
|
trigger: schedule::ScheduleTrigger,
|
||||||
|
#[serde(default = "default_enabled")]
|
||||||
|
enabled: bool,
|
||||||
|
#[serde(default)]
|
||||||
|
created_by: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_enabled() -> bool { true }
|
||||||
|
|
||||||
|
async fn create_schedule(
|
||||||
|
State(state): State<IngestState>,
|
||||||
|
Json(req): Json<CreateScheduleRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
if state.schedules.get(&req.id).await.is_some() {
|
||||||
|
return Err((StatusCode::CONFLICT, format!("schedule '{}' already exists", req.id)));
|
||||||
|
}
|
||||||
|
let now = chrono::Utc::now();
|
||||||
|
let def = schedule::ScheduleDef {
|
||||||
|
id: req.id,
|
||||||
|
kind: req.kind,
|
||||||
|
trigger: req.trigger,
|
||||||
|
enabled: req.enabled,
|
||||||
|
created_at: now,
|
||||||
|
created_by: req.created_by,
|
||||||
|
// Fire on the next tick — operators usually want to verify
|
||||||
|
// their schedule works without waiting for the full interval.
|
||||||
|
next_run_at: now,
|
||||||
|
last_run_at: None,
|
||||||
|
last_outcome: None,
|
||||||
|
run_count: 0,
|
||||||
|
failure_count: 0,
|
||||||
|
};
|
||||||
|
match state.schedules.put(def).await {
|
||||||
|
Ok(d) => Ok((StatusCode::CREATED, Json(d))),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_schedules(State(state): State<IngestState>) -> impl IntoResponse {
|
||||||
|
Json(state.schedules.list().await)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_schedule(
|
||||||
|
State(state): State<IngestState>,
|
||||||
|
Path(id): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match state.schedules.get(&id).await {
|
||||||
|
Some(d) => Ok(Json(d)),
|
||||||
|
None => Err((StatusCode::NOT_FOUND, format!("schedule '{id}' not found"))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct PatchScheduleRequest {
|
||||||
|
/// Toggle on/off without changing anything else.
|
||||||
|
#[serde(default)]
|
||||||
|
enabled: Option<bool>,
|
||||||
|
/// Replace the trigger spec — useful for bumping intervals.
|
||||||
|
#[serde(default)]
|
||||||
|
trigger: Option<schedule::ScheduleTrigger>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn patch_schedule(
|
||||||
|
State(state): State<IngestState>,
|
||||||
|
Path(id): Path<String>,
|
||||||
|
Json(req): Json<PatchScheduleRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let Some(mut def) = state.schedules.get(&id).await else {
|
||||||
|
return Err((StatusCode::NOT_FOUND, format!("schedule '{id}' not found")));
|
||||||
|
};
|
||||||
|
if let Some(e) = req.enabled { def.enabled = e; }
|
||||||
|
if let Some(t) = req.trigger { def.trigger = t; }
|
||||||
|
match state.schedules.put(def).await {
|
||||||
|
Ok(d) => Ok(Json(d)),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_schedule(
|
||||||
|
State(state): State<IngestState>,
|
||||||
|
Path(id): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match state.schedules.delete(&id).await {
|
||||||
|
Ok(true) => Ok(Json(serde_json::json!({ "deleted": id }))),
|
||||||
|
Ok(false) => Err((StatusCode::NOT_FOUND, format!("schedule '{id}' not found"))),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Force a schedule to fire right now, regardless of its next_run_at.
|
||||||
|
/// Records the outcome and reschedules. Useful for "test my new schedule
|
||||||
|
/// works without waiting an hour."
|
||||||
|
async fn run_schedule_now(
|
||||||
|
State(state): State<IngestState>,
|
||||||
|
Path(id): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let Some(def) = state.schedules.get(&id).await else {
|
||||||
|
return Err((StatusCode::NOT_FOUND, format!("schedule '{id}' not found")));
|
||||||
|
};
|
||||||
|
let deps = schedule::SchedulerIngestDeps {
|
||||||
|
store: state.store.clone(),
|
||||||
|
registry: state.registry.clone(),
|
||||||
|
buckets: state.buckets.clone(),
|
||||||
|
agent_handle: Some(state.agent_handle.clone()),
|
||||||
|
index_registry: state.index_registry.clone(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let t0 = std::time::Instant::now();
|
||||||
|
let result = schedule::run_schedule(&def, &deps).await;
|
||||||
|
let elapsed = t0.elapsed().as_secs_f32();
|
||||||
|
let outcome = schedule::RunOutcome {
|
||||||
|
at: chrono::Utc::now(),
|
||||||
|
success: result.is_ok(),
|
||||||
|
message: match &result {
|
||||||
|
Ok(s) => s.clone(),
|
||||||
|
Err(e) => e.clone(),
|
||||||
|
},
|
||||||
|
rows: result.as_ref().ok().and_then(|m: &String| {
|
||||||
|
m.split_whitespace()
|
||||||
|
.find_map(|tok| tok.strip_prefix("rows=").and_then(|n| n.parse().ok()))
|
||||||
|
}),
|
||||||
|
duration_secs: elapsed,
|
||||||
|
};
|
||||||
|
if let Err(e) = state.schedules.record_run(&id, outcome.clone()).await {
|
||||||
|
tracing::warn!("schedule {id}: record_run failed: {e}");
|
||||||
|
}
|
||||||
|
Ok(Json(outcome))
|
||||||
|
}
|
||||||
@@ -0,0 +1,122 @@
|
|||||||
|
/// File watcher: monitors a directory for new files and auto-ingests them.
|
||||||
|
/// Drop a CSV, JSON, PDF, or text file → auto-detected → Parquet → queryable.
|
||||||
|
/// Processed files are moved to a "processed/" subdirectory.
|
||||||
|
|
||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::Duration;
|
||||||
|
use tokio::time;
|
||||||
|
|
||||||
|
use catalogd::registry::Registry;
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
|
||||||
|
use crate::pipeline;
|
||||||
|
|
||||||
|
/// Watch config.
|
||||||
|
pub struct WatchConfig {
|
||||||
|
pub watch_dir: PathBuf,
|
||||||
|
pub processed_dir: PathBuf,
|
||||||
|
pub failed_dir: PathBuf,
|
||||||
|
pub poll_interval: Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for WatchConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
watch_dir: PathBuf::from("./inbox"),
|
||||||
|
processed_dir: PathBuf::from("./inbox/processed"),
|
||||||
|
failed_dir: PathBuf::from("./inbox/failed"),
|
||||||
|
poll_interval: Duration::from_secs(10),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Start the file watcher as a background task.
|
||||||
|
pub fn spawn_watcher(
|
||||||
|
config: WatchConfig,
|
||||||
|
store: Arc<dyn ObjectStore>,
|
||||||
|
registry: Registry,
|
||||||
|
) -> tokio::task::JoinHandle<()> {
|
||||||
|
// Create directories
|
||||||
|
std::fs::create_dir_all(&config.watch_dir).ok();
|
||||||
|
std::fs::create_dir_all(&config.processed_dir).ok();
|
||||||
|
std::fs::create_dir_all(&config.failed_dir).ok();
|
||||||
|
|
||||||
|
tracing::info!("file watcher started: watching {}", config.watch_dir.display());
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let mut interval = time::interval(config.poll_interval);
|
||||||
|
loop {
|
||||||
|
interval.tick().await;
|
||||||
|
if let Err(e) = process_inbox(&config, &store, ®istry).await {
|
||||||
|
tracing::error!("watcher error: {e}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Scan inbox directory and process any new files.
|
||||||
|
async fn process_inbox(
|
||||||
|
config: &WatchConfig,
|
||||||
|
store: &Arc<dyn ObjectStore>,
|
||||||
|
registry: &Registry,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let entries = std::fs::read_dir(&config.watch_dir)
|
||||||
|
.map_err(|e| format!("can't read inbox: {e}"))?;
|
||||||
|
|
||||||
|
for entry in entries {
|
||||||
|
let entry = match entry {
|
||||||
|
Ok(e) => e,
|
||||||
|
Err(_) => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
let path = entry.path();
|
||||||
|
|
||||||
|
// Skip directories and hidden files
|
||||||
|
if path.is_dir() || path.file_name().map_or(true, |n| n.to_string_lossy().starts_with('.')) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let filename = path.file_name().unwrap().to_string_lossy().to_string();
|
||||||
|
tracing::info!("watcher: found new file '{}'", filename);
|
||||||
|
|
||||||
|
// Read file
|
||||||
|
let content = match std::fs::read(&path) {
|
||||||
|
Ok(c) => c,
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!("watcher: can't read '{}': {e}", filename);
|
||||||
|
move_file(&path, &config.failed_dir.join(&filename));
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Ingest
|
||||||
|
match pipeline::ingest_file(&filename, &content, None, store, registry).await {
|
||||||
|
Ok(result) => {
|
||||||
|
if result.deduplicated {
|
||||||
|
tracing::info!("watcher: '{}' already ingested, skipping", filename);
|
||||||
|
} else {
|
||||||
|
tracing::info!("watcher: ingested '{}' → {} ({} rows, {} cols)",
|
||||||
|
filename, result.dataset_name, result.rows, result.columns);
|
||||||
|
}
|
||||||
|
// Move to processed
|
||||||
|
move_file(&path, &config.processed_dir.join(&filename));
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!("watcher: failed to ingest '{}': {e}", filename);
|
||||||
|
move_file(&path, &config.failed_dir.join(&filename));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn move_file(from: &Path, to: &Path) {
|
||||||
|
if let Err(e) = std::fs::rename(from, to) {
|
||||||
|
// rename fails across filesystems, try copy+delete
|
||||||
|
if let Err(e2) = std::fs::copy(from, to).and_then(|_| std::fs::remove_file(from)) {
|
||||||
|
tracing::warn!("can't move {} → {}: {e}, copy failed: {e2}", from.display(), to.display());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,18 @@
|
|||||||
|
[package]
|
||||||
|
name = "journald"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2024"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
shared = { path = "../shared" }
|
||||||
|
storaged = { path = "../storaged" }
|
||||||
|
tokio = { workspace = true }
|
||||||
|
axum = { workspace = true }
|
||||||
|
serde = { workspace = true }
|
||||||
|
serde_json = { workspace = true }
|
||||||
|
tracing = { workspace = true }
|
||||||
|
arrow = { workspace = true }
|
||||||
|
parquet = { workspace = true }
|
||||||
|
bytes = { workspace = true }
|
||||||
|
chrono = { workspace = true }
|
||||||
|
object_store = { workspace = true }
|
||||||
@@ -0,0 +1,322 @@
|
|||||||
|
/// Append-only event journal.
|
||||||
|
/// Every mutation to any dataset is recorded here as an immutable event.
|
||||||
|
/// Events are never modified or deleted — this IS the audit trail.
|
||||||
|
///
|
||||||
|
/// Storage: events buffer in memory, flush to Parquet periodically.
|
||||||
|
/// Query: load Parquet files, filter by entity/field/actor/time.
|
||||||
|
|
||||||
|
use arrow::array::{ArrayRef, RecordBatch, StringArray, UInt64Array};
|
||||||
|
use arrow::datatypes::{DataType, Field, Schema};
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
use shared::arrow_helpers::record_batch_to_parquet;
|
||||||
|
use storaged::ops;
|
||||||
|
|
||||||
|
/// A single mutation event. Immutable once created.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct Event {
|
||||||
|
pub event_id: String,
|
||||||
|
pub timestamp: DateTime<Utc>,
|
||||||
|
pub entity_type: String, // "candidate", "placement", "client", etc.
|
||||||
|
pub entity_id: String, // "CAND-00342"
|
||||||
|
pub field: String, // "phone", "*" for insert/delete
|
||||||
|
pub action: String, // "insert", "update", "delete"
|
||||||
|
pub old_value: String, // "" for inserts
|
||||||
|
pub new_value: String, // "" for deletes
|
||||||
|
pub actor: String, // "Sarah", "ingest_pipeline", "auto_matcher"
|
||||||
|
pub source: String, // "api", "ingest", "workspace", "agent"
|
||||||
|
pub workspace_id: String, // "" if not workspace-scoped
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Arrow schema for journal Parquet files.
|
||||||
|
fn journal_schema() -> Arc<Schema> {
|
||||||
|
Arc::new(Schema::new(vec![
|
||||||
|
Field::new("event_id", DataType::Utf8, false),
|
||||||
|
Field::new("timestamp", DataType::Utf8, false),
|
||||||
|
Field::new("entity_type", DataType::Utf8, false),
|
||||||
|
Field::new("entity_id", DataType::Utf8, false),
|
||||||
|
Field::new("field", DataType::Utf8, false),
|
||||||
|
Field::new("action", DataType::Utf8, false),
|
||||||
|
Field::new("old_value", DataType::Utf8, true),
|
||||||
|
Field::new("new_value", DataType::Utf8, true),
|
||||||
|
Field::new("actor", DataType::Utf8, false),
|
||||||
|
Field::new("source", DataType::Utf8, false),
|
||||||
|
Field::new("workspace_id", DataType::Utf8, true),
|
||||||
|
]))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The event journal — buffers events in memory, flushes to Parquet.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct Journal {
|
||||||
|
buffer: Arc<RwLock<Vec<Event>>>,
|
||||||
|
store: Arc<dyn ObjectStore>,
|
||||||
|
flush_threshold: usize,
|
||||||
|
event_counter: Arc<std::sync::atomic::AtomicU64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Journal {
|
||||||
|
pub fn new(store: Arc<dyn ObjectStore>, flush_threshold: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
buffer: Arc::new(RwLock::new(Vec::new())),
|
||||||
|
store,
|
||||||
|
flush_threshold,
|
||||||
|
event_counter: Arc::new(std::sync::atomic::AtomicU64::new(0)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Record a single event. Triggers flush if buffer exceeds threshold.
|
||||||
|
pub async fn record(&self, mut event: Event) -> Result<(), String> {
|
||||||
|
let seq = self.event_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||||
|
event.event_id = format!("evt-{}-{seq}", event.timestamp.timestamp_millis());
|
||||||
|
|
||||||
|
let should_flush = {
|
||||||
|
let mut buf = self.buffer.write().await;
|
||||||
|
buf.push(event);
|
||||||
|
buf.len() >= self.flush_threshold
|
||||||
|
};
|
||||||
|
|
||||||
|
if should_flush {
|
||||||
|
self.flush().await?;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Record a batch insert (multiple rows added to a dataset).
|
||||||
|
pub async fn record_insert(
|
||||||
|
&self,
|
||||||
|
entity_type: &str,
|
||||||
|
entity_id: &str,
|
||||||
|
actor: &str,
|
||||||
|
source: &str,
|
||||||
|
workspace_id: &str,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
self.record(Event {
|
||||||
|
event_id: String::new(),
|
||||||
|
timestamp: Utc::now(),
|
||||||
|
entity_type: entity_type.to_string(),
|
||||||
|
entity_id: entity_id.to_string(),
|
||||||
|
field: "*".to_string(),
|
||||||
|
action: "insert".to_string(),
|
||||||
|
old_value: String::new(),
|
||||||
|
new_value: String::new(),
|
||||||
|
actor: actor.to_string(),
|
||||||
|
source: source.to_string(),
|
||||||
|
workspace_id: workspace_id.to_string(),
|
||||||
|
}).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Record a field update.
|
||||||
|
pub async fn record_update(
|
||||||
|
&self,
|
||||||
|
entity_type: &str,
|
||||||
|
entity_id: &str,
|
||||||
|
field: &str,
|
||||||
|
old_value: &str,
|
||||||
|
new_value: &str,
|
||||||
|
actor: &str,
|
||||||
|
source: &str,
|
||||||
|
workspace_id: &str,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
self.record(Event {
|
||||||
|
event_id: String::new(),
|
||||||
|
timestamp: Utc::now(),
|
||||||
|
entity_type: entity_type.to_string(),
|
||||||
|
entity_id: entity_id.to_string(),
|
||||||
|
field: field.to_string(),
|
||||||
|
action: "update".to_string(),
|
||||||
|
old_value: old_value.to_string(),
|
||||||
|
new_value: new_value.to_string(),
|
||||||
|
actor: actor.to_string(),
|
||||||
|
source: source.to_string(),
|
||||||
|
workspace_id: workspace_id.to_string(),
|
||||||
|
}).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Record a batch of events from a dataset ingest.
|
||||||
|
pub async fn record_ingest(
|
||||||
|
&self,
|
||||||
|
entity_type: &str,
|
||||||
|
count: usize,
|
||||||
|
actor: &str,
|
||||||
|
source_file: &str,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
self.record(Event {
|
||||||
|
event_id: String::new(),
|
||||||
|
timestamp: Utc::now(),
|
||||||
|
entity_type: entity_type.to_string(),
|
||||||
|
entity_id: format!("batch:{count}"),
|
||||||
|
field: "*".to_string(),
|
||||||
|
action: "ingest".to_string(),
|
||||||
|
old_value: String::new(),
|
||||||
|
new_value: format!("{count} rows from {source_file}"),
|
||||||
|
actor: actor.to_string(),
|
||||||
|
source: "ingest".to_string(),
|
||||||
|
workspace_id: String::new(),
|
||||||
|
}).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Flush buffer to Parquet. Called automatically at threshold or manually.
|
||||||
|
pub async fn flush(&self) -> Result<usize, String> {
|
||||||
|
let events = {
|
||||||
|
let mut buf = self.buffer.write().await;
|
||||||
|
if buf.is_empty() { return Ok(0); }
|
||||||
|
std::mem::take(&mut *buf)
|
||||||
|
};
|
||||||
|
|
||||||
|
let n = events.len();
|
||||||
|
let batch = events_to_batch(&events)?;
|
||||||
|
let parquet = record_batch_to_parquet(&batch)?;
|
||||||
|
|
||||||
|
let ts = Utc::now().format("%Y%m%d_%H%M%S_%3f");
|
||||||
|
let key = format!("journal/{ts}.parquet");
|
||||||
|
ops::put(&self.store, &key, parquet).await?;
|
||||||
|
|
||||||
|
tracing::info!("journal: flushed {n} events to {key}");
|
||||||
|
Ok(n)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Query journal: get all events for a specific entity.
|
||||||
|
pub async fn get_entity_history(&self, entity_id: &str) -> Result<Vec<Event>, String> {
|
||||||
|
let mut all_events = Vec::new();
|
||||||
|
|
||||||
|
// Check in-memory buffer first
|
||||||
|
{
|
||||||
|
let buf = self.buffer.read().await;
|
||||||
|
for e in buf.iter() {
|
||||||
|
if e.entity_id == entity_id {
|
||||||
|
all_events.push(e.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Load from persisted journal files
|
||||||
|
let keys = ops::list(&self.store, Some("journal/")).await?;
|
||||||
|
for key in &keys {
|
||||||
|
if !key.ends_with(".parquet") { continue; }
|
||||||
|
let data = ops::get(&self.store, key).await?;
|
||||||
|
let (_, batches) = shared::arrow_helpers::parquet_to_record_batches(&data)?;
|
||||||
|
for batch in &batches {
|
||||||
|
let events = batch_to_events(batch)?;
|
||||||
|
for e in events {
|
||||||
|
if e.entity_id == entity_id {
|
||||||
|
all_events.push(e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
all_events.sort_by_key(|e| e.timestamp);
|
||||||
|
Ok(all_events)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Query journal: get recent events (last N).
|
||||||
|
pub async fn get_recent(&self, limit: usize) -> Result<Vec<Event>, String> {
|
||||||
|
let mut all_events = Vec::new();
|
||||||
|
|
||||||
|
// Buffer
|
||||||
|
{
|
||||||
|
let buf = self.buffer.read().await;
|
||||||
|
all_events.extend(buf.iter().cloned());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Persisted (load most recent files first)
|
||||||
|
let mut keys = ops::list(&self.store, Some("journal/")).await?;
|
||||||
|
keys.sort();
|
||||||
|
keys.reverse();
|
||||||
|
|
||||||
|
for key in &keys {
|
||||||
|
if !key.ends_with(".parquet") { continue; }
|
||||||
|
if all_events.len() >= limit { break; }
|
||||||
|
let data = ops::get(&self.store, key).await?;
|
||||||
|
let (_, batches) = shared::arrow_helpers::parquet_to_record_batches(&data)?;
|
||||||
|
for batch in &batches {
|
||||||
|
all_events.extend(batch_to_events(batch)?);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
all_events.sort_by_key(|e| e.timestamp);
|
||||||
|
all_events.reverse();
|
||||||
|
all_events.truncate(limit);
|
||||||
|
Ok(all_events)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get buffer stats.
|
||||||
|
pub async fn stats(&self) -> JournalStats {
|
||||||
|
let buf = self.buffer.read().await;
|
||||||
|
let keys = ops::list(&self.store, Some("journal/")).await.unwrap_or_default();
|
||||||
|
JournalStats {
|
||||||
|
buffer_events: buf.len(),
|
||||||
|
flush_threshold: self.flush_threshold,
|
||||||
|
persisted_files: keys.len(),
|
||||||
|
total_events_created: self.event_counter.load(std::sync::atomic::Ordering::Relaxed),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct JournalStats {
|
||||||
|
pub buffer_events: usize,
|
||||||
|
pub flush_threshold: usize,
|
||||||
|
pub persisted_files: usize,
|
||||||
|
pub total_events_created: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn events_to_batch(events: &[Event]) -> Result<RecordBatch, String> {
|
||||||
|
let schema = journal_schema();
|
||||||
|
let arrays: Vec<ArrayRef> = vec![
|
||||||
|
Arc::new(StringArray::from(events.iter().map(|e| e.event_id.as_str()).collect::<Vec<_>>())),
|
||||||
|
Arc::new(StringArray::from(events.iter().map(|e| e.timestamp.to_rfc3339()).collect::<Vec<_>>())),
|
||||||
|
Arc::new(StringArray::from(events.iter().map(|e| e.entity_type.as_str()).collect::<Vec<_>>())),
|
||||||
|
Arc::new(StringArray::from(events.iter().map(|e| e.entity_id.as_str()).collect::<Vec<_>>())),
|
||||||
|
Arc::new(StringArray::from(events.iter().map(|e| e.field.as_str()).collect::<Vec<_>>())),
|
||||||
|
Arc::new(StringArray::from(events.iter().map(|e| e.action.as_str()).collect::<Vec<_>>())),
|
||||||
|
Arc::new(StringArray::from(events.iter().map(|e| e.old_value.as_str()).collect::<Vec<_>>())),
|
||||||
|
Arc::new(StringArray::from(events.iter().map(|e| e.new_value.as_str()).collect::<Vec<_>>())),
|
||||||
|
Arc::new(StringArray::from(events.iter().map(|e| e.actor.as_str()).collect::<Vec<_>>())),
|
||||||
|
Arc::new(StringArray::from(events.iter().map(|e| e.source.as_str()).collect::<Vec<_>>())),
|
||||||
|
Arc::new(StringArray::from(events.iter().map(|e| e.workspace_id.as_str()).collect::<Vec<_>>())),
|
||||||
|
];
|
||||||
|
RecordBatch::try_new(schema, arrays).map_err(|e| format!("batch error: {e}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn batch_to_events(batch: &RecordBatch) -> Result<Vec<Event>, String> {
|
||||||
|
let get_col = |i: usize| -> Result<&StringArray, String> {
|
||||||
|
batch.column(i).as_any().downcast_ref::<StringArray>()
|
||||||
|
.ok_or_else(|| format!("column {i} not string"))
|
||||||
|
};
|
||||||
|
|
||||||
|
let event_ids = get_col(0)?;
|
||||||
|
let timestamps = get_col(1)?;
|
||||||
|
let entity_types = get_col(2)?;
|
||||||
|
let entity_ids = get_col(3)?;
|
||||||
|
let fields = get_col(4)?;
|
||||||
|
let actions = get_col(5)?;
|
||||||
|
let old_values = get_col(6)?;
|
||||||
|
let new_values = get_col(7)?;
|
||||||
|
let actors = get_col(8)?;
|
||||||
|
let sources = get_col(9)?;
|
||||||
|
let workspace_ids = get_col(10)?;
|
||||||
|
|
||||||
|
let mut events = Vec::with_capacity(batch.num_rows());
|
||||||
|
for i in 0..batch.num_rows() {
|
||||||
|
events.push(Event {
|
||||||
|
event_id: event_ids.value(i).to_string(),
|
||||||
|
timestamp: timestamps.value(i).parse().unwrap_or_else(|_| Utc::now()),
|
||||||
|
entity_type: entity_types.value(i).to_string(),
|
||||||
|
entity_id: entity_ids.value(i).to_string(),
|
||||||
|
field: fields.value(i).to_string(),
|
||||||
|
action: actions.value(i).to_string(),
|
||||||
|
old_value: old_values.value(i).to_string(),
|
||||||
|
new_value: new_values.value(i).to_string(),
|
||||||
|
actor: actors.value(i).to_string(),
|
||||||
|
source: sources.value(i).to_string(),
|
||||||
|
workspace_id: workspace_ids.value(i).to_string(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Ok(events)
|
||||||
|
}
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
pub mod journal;
|
||||||
|
pub mod service;
|
||||||
@@ -0,0 +1,131 @@
|
|||||||
|
use axum::{
|
||||||
|
Json, Router,
|
||||||
|
extract::{Path, Query, State},
|
||||||
|
http::StatusCode,
|
||||||
|
response::IntoResponse,
|
||||||
|
routing::{get, post},
|
||||||
|
};
|
||||||
|
use serde::Deserialize;
|
||||||
|
|
||||||
|
use crate::journal::{Event, Journal};
|
||||||
|
|
||||||
|
pub fn router(journal: Journal) -> Router {
|
||||||
|
Router::new()
|
||||||
|
.route("/health", get(health))
|
||||||
|
.route("/event", post(record_event))
|
||||||
|
.route("/update", post(record_update))
|
||||||
|
.route("/flush", post(flush))
|
||||||
|
.route("/history/{entity_id}", get(entity_history))
|
||||||
|
.route("/recent", get(recent_events))
|
||||||
|
.route("/stats", get(stats))
|
||||||
|
.with_state(journal)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn health() -> &'static str {
|
||||||
|
"journald ok"
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Record events ---
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct RecordEventRequest {
|
||||||
|
entity_type: String,
|
||||||
|
entity_id: String,
|
||||||
|
action: String,
|
||||||
|
field: Option<String>,
|
||||||
|
old_value: Option<String>,
|
||||||
|
new_value: Option<String>,
|
||||||
|
actor: String,
|
||||||
|
source: String,
|
||||||
|
workspace_id: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_event(
|
||||||
|
State(journal): State<Journal>,
|
||||||
|
Json(req): Json<RecordEventRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let event = Event {
|
||||||
|
event_id: String::new(), // assigned by journal
|
||||||
|
timestamp: chrono::Utc::now(),
|
||||||
|
entity_type: req.entity_type,
|
||||||
|
entity_id: req.entity_id,
|
||||||
|
field: req.field.unwrap_or_else(|| "*".to_string()),
|
||||||
|
action: req.action,
|
||||||
|
old_value: req.old_value.unwrap_or_default(),
|
||||||
|
new_value: req.new_value.unwrap_or_default(),
|
||||||
|
actor: req.actor,
|
||||||
|
source: req.source,
|
||||||
|
workspace_id: req.workspace_id.unwrap_or_default(),
|
||||||
|
};
|
||||||
|
|
||||||
|
match journal.record(event).await {
|
||||||
|
Ok(()) => Ok((StatusCode::CREATED, "event recorded")),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct UpdateRequest {
|
||||||
|
entity_type: String,
|
||||||
|
entity_id: String,
|
||||||
|
field: String,
|
||||||
|
old_value: String,
|
||||||
|
new_value: String,
|
||||||
|
actor: String,
|
||||||
|
workspace_id: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_update(
|
||||||
|
State(journal): State<Journal>,
|
||||||
|
Json(req): Json<UpdateRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match journal.record_update(
|
||||||
|
&req.entity_type, &req.entity_id, &req.field,
|
||||||
|
&req.old_value, &req.new_value, &req.actor, "api",
|
||||||
|
&req.workspace_id.unwrap_or_default(),
|
||||||
|
).await {
|
||||||
|
Ok(()) => Ok((StatusCode::CREATED, "update recorded")),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Flush ---
|
||||||
|
|
||||||
|
async fn flush(State(journal): State<Journal>) -> impl IntoResponse {
|
||||||
|
match journal.flush().await {
|
||||||
|
Ok(n) => Ok(format!("flushed {n} events")),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Query ---
|
||||||
|
|
||||||
|
async fn entity_history(
|
||||||
|
State(journal): State<Journal>,
|
||||||
|
Path(entity_id): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match journal.get_entity_history(&entity_id).await {
|
||||||
|
Ok(events) => Ok(Json(events)),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct RecentQuery {
|
||||||
|
limit: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn recent_events(
|
||||||
|
State(journal): State<Journal>,
|
||||||
|
Query(q): Query<RecentQuery>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let limit = q.limit.unwrap_or(50);
|
||||||
|
match journal.get_recent(limit).await {
|
||||||
|
Ok(events) => Ok(Json(events)),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn stats(State(journal): State<Journal>) -> impl IntoResponse {
|
||||||
|
Json(journal.stats().await)
|
||||||
|
}
|
||||||
@@ -0,0 +1,42 @@
|
|||||||
|
[package]
|
||||||
|
name = "lance-bench"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2024"
|
||||||
|
|
||||||
|
# Standalone pilot for Phase B (see docs/EXECUTION_PLAN.md).
|
||||||
|
# Deliberately NOT sharing workspace deps — Lance 4.x pulls in its own
|
||||||
|
# DataFusion and Arrow versions incompatible with the rest of the stack.
|
||||||
|
# Isolating the pilot means we don't force a workspace-wide upgrade until
|
||||||
|
# we've decided Lance is worth it.
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
# Only the features we actually need — the default brings in AWS/Azure/GCP/HF etc
|
||||||
|
# which is ~200 extra crates we don't care about for a local pilot.
|
||||||
|
lance = { version = "4.0", default-features = false }
|
||||||
|
# Lance exposes DatasetIndexExt, IndexType, and IvfBuildParams through
|
||||||
|
# its sub-crates which must be imported directly — lance itself doesn't
|
||||||
|
# re-export them at a convenient path.
|
||||||
|
lance-index = { version = "4.0", default-features = false }
|
||||||
|
lance-linalg = { version = "4.0", default-features = false }
|
||||||
|
|
||||||
|
# Arrow re-exported by Lance; pin to a range Lance picks so types match.
|
||||||
|
arrow = "57"
|
||||||
|
arrow-array = "57"
|
||||||
|
arrow-schema = "57"
|
||||||
|
|
||||||
|
# Also need to read the EXISTING Parquet vector files so we can compare.
|
||||||
|
# These live in data/vectors/*.parquet. Lance's internal Parquet reading
|
||||||
|
# might differ from ours; using our format's Arrow/Parquet versions for
|
||||||
|
# the read side keeps the inputs identical.
|
||||||
|
parquet = "57"
|
||||||
|
|
||||||
|
tokio = { version = "1", features = ["full"] }
|
||||||
|
futures = "0.3"
|
||||||
|
serde = { version = "1", features = ["derive"] }
|
||||||
|
serde_json = "1"
|
||||||
|
anyhow = "1"
|
||||||
|
bytes = "1"
|
||||||
|
|
||||||
|
[[bin]]
|
||||||
|
name = "lance-bench"
|
||||||
|
path = "src/main.rs"
|
||||||
@@ -0,0 +1,633 @@
|
|||||||
|
//! Phase B: Lance pilot benchmark.
|
||||||
|
//!
|
||||||
|
//! Standalone binary that compares Lance vector storage against our
|
||||||
|
//! Parquet-with-binary-blob + in-RAM HNSW approach. See
|
||||||
|
//! docs/EXECUTION_PLAN.md for the decision rules this fuels.
|
||||||
|
//!
|
||||||
|
//! Inputs:
|
||||||
|
//! data/vectors/resumes_100k_v2.parquet — existing 100K × 768d embeddings
|
||||||
|
//!
|
||||||
|
//! Output:
|
||||||
|
//! A JSON report printed to stdout with measurements for:
|
||||||
|
//! - Cold load time (parquet → arrow) vs Lance open + scan
|
||||||
|
//! - Disk size
|
||||||
|
//! - Vector search latency (p50 / p95 / p99)
|
||||||
|
//! - Single-row random access
|
||||||
|
//! - Append cost (adding 10K rows)
|
||||||
|
//!
|
||||||
|
//! Usage:
|
||||||
|
//! cargo run --bin lance-bench -- \
|
||||||
|
//! --parquet data/vectors/resumes_100k_v2.parquet \
|
||||||
|
//! --lance-out /tmp/lance_resumes_100k_v2 \
|
||||||
|
//! --json-out /tmp/lance_bench.json
|
||||||
|
|
||||||
|
use anyhow::{Context, Result};
|
||||||
|
use arrow_array::{Array, ArrayRef, BinaryArray, FixedSizeListArray, Float32Array, RecordBatch, RecordBatchIterator};
|
||||||
|
use arrow_schema::{DataType, Field, Schema};
|
||||||
|
use serde::Serialize;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
#[derive(Debug, Serialize)]
|
||||||
|
struct BenchReport {
|
||||||
|
vectors: usize,
|
||||||
|
dimensions: usize,
|
||||||
|
parquet_path: String,
|
||||||
|
lance_path: String,
|
||||||
|
|
||||||
|
// Parquet baseline
|
||||||
|
parquet_disk_bytes: u64,
|
||||||
|
parquet_cold_load_secs: f32,
|
||||||
|
|
||||||
|
// Lance numbers
|
||||||
|
lance_write_secs: f32,
|
||||||
|
lance_disk_bytes: u64,
|
||||||
|
lance_cold_open_secs: f32,
|
||||||
|
|
||||||
|
// Index + search
|
||||||
|
lance_index_build_secs: Option<f32>,
|
||||||
|
lance_index_disk_bytes: Option<u64>,
|
||||||
|
lance_search_p50_us: Option<f32>,
|
||||||
|
lance_search_p95_us: Option<f32>,
|
||||||
|
lance_search_p99_us: Option<f32>,
|
||||||
|
|
||||||
|
// Architectural features Parquet+sidecar can't cheaply do
|
||||||
|
lance_random_row_access_us: Option<f32>, // fetch one row by row_id
|
||||||
|
parquet_random_row_access_us: Option<f32>, // for comparison — full scan cost
|
||||||
|
lance_append_10k_secs: Option<f32>, // add 10K new rows
|
||||||
|
lance_append_disk_bytes_added: Option<u64>,
|
||||||
|
|
||||||
|
// Head-to-head reference (from our own measurements)
|
||||||
|
reference_hnsw_p50_us: f32,
|
||||||
|
reference_hnsw_p95_us: f32,
|
||||||
|
reference_brute_force_us: f32,
|
||||||
|
reference_hnsw_build_secs: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() -> Result<()> {
|
||||||
|
// Simple positional args: parquet_in, lance_out.
|
||||||
|
let args: Vec<String> = std::env::args().collect();
|
||||||
|
let parquet_path = args
|
||||||
|
.get(1)
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or_else(|| "data/vectors/resumes_100k_v2.parquet".to_string());
|
||||||
|
let lance_path = args
|
||||||
|
.get(2)
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or_else(|| "/tmp/lance_bench_dataset".to_string());
|
||||||
|
|
||||||
|
eprintln!("=== Phase B Lance pilot ===");
|
||||||
|
eprintln!("input parquet: {}", parquet_path);
|
||||||
|
eprintln!("output lance: {}", lance_path);
|
||||||
|
|
||||||
|
// --- 1. Cold-load the existing Parquet vector index into memory
|
||||||
|
eprintln!("\n[1/4] reading Parquet baseline...");
|
||||||
|
let t0 = Instant::now();
|
||||||
|
let (schema, batches, total_rows) = read_parquet_vectors(&parquet_path)
|
||||||
|
.context("read parquet")?;
|
||||||
|
let parquet_cold_load_secs = t0.elapsed().as_secs_f32();
|
||||||
|
let parquet_disk_bytes = std::fs::metadata(&parquet_path)?.len();
|
||||||
|
|
||||||
|
let dims = detect_vector_dims(&batches)?;
|
||||||
|
eprintln!(
|
||||||
|
" loaded {} rows, {} columns, vectors={}d, disk={:.1} MB, cold load={:.2}s",
|
||||||
|
total_rows,
|
||||||
|
schema.fields().len(),
|
||||||
|
dims,
|
||||||
|
parquet_disk_bytes as f64 / 1_000_000.0,
|
||||||
|
parquet_cold_load_secs,
|
||||||
|
);
|
||||||
|
|
||||||
|
// --- 2. Convert from binary-blob-of-f32 to Lance's FixedSizeList<Float32>
|
||||||
|
eprintln!("\n[2/4] converting binary-blob vectors to Arrow FixedSizeList...");
|
||||||
|
let t0 = Instant::now();
|
||||||
|
let (lance_schema, lance_batches) = convert_to_fixed_size_list(&schema, batches, dims)?;
|
||||||
|
eprintln!(" conversion took {:.2}s", t0.elapsed().as_secs_f32());
|
||||||
|
|
||||||
|
// --- 3. Write as Lance dataset
|
||||||
|
eprintln!("\n[3/4] writing Lance dataset...");
|
||||||
|
let t0 = Instant::now();
|
||||||
|
// Clean up any prior run
|
||||||
|
let _ = std::fs::remove_dir_all(&lance_path);
|
||||||
|
write_lance_dataset(&lance_path, lance_schema.clone(), lance_batches).await?;
|
||||||
|
let lance_write_secs = t0.elapsed().as_secs_f32();
|
||||||
|
let lance_disk_bytes = dir_size_bytes(&lance_path);
|
||||||
|
eprintln!(
|
||||||
|
" write took {:.2}s, disk={:.1} MB",
|
||||||
|
lance_write_secs,
|
||||||
|
lance_disk_bytes as f64 / 1_000_000.0,
|
||||||
|
);
|
||||||
|
|
||||||
|
// --- 4. Cold open + scan the Lance dataset
|
||||||
|
eprintln!("\n[4/6] cold-opening Lance dataset...");
|
||||||
|
let t0 = Instant::now();
|
||||||
|
let scanned_rows = cold_open_and_scan_lance(&lance_path).await?;
|
||||||
|
let lance_cold_open_secs = t0.elapsed().as_secs_f32();
|
||||||
|
eprintln!(
|
||||||
|
" open + full scan: {} rows in {:.2}s",
|
||||||
|
scanned_rows, lance_cold_open_secs,
|
||||||
|
);
|
||||||
|
|
||||||
|
// --- 5. Build a vector index on the Lance dataset
|
||||||
|
eprintln!("\n[5/6] building Lance vector index (IVF_PQ)...");
|
||||||
|
let t0 = Instant::now();
|
||||||
|
let index_built = build_lance_vector_index(&lance_path, dims).await;
|
||||||
|
let (lance_index_build_secs, lance_index_disk_bytes) = match index_built {
|
||||||
|
Ok(()) => {
|
||||||
|
let secs = t0.elapsed().as_secs_f32();
|
||||||
|
let disk = dir_size_bytes(&lance_path) - lance_disk_bytes;
|
||||||
|
eprintln!(" built in {:.2}s, index adds {:.1} MB on disk", secs, disk as f64 / 1e6);
|
||||||
|
(Some(secs), Some(disk))
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
eprintln!(" index build failed: {e:#}");
|
||||||
|
(None, None)
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// --- 6. Run search queries, measure latency
|
||||||
|
eprintln!("\n[6/6] running vector search benchmarks...");
|
||||||
|
let search_stats = if lance_index_build_secs.is_some() {
|
||||||
|
run_search_benchmarks(&lance_path, dims).await.ok()
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
let (lance_search_p50, lance_search_p95, lance_search_p99) = match search_stats {
|
||||||
|
Some((p50, p95, p99)) => {
|
||||||
|
eprintln!(" p50={:.0}us p95={:.0}us p99={:.0}us", p50, p95, p99);
|
||||||
|
(Some(p50), Some(p95), Some(p99))
|
||||||
|
}
|
||||||
|
None => (None, None, None),
|
||||||
|
};
|
||||||
|
|
||||||
|
// --- Random access comparison
|
||||||
|
eprintln!("\n[7/8] random row access — Lance vs full-scan Parquet...");
|
||||||
|
let lance_random = measure_random_access_lance(&lance_path).await.ok();
|
||||||
|
let parquet_random = measure_random_access_parquet(&parquet_path).ok();
|
||||||
|
if let Some(us) = lance_random {
|
||||||
|
eprintln!(" Lance random-fetch avg: {:.0}us", us);
|
||||||
|
}
|
||||||
|
if let Some(us) = parquet_random {
|
||||||
|
eprintln!(" Parquet full-scan-to-row avg: {:.0}us", us);
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Append cost
|
||||||
|
eprintln!("\n[8/8] append 10K new rows to existing dataset...");
|
||||||
|
let t0 = Instant::now();
|
||||||
|
let pre_append_bytes = dir_size_bytes(&lance_path);
|
||||||
|
let append_result = append_10k_rows(&lance_path, dims).await;
|
||||||
|
let (lance_append_secs, lance_append_bytes) = match append_result {
|
||||||
|
Ok(()) => {
|
||||||
|
let secs = t0.elapsed().as_secs_f32();
|
||||||
|
let bytes = dir_size_bytes(&lance_path).saturating_sub(pre_append_bytes);
|
||||||
|
eprintln!(" append took {:.2}s, added {:.1} MB", secs, bytes as f64 / 1e6);
|
||||||
|
(Some(secs), Some(bytes))
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
eprintln!(" append failed: {e:#}");
|
||||||
|
(None, None)
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// --- Report
|
||||||
|
let report = BenchReport {
|
||||||
|
vectors: total_rows,
|
||||||
|
dimensions: dims,
|
||||||
|
parquet_path: parquet_path.clone(),
|
||||||
|
lance_path: lance_path.clone(),
|
||||||
|
parquet_disk_bytes,
|
||||||
|
parquet_cold_load_secs,
|
||||||
|
lance_write_secs,
|
||||||
|
lance_disk_bytes,
|
||||||
|
lance_cold_open_secs,
|
||||||
|
lance_index_build_secs,
|
||||||
|
lance_index_disk_bytes,
|
||||||
|
lance_search_p50_us: lance_search_p50,
|
||||||
|
lance_search_p95_us: lance_search_p95,
|
||||||
|
lance_search_p99_us: lance_search_p99,
|
||||||
|
lance_random_row_access_us: lance_random,
|
||||||
|
parquet_random_row_access_us: parquet_random,
|
||||||
|
lance_append_10k_secs: lance_append_secs,
|
||||||
|
lance_append_disk_bytes_added: lance_append_bytes,
|
||||||
|
// From our Phase 15 trial on the SAME index (ec=80 es=30, recall=1.00):
|
||||||
|
reference_hnsw_p50_us: 873.0,
|
||||||
|
reference_hnsw_p95_us: 1413.0,
|
||||||
|
reference_brute_force_us: 43983.0,
|
||||||
|
reference_hnsw_build_secs: 230.0,
|
||||||
|
};
|
||||||
|
|
||||||
|
let json = serde_json::to_string_pretty(&report)?;
|
||||||
|
println!("{}", json);
|
||||||
|
|
||||||
|
eprintln!("\n=== Summary ===");
|
||||||
|
eprintln!(" Parquet cold load: {:.2}s", report.parquet_cold_load_secs);
|
||||||
|
eprintln!(" Lance cold open: {:.2}s ({})",
|
||||||
|
report.lance_cold_open_secs,
|
||||||
|
format_ratio(report.parquet_cold_load_secs, report.lance_cold_open_secs));
|
||||||
|
eprintln!(" Parquet disk: {:.1} MB", report.parquet_disk_bytes as f64 / 1e6);
|
||||||
|
eprintln!(" Lance disk: {:.1} MB ({})",
|
||||||
|
report.lance_disk_bytes as f64 / 1e6,
|
||||||
|
format_ratio(report.parquet_disk_bytes as f32, report.lance_disk_bytes as f32));
|
||||||
|
if let (Some(p50), Some(p95)) = (report.lance_search_p50_us, report.lance_search_p95_us) {
|
||||||
|
eprintln!(" Lance search p50: {:.0}us vs our HNSW {:.0}us ({})",
|
||||||
|
p50, report.reference_hnsw_p50_us,
|
||||||
|
format_ratio(report.reference_hnsw_p50_us, p50));
|
||||||
|
eprintln!(" Lance search p95: {:.0}us vs our HNSW {:.0}us ({})",
|
||||||
|
p95, report.reference_hnsw_p95_us,
|
||||||
|
format_ratio(report.reference_hnsw_p95_us, p95));
|
||||||
|
eprintln!(" Speedup vs brute force: {:.1}× (Lance) vs {:.1}× (HNSW)",
|
||||||
|
report.reference_brute_force_us / p50,
|
||||||
|
report.reference_brute_force_us / report.reference_hnsw_p50_us);
|
||||||
|
}
|
||||||
|
if let Some(build) = report.lance_index_build_secs {
|
||||||
|
eprintln!(" Index build: {:.1}s (Lance IVF_PQ) vs {:.0}s (our HNSW ec=80) ({}× faster)",
|
||||||
|
build, report.reference_hnsw_build_secs, report.reference_hnsw_build_secs / build);
|
||||||
|
}
|
||||||
|
if let (Some(lance_us), Some(parquet_us)) = (report.lance_random_row_access_us, report.parquet_random_row_access_us) {
|
||||||
|
eprintln!(" Random row access: {:.0}us (Lance) vs {:.0}us (Parquet scan) ({})",
|
||||||
|
lance_us, parquet_us, format_ratio(parquet_us, lance_us));
|
||||||
|
}
|
||||||
|
if let Some(append_secs) = report.lance_append_10k_secs {
|
||||||
|
eprintln!(" Append 10K rows: {:.2}s (Lance native) [Parquet would require full rewrite]",
|
||||||
|
append_secs);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_ratio(baseline: f32, candidate: f32) -> String {
|
||||||
|
if candidate == 0.0 { return "inf".into(); }
|
||||||
|
let ratio = baseline / candidate;
|
||||||
|
if ratio >= 1.0 {
|
||||||
|
format!("{:.2}× faster/smaller", ratio)
|
||||||
|
} else {
|
||||||
|
format!("{:.2}× slower/larger", 1.0 / ratio)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn dir_size_bytes(path: &str) -> u64 {
|
||||||
|
fn recurse(p: &std::path::Path) -> u64 {
|
||||||
|
let Ok(meta) = std::fs::metadata(p) else { return 0; };
|
||||||
|
if meta.is_file() { return meta.len(); }
|
||||||
|
let Ok(entries) = std::fs::read_dir(p) else { return 0; };
|
||||||
|
entries
|
||||||
|
.filter_map(|e| e.ok())
|
||||||
|
.map(|e| recurse(&e.path()))
|
||||||
|
.sum()
|
||||||
|
}
|
||||||
|
recurse(std::path::Path::new(path))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Read the existing vector Parquet (binary-blob format: source, doc_id,
|
||||||
|
/// chunk_idx, chunk_text, vector as Binary bytes).
|
||||||
|
fn read_parquet_vectors(path: &str) -> Result<(Arc<Schema>, Vec<RecordBatch>, usize)> {
|
||||||
|
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
|
||||||
|
use std::fs::File;
|
||||||
|
|
||||||
|
let file = File::open(path).with_context(|| format!("open {path}"))?;
|
||||||
|
let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
|
||||||
|
let schema = builder.schema().clone();
|
||||||
|
let reader = builder.build()?;
|
||||||
|
let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>()?;
|
||||||
|
let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
|
||||||
|
Ok((schema, batches, rows))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn detect_vector_dims(batches: &[RecordBatch]) -> Result<usize> {
|
||||||
|
for batch in batches {
|
||||||
|
let vector_col_idx = batch
|
||||||
|
.schema()
|
||||||
|
.index_of("vector")
|
||||||
|
.context("no 'vector' column in parquet")?;
|
||||||
|
let col = batch.column(vector_col_idx);
|
||||||
|
if let Some(binary) = col.as_any().downcast_ref::<BinaryArray>() {
|
||||||
|
for i in 0..binary.len() {
|
||||||
|
if !binary.is_null(i) {
|
||||||
|
let bytes = binary.value(i);
|
||||||
|
return Ok(bytes.len() / 4); // f32 = 4 bytes
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
anyhow::bail!("could not determine vector dimensions")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Convert our binary-blob vector representation into Arrow's native
|
||||||
|
/// FixedSizeList<Float32> — that's what Lance expects for vector columns.
|
||||||
|
fn convert_to_fixed_size_list(
|
||||||
|
schema: &Arc<Schema>,
|
||||||
|
batches: Vec<RecordBatch>,
|
||||||
|
dims: usize,
|
||||||
|
) -> Result<(Arc<Schema>, Vec<RecordBatch>)> {
|
||||||
|
// New schema keeps everything identical but replaces the vector column
|
||||||
|
// with a FixedSizeList<Float32, dims>.
|
||||||
|
let new_fields: Vec<Arc<Field>> = schema
|
||||||
|
.fields()
|
||||||
|
.iter()
|
||||||
|
.map(|f| {
|
||||||
|
if f.name() == "vector" {
|
||||||
|
Arc::new(Field::new(
|
||||||
|
"vector",
|
||||||
|
DataType::FixedSizeList(
|
||||||
|
Arc::new(Field::new("item", DataType::Float32, true)),
|
||||||
|
dims as i32,
|
||||||
|
),
|
||||||
|
false,
|
||||||
|
))
|
||||||
|
} else {
|
||||||
|
f.clone()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let new_schema = Arc::new(Schema::new(new_fields));
|
||||||
|
|
||||||
|
let mut new_batches = Vec::with_capacity(batches.len());
|
||||||
|
for batch in batches {
|
||||||
|
let vector_idx = batch.schema().index_of("vector")?;
|
||||||
|
let mut new_arrays: Vec<ArrayRef> = Vec::with_capacity(batch.num_columns());
|
||||||
|
for (i, col) in batch.columns().iter().enumerate() {
|
||||||
|
if i == vector_idx {
|
||||||
|
let binary = col
|
||||||
|
.as_any()
|
||||||
|
.downcast_ref::<BinaryArray>()
|
||||||
|
.context("vector column must be Binary")?;
|
||||||
|
let fsl = binary_to_fixed_size_list(binary, dims)?;
|
||||||
|
new_arrays.push(Arc::new(fsl));
|
||||||
|
} else {
|
||||||
|
new_arrays.push(col.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
new_batches.push(RecordBatch::try_new(new_schema.clone(), new_arrays)?);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok((new_schema, new_batches))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn binary_to_fixed_size_list(binary: &BinaryArray, dims: usize) -> Result<FixedSizeListArray> {
|
||||||
|
let n = binary.len();
|
||||||
|
let mut all_floats: Vec<f32> = Vec::with_capacity(n * dims);
|
||||||
|
for i in 0..n {
|
||||||
|
if binary.is_null(i) {
|
||||||
|
all_floats.extend(std::iter::repeat(0.0).take(dims));
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let bytes = binary.value(i);
|
||||||
|
if bytes.len() != dims * 4 {
|
||||||
|
anyhow::bail!(
|
||||||
|
"row {} has {} bytes, expected {} ({} × f32)",
|
||||||
|
i, bytes.len(), dims * 4, dims,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
for chunk in bytes.chunks_exact(4) {
|
||||||
|
all_floats.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let values = Float32Array::from(all_floats);
|
||||||
|
let field = Arc::new(Field::new("item", DataType::Float32, true));
|
||||||
|
FixedSizeListArray::try_new(field, dims as i32, Arc::new(values), None)
|
||||||
|
.context("build FixedSizeListArray")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Write batches into a Lance dataset at the given path.
|
||||||
|
async fn write_lance_dataset(
|
||||||
|
path: &str,
|
||||||
|
schema: Arc<Schema>,
|
||||||
|
batches: Vec<RecordBatch>,
|
||||||
|
) -> Result<()> {
|
||||||
|
use lance::dataset::{Dataset, WriteParams};
|
||||||
|
|
||||||
|
let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema);
|
||||||
|
Dataset::write(reader, path, Some(WriteParams::default()))
|
||||||
|
.await
|
||||||
|
.context("Dataset::write")?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Open a Lance dataset cold (from disk) and scan it fully — measuring the
|
||||||
|
/// equivalent of our "load embeddings from Parquet" cost.
|
||||||
|
async fn cold_open_and_scan_lance(path: &str) -> Result<usize> {
|
||||||
|
use futures::StreamExt;
|
||||||
|
use lance::dataset::Dataset;
|
||||||
|
|
||||||
|
let dataset = Dataset::open(path).await.context("Dataset::open")?;
|
||||||
|
let scanner = dataset.scan();
|
||||||
|
let mut stream = scanner.try_into_stream().await?;
|
||||||
|
let mut total = 0usize;
|
||||||
|
while let Some(batch) = stream.next().await {
|
||||||
|
let batch = batch?;
|
||||||
|
total += batch.num_rows();
|
||||||
|
}
|
||||||
|
Ok(total)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build an IVF_PQ vector index on the `vector` column. IVF_PQ (Inverted File
|
||||||
|
/// with Product Quantization) is Lance's native ANN index — comparable to
|
||||||
|
/// HNSW in intent, but on-disk and compatible with Lance's random-access
|
||||||
|
/// model.
|
||||||
|
async fn build_lance_vector_index(path: &str, _dims: usize) -> Result<()> {
|
||||||
|
use lance::dataset::Dataset;
|
||||||
|
use lance::index::vector::VectorIndexParams;
|
||||||
|
use lance_index::{DatasetIndexExt, IndexType};
|
||||||
|
use lance_linalg::distance::MetricType;
|
||||||
|
|
||||||
|
let mut dataset = Dataset::open(path).await?;
|
||||||
|
|
||||||
|
// IVF_PQ with ~sqrt(N) partitions is a reasonable default for 100K.
|
||||||
|
// num_sub_vectors must divide dims evenly: 768/48 = 16 dims per subvector.
|
||||||
|
// num_bits = 8 gives 256 codes per subvector (good recall/size trade).
|
||||||
|
// max_iterations = 50 is plenty for this scale.
|
||||||
|
let params = VectorIndexParams::ivf_pq(
|
||||||
|
316, // num_partitions (~sqrt(100000))
|
||||||
|
8, // num_bits
|
||||||
|
48, // num_sub_vectors
|
||||||
|
MetricType::Cosine,
|
||||||
|
50, // max_iterations
|
||||||
|
);
|
||||||
|
|
||||||
|
dataset
|
||||||
|
.create_index(
|
||||||
|
&["vector"],
|
||||||
|
IndexType::Vector,
|
||||||
|
Some("vec_idx".into()),
|
||||||
|
¶ms,
|
||||||
|
true,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.context("create_index")?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run N vector searches against the Lance dataset and return (p50, p95, p99) latencies in us.
|
||||||
|
/// Uses a handful of random rows as queries — same pattern as our harness::synthetic_from_chunks.
|
||||||
|
async fn run_search_benchmarks(path: &str, _dims: usize) -> Result<(f32, f32, f32)> {
|
||||||
|
use futures::StreamExt;
|
||||||
|
use lance::dataset::Dataset;
|
||||||
|
|
||||||
|
let dataset = Dataset::open(path).await?;
|
||||||
|
|
||||||
|
// Pick 20 representative query vectors from the data itself.
|
||||||
|
// (Synthetic — same pattern as our existing harness.)
|
||||||
|
let query_vectors = sample_query_vectors(&dataset, 20).await?;
|
||||||
|
|
||||||
|
let mut latencies_us: Vec<f32> = Vec::with_capacity(query_vectors.len());
|
||||||
|
for (i, qv) in query_vectors.iter().enumerate() {
|
||||||
|
let qarr = Arc::new(Float32Array::from(qv.clone())) as ArrayRef;
|
||||||
|
|
||||||
|
let t0 = Instant::now();
|
||||||
|
let mut scanner = dataset.scan();
|
||||||
|
scanner
|
||||||
|
.nearest("vector", qarr.as_any().downcast_ref::<Float32Array>().unwrap(), 10)
|
||||||
|
.context("scanner.nearest")?;
|
||||||
|
let mut stream = scanner.try_into_stream().await?;
|
||||||
|
let mut hits = 0;
|
||||||
|
while let Some(batch) = stream.next().await {
|
||||||
|
let batch = batch?;
|
||||||
|
hits += batch.num_rows();
|
||||||
|
}
|
||||||
|
let us = t0.elapsed().as_micros() as f32;
|
||||||
|
latencies_us.push(us);
|
||||||
|
if i == 0 {
|
||||||
|
eprintln!(" first query: {} hits in {:.0}us (includes any lazy init)", hits, us);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
latencies_us.sort_by(|a, b| a.partial_cmp(b).unwrap());
|
||||||
|
let p = |pct: f32| -> f32 {
|
||||||
|
let idx = ((latencies_us.len() as f32 - 1.0) * pct).round() as usize;
|
||||||
|
latencies_us[idx.min(latencies_us.len() - 1)]
|
||||||
|
};
|
||||||
|
Ok((p(0.50), p(0.95), p(0.99)))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Random row access via Lance's `take` — fetch 20 random rows by index, measure avg latency.
|
||||||
|
async fn measure_random_access_lance(path: &str) -> Result<f32> {
|
||||||
|
use lance::dataset::Dataset;
|
||||||
|
let dataset = Dataset::open(path).await?;
|
||||||
|
let n = dataset.count_rows(None).await?;
|
||||||
|
let indices: Vec<u64> = (0..20).map(|i| ((i as u64) * (n as u64 / 23)) % (n as u64)).collect();
|
||||||
|
|
||||||
|
// Full-schema projection — Lance's Schema implements Into<ProjectionRequest>.
|
||||||
|
let schema = dataset.schema().clone();
|
||||||
|
let mut total_us: u128 = 0;
|
||||||
|
for idx in &indices {
|
||||||
|
let t0 = Instant::now();
|
||||||
|
let _batch = dataset.take(&[*idx], schema.clone()).await?;
|
||||||
|
total_us += t0.elapsed().as_micros();
|
||||||
|
}
|
||||||
|
Ok(total_us as f32 / indices.len() as f32)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Random row access for Parquet — full scan + filter. There's no random-access
|
||||||
|
/// primitive in vanilla Parquet, so this is the cost of finding one specific row.
|
||||||
|
/// This is the cost our current design pays for "get doc X's full text for RAG."
|
||||||
|
fn measure_random_access_parquet(path: &str) -> Result<f32> {
|
||||||
|
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
|
||||||
|
use std::fs::File;
|
||||||
|
|
||||||
|
// We simulate 5 lookups — full scan each time. 20 would be painful.
|
||||||
|
let iters = 5;
|
||||||
|
let mut total_us: u128 = 0;
|
||||||
|
for _ in 0..iters {
|
||||||
|
let t0 = Instant::now();
|
||||||
|
let file = File::open(path)?;
|
||||||
|
let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
|
||||||
|
let reader = builder.build()?;
|
||||||
|
// Iterate until we've conceptually found a row — we stop early if
|
||||||
|
// we wanted row 50000, but we have to at least read its batch.
|
||||||
|
let mut seen = 0usize;
|
||||||
|
for b in reader {
|
||||||
|
let b = b?;
|
||||||
|
seen += b.num_rows();
|
||||||
|
if seen > 50000 { break; }
|
||||||
|
}
|
||||||
|
total_us += t0.elapsed().as_micros();
|
||||||
|
}
|
||||||
|
Ok(total_us as f32 / iters as f32)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append 10K new rows to the existing Lance dataset.
|
||||||
|
/// Measures the "ingest delta" cost without full rewrite.
|
||||||
|
async fn append_10k_rows(path: &str, dims: usize) -> Result<()> {
|
||||||
|
use lance::dataset::{Dataset, WriteMode, WriteParams};
|
||||||
|
|
||||||
|
let dataset = Dataset::open(path).await?;
|
||||||
|
let schema = dataset.schema();
|
||||||
|
let arrow_schema: Arc<Schema> = Arc::new(schema.into());
|
||||||
|
|
||||||
|
// Build a 10K row batch with random-ish data matching the existing schema.
|
||||||
|
let n = 10_000;
|
||||||
|
let arrays: Vec<ArrayRef> = arrow_schema
|
||||||
|
.fields()
|
||||||
|
.iter()
|
||||||
|
.map(|f| -> Result<ArrayRef> {
|
||||||
|
match f.data_type() {
|
||||||
|
DataType::Utf8 => {
|
||||||
|
let vals: Vec<String> = (0..n).map(|i| format!("appended-{}", i)).collect();
|
||||||
|
Ok(Arc::new(arrow_array::StringArray::from(vals)))
|
||||||
|
}
|
||||||
|
DataType::Int32 => {
|
||||||
|
let vals: Vec<i32> = (0..n as i32).collect();
|
||||||
|
Ok(Arc::new(arrow_array::Int32Array::from(vals)))
|
||||||
|
}
|
||||||
|
DataType::FixedSizeList(_, _) => {
|
||||||
|
let floats: Vec<f32> = (0..n * dims).map(|i| (i as f32).sin()).collect();
|
||||||
|
let values = Float32Array::from(floats);
|
||||||
|
let field = Arc::new(Field::new("item", DataType::Float32, true));
|
||||||
|
let fsl = FixedSizeListArray::try_new(field, dims as i32, Arc::new(values), None)?;
|
||||||
|
Ok(Arc::new(fsl))
|
||||||
|
}
|
||||||
|
other => anyhow::bail!("unsupported append column type: {:?}", other),
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect::<Result<Vec<_>>>()?;
|
||||||
|
|
||||||
|
let batch = RecordBatch::try_new(arrow_schema.clone(), arrays)?;
|
||||||
|
let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), arrow_schema);
|
||||||
|
let params = WriteParams { mode: WriteMode::Append, ..Default::default() };
|
||||||
|
Dataset::write(reader, path, Some(params)).await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Grab a few existing vectors from the dataset to use as self-similar queries.
|
||||||
|
async fn sample_query_vectors(
|
||||||
|
dataset: &lance::dataset::Dataset,
|
||||||
|
count: usize,
|
||||||
|
) -> Result<Vec<Vec<f32>>> {
|
||||||
|
use futures::StreamExt;
|
||||||
|
|
||||||
|
// Just take the first `count` rows; good enough for latency measurement.
|
||||||
|
let scanner = dataset.scan();
|
||||||
|
let mut scanner = scanner;
|
||||||
|
scanner.limit(Some(count as i64), None)?;
|
||||||
|
scanner.project(&["vector"])?;
|
||||||
|
let mut stream = scanner.try_into_stream().await?;
|
||||||
|
|
||||||
|
let mut out = Vec::with_capacity(count);
|
||||||
|
while let Some(batch) = stream.next().await {
|
||||||
|
let batch = batch?;
|
||||||
|
let vector_col = batch
|
||||||
|
.column(0)
|
||||||
|
.as_any()
|
||||||
|
.downcast_ref::<FixedSizeListArray>()
|
||||||
|
.context("vector column must be FixedSizeList")?;
|
||||||
|
|
||||||
|
for row in 0..vector_col.len() {
|
||||||
|
if out.len() >= count { break; }
|
||||||
|
let values = vector_col.value(row);
|
||||||
|
let f32_arr = values
|
||||||
|
.as_any()
|
||||||
|
.downcast_ref::<Float32Array>()
|
||||||
|
.context("inner array must be Float32")?;
|
||||||
|
let mut v = Vec::with_capacity(f32_arr.len());
|
||||||
|
for i in 0..f32_arr.len() {
|
||||||
|
v.push(f32_arr.value(i));
|
||||||
|
}
|
||||||
|
out.push(v);
|
||||||
|
}
|
||||||
|
if out.len() >= count { break; }
|
||||||
|
}
|
||||||
|
Ok(out)
|
||||||
|
}
|
||||||
|
|
||||||
@@ -15,6 +15,8 @@ tracing = { workspace = true }
|
|||||||
datafusion = { workspace = true }
|
datafusion = { workspace = true }
|
||||||
object_store = { workspace = true }
|
object_store = { workspace = true }
|
||||||
arrow = { workspace = true }
|
arrow = { workspace = true }
|
||||||
|
parquet = { workspace = true }
|
||||||
bytes = { workspace = true }
|
bytes = { workspace = true }
|
||||||
futures = { workspace = true }
|
futures = { workspace = true }
|
||||||
url = { workspace = true }
|
url = { workspace = true }
|
||||||
|
chrono = { workspace = true }
|
||||||
|
|||||||
@@ -0,0 +1,135 @@
|
|||||||
|
/// In-memory cache for hot datasets.
|
||||||
|
/// Pinned datasets are loaded as Arrow RecordBatches and served from RAM.
|
||||||
|
/// LRU eviction keeps memory bounded.
|
||||||
|
|
||||||
|
use arrow::array::RecordBatch;
|
||||||
|
use arrow::datatypes::SchemaRef;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
/// A cached dataset — schema + all batches in memory.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct CachedDataset {
|
||||||
|
pub name: String,
|
||||||
|
pub schema: SchemaRef,
|
||||||
|
pub batches: Vec<RecordBatch>,
|
||||||
|
pub row_count: usize,
|
||||||
|
pub size_bytes: usize, // approximate
|
||||||
|
pub last_accessed: std::time::Instant,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Memory cache with LRU eviction.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct MemCache {
|
||||||
|
datasets: Arc<RwLock<HashMap<String, CachedDataset>>>,
|
||||||
|
max_bytes: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl MemCache {
|
||||||
|
pub fn new(max_bytes: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
datasets: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
max_bytes,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Load a dataset into the cache.
|
||||||
|
pub async fn put(&self, name: &str, schema: SchemaRef, batches: Vec<RecordBatch>) {
|
||||||
|
let row_count: usize = batches.iter().map(|b| b.num_rows()).sum();
|
||||||
|
let size_bytes: usize = batches.iter()
|
||||||
|
.map(|b| b.get_array_memory_size())
|
||||||
|
.sum();
|
||||||
|
|
||||||
|
// Evict if needed
|
||||||
|
self.evict_to_fit(size_bytes).await;
|
||||||
|
|
||||||
|
let entry = CachedDataset {
|
||||||
|
name: name.to_string(),
|
||||||
|
schema,
|
||||||
|
batches,
|
||||||
|
row_count,
|
||||||
|
size_bytes,
|
||||||
|
last_accessed: std::time::Instant::now(),
|
||||||
|
};
|
||||||
|
|
||||||
|
tracing::info!("cached '{}': {} rows, {:.1} MB",
|
||||||
|
name, row_count, size_bytes as f64 / 1024.0 / 1024.0);
|
||||||
|
|
||||||
|
let mut cache = self.datasets.write().await;
|
||||||
|
cache.insert(name.to_string(), entry);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get a cached dataset, updating access time.
|
||||||
|
pub async fn get(&self, name: &str) -> Option<CachedDataset> {
|
||||||
|
let mut cache = self.datasets.write().await;
|
||||||
|
if let Some(entry) = cache.get_mut(name) {
|
||||||
|
entry.last_accessed = std::time::Instant::now();
|
||||||
|
Some(entry.clone())
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if a dataset is cached.
|
||||||
|
pub async fn contains(&self, name: &str) -> bool {
|
||||||
|
self.datasets.read().await.contains_key(name)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Evict a specific dataset.
|
||||||
|
pub async fn evict(&self, name: &str) -> bool {
|
||||||
|
let mut cache = self.datasets.write().await;
|
||||||
|
cache.remove(name).is_some()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get cache stats.
|
||||||
|
pub async fn stats(&self) -> CacheStats {
|
||||||
|
let cache = self.datasets.read().await;
|
||||||
|
let total_bytes: usize = cache.values().map(|d| d.size_bytes).sum();
|
||||||
|
let total_rows: usize = cache.values().map(|d| d.row_count).sum();
|
||||||
|
CacheStats {
|
||||||
|
datasets: cache.len(),
|
||||||
|
total_bytes,
|
||||||
|
total_rows,
|
||||||
|
max_bytes: self.max_bytes,
|
||||||
|
names: cache.keys().cloned().collect(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Evict least-recently-used entries to make room for `needed_bytes`.
|
||||||
|
async fn evict_to_fit(&self, needed_bytes: usize) {
|
||||||
|
let mut cache = self.datasets.write().await;
|
||||||
|
let current: usize = cache.values().map(|d| d.size_bytes).sum();
|
||||||
|
|
||||||
|
if current + needed_bytes <= self.max_bytes {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sort by last accessed (oldest first)
|
||||||
|
let mut entries: Vec<(String, std::time::Instant, usize)> = cache.iter()
|
||||||
|
.map(|(k, v)| (k.clone(), v.last_accessed, v.size_bytes))
|
||||||
|
.collect();
|
||||||
|
entries.sort_by_key(|(_, t, _)| *t);
|
||||||
|
|
||||||
|
let mut freed = 0usize;
|
||||||
|
let target = (current + needed_bytes).saturating_sub(self.max_bytes);
|
||||||
|
|
||||||
|
for (name, _, size) in entries {
|
||||||
|
if freed >= target {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
tracing::info!("evicting '{}' from cache ({:.1} MB)", name, size as f64 / 1024.0 / 1024.0);
|
||||||
|
cache.remove(&name);
|
||||||
|
freed += size;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct CacheStats {
|
||||||
|
pub datasets: usize,
|
||||||
|
pub total_bytes: usize,
|
||||||
|
pub total_rows: usize,
|
||||||
|
pub max_bytes: usize,
|
||||||
|
pub names: Vec<String>,
|
||||||
|
}
|
||||||
+268
-22
@@ -1,80 +1,326 @@
|
|||||||
use catalogd::registry::Registry;
|
use catalogd::registry::Registry;
|
||||||
use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl};
|
use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl};
|
||||||
|
use datafusion::datasource::MemTable;
|
||||||
use datafusion::datasource::file_format::parquet::ParquetFormat;
|
use datafusion::datasource::file_format::parquet::ParquetFormat;
|
||||||
use datafusion::prelude::*;
|
use datafusion::prelude::*;
|
||||||
use object_store::ObjectStore;
|
use object_store::ObjectStore;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
use storaged::registry::BucketRegistry;
|
||||||
use url::Url;
|
use url::Url;
|
||||||
|
|
||||||
|
use crate::cache::MemCache;
|
||||||
|
use crate::delta;
|
||||||
|
|
||||||
const STORE_SCHEME: &str = "lakehouse";
|
const STORE_SCHEME: &str = "lakehouse";
|
||||||
|
|
||||||
/// Query engine that runs DataFusion over catalog-registered Parquet datasets.
|
/// URL scheme used to register a bucket's object store with DataFusion.
|
||||||
|
/// Bucket names safely embed in URL host components after sanitizing
|
||||||
|
/// `:` (used in `profile:user`) to `-`.
|
||||||
|
fn bucket_scheme(bucket: &str) -> String {
|
||||||
|
format!("{STORE_SCHEME}-{}", bucket.replace(':', "-"))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a `ListingTableUrl`-shaped string for an object in a bucket.
|
||||||
|
fn bucket_object_url(bucket: &str, key: &str) -> String {
|
||||||
|
format!("{}://data/{}", bucket_scheme(bucket), key)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build the registration base URL for a bucket.
|
||||||
|
fn bucket_base_url(bucket: &str) -> String {
|
||||||
|
format!("{}://data/", bucket_scheme(bucket))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Query engine with in-memory cache and delta merge support.
|
||||||
|
/// Federation layer 2: holds a `BucketRegistry` so cross-bucket queries
|
||||||
|
/// route reads to the right ObjectStore based on `ObjectRef.bucket`.
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
pub struct QueryEngine {
|
pub struct QueryEngine {
|
||||||
registry: Registry,
|
registry: Registry,
|
||||||
|
buckets: Arc<BucketRegistry>,
|
||||||
|
/// Primary bucket store cached for backward-compat callers (compact,
|
||||||
|
/// workspace ops) that haven't been migrated to bucket-aware paths.
|
||||||
store: Arc<dyn ObjectStore>,
|
store: Arc<dyn ObjectStore>,
|
||||||
|
cache: MemCache,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl QueryEngine {
|
impl QueryEngine {
|
||||||
pub fn new(registry: Registry, store: Arc<dyn ObjectStore>) -> Self {
|
pub fn new(registry: Registry, buckets: Arc<BucketRegistry>, cache: MemCache) -> Self {
|
||||||
Self { registry, store }
|
let store = buckets.default_store();
|
||||||
|
Self { registry, buckets, store, cache }
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Execute a SQL query. Resolves all catalog datasets as tables before execution.
|
pub fn cache(&self) -> &MemCache {
|
||||||
|
&self.cache
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Backward-compat: returns the primary bucket's store. Bucket-aware
|
||||||
|
/// callers should use `engine.buckets()` directly.
|
||||||
|
pub fn store(&self) -> &Arc<dyn ObjectStore> {
|
||||||
|
&self.store
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn buckets(&self) -> &Arc<BucketRegistry> {
|
||||||
|
&self.buckets
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Catalog registry — used by compact/tombstone integration.
|
||||||
|
pub fn catalog(&self) -> &Registry {
|
||||||
|
&self.registry
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute a SQL query. Uses cache for hot data, falls back to Parquet.
|
||||||
pub async fn query(&self, sql: &str) -> Result<Vec<arrow::array::RecordBatch>, String> {
|
pub async fn query(&self, sql: &str) -> Result<Vec<arrow::array::RecordBatch>, String> {
|
||||||
let ctx = self.build_context().await?;
|
let ctx = self.build_context().await?;
|
||||||
|
|
||||||
let df = ctx.sql(sql).await.map_err(|e| format!("SQL error: {e}"))?;
|
let df = ctx.sql(sql).await.map_err(|e| format!("SQL error: {e}"))?;
|
||||||
let batches = df.collect().await.map_err(|e| format!("execution error: {e}"))?;
|
let batches = df.collect().await.map_err(|e| format!("execution error: {e}"))?;
|
||||||
Ok(batches)
|
Ok(batches)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Build a SessionContext with all catalog datasets registered as tables.
|
/// Pin a dataset into the memory cache.
|
||||||
|
pub async fn pin_dataset(&self, name: &str) -> Result<(), String> {
|
||||||
|
let ctx = SessionContext::new();
|
||||||
|
|
||||||
|
// Register every bucket so a multi-bucket dataset can be pinned.
|
||||||
|
for info in self.buckets.list().await {
|
||||||
|
let url = Url::parse(&bucket_base_url(&info.name))
|
||||||
|
.map_err(|e| format!("invalid store url: {e}"))?;
|
||||||
|
let store = self.buckets.get(&info.name)
|
||||||
|
.map_err(|e| format!("registry inconsistency: {e}"))?;
|
||||||
|
ctx.runtime_env().register_object_store(&url, store);
|
||||||
|
}
|
||||||
|
|
||||||
|
let dataset = self.registry.get_by_name(name).await
|
||||||
|
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||||
|
|
||||||
|
if dataset.objects.is_empty() {
|
||||||
|
return Err(format!("dataset '{name}' has no objects"));
|
||||||
|
}
|
||||||
|
|
||||||
|
let opts = ListingOptions::new(Arc::new(ParquetFormat::default()));
|
||||||
|
let table_paths: Vec<ListingTableUrl> = dataset.objects.iter()
|
||||||
|
.filter_map(|o| {
|
||||||
|
let bucket = if o.bucket.is_empty() { "primary" } else { &o.bucket };
|
||||||
|
ListingTableUrl::parse(&bucket_object_url(bucket, &o.key)).ok()
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let schema = opts.infer_schema(&ctx.state(), &table_paths[0]).await
|
||||||
|
.map_err(|e| format!("schema inference: {e}"))?;
|
||||||
|
let config = ListingTableConfig::new_with_multi_paths(table_paths)
|
||||||
|
.with_listing_options(opts)
|
||||||
|
.with_schema(schema.clone());
|
||||||
|
let table = ListingTable::try_new(config).map_err(|e| format!("table: {e}"))?;
|
||||||
|
|
||||||
|
ctx.register_table(name, Arc::new(table)).map_err(|e| format!("register: {e}"))?;
|
||||||
|
let df = ctx.sql(&format!("SELECT * FROM {name}")).await.map_err(|e| format!("read: {e}"))?;
|
||||||
|
let batches = df.collect().await.map_err(|e| format!("collect: {e}"))?;
|
||||||
|
|
||||||
|
self.cache.put(name, schema, batches).await;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
async fn build_context(&self) -> Result<SessionContext, String> {
|
async fn build_context(&self) -> Result<SessionContext, String> {
|
||||||
let ctx = SessionContext::new();
|
let ctx = SessionContext::new();
|
||||||
|
|
||||||
// Register the object store under a custom scheme to avoid path doubling.
|
// Phase E: snapshot tombstones by dataset before registering tables
|
||||||
// The store already has the root prefix (e.g. ./data), so paths are relative keys.
|
// so we can wrap tombstoned tables in a filter view. The underlying
|
||||||
let base_url = Url::parse(&format!("{STORE_SCHEME}://data/"))
|
// base table is registered under an internal name `__raw__{dataset}`
|
||||||
.map_err(|e| format!("invalid store url: {e}"))?;
|
// and the public `{dataset}` name becomes the filtered view.
|
||||||
ctx.runtime_env().register_object_store(&base_url, self.store.clone());
|
let all_dataset_names: Vec<String> = self.registry.list().await
|
||||||
|
.iter().map(|d| d.name.clone()).collect();
|
||||||
|
let tombstones_by_dataset = self
|
||||||
|
.registry
|
||||||
|
.tombstones()
|
||||||
|
.all_grouped(&all_dataset_names)
|
||||||
|
.await
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
// Federation layer 2: register every configured bucket as its own
|
||||||
|
// DataFusion ObjectStore under a distinct URL scheme. Each
|
||||||
|
// dataset's ObjectRef.bucket determines which store DataFusion
|
||||||
|
// routes to at scan time.
|
||||||
|
let bucket_infos = self.buckets.list().await;
|
||||||
|
for info in &bucket_infos {
|
||||||
|
let url_str = bucket_base_url(&info.name);
|
||||||
|
let url = Url::parse(&url_str)
|
||||||
|
.map_err(|e| format!("invalid store url '{url_str}': {e}"))?;
|
||||||
|
// unknown-bucket here would be a config invariant violation,
|
||||||
|
// since we just listed them — propagate as an error to surface it.
|
||||||
|
let store = self.buckets.get(&info.name)
|
||||||
|
.map_err(|e| format!("registry inconsistency: {e}"))?;
|
||||||
|
ctx.runtime_env().register_object_store(&url, store);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Backward-compat: also register the legacy `lakehouse://data/`
|
||||||
|
// URL pointing at primary, in case any code path still constructs
|
||||||
|
// ListingTableUrls without a bucket prefix.
|
||||||
|
let legacy_url = Url::parse(&format!("{STORE_SCHEME}://data/"))
|
||||||
|
.map_err(|e| format!("invalid legacy store url: {e}"))?;
|
||||||
|
ctx.runtime_env().register_object_store(&legacy_url, self.store.clone());
|
||||||
|
|
||||||
// Register each catalog dataset as a table
|
|
||||||
let datasets = self.registry.list().await;
|
let datasets = self.registry.list().await;
|
||||||
for dataset in &datasets {
|
for dataset in &datasets {
|
||||||
if dataset.objects.is_empty() {
|
if dataset.objects.is_empty() {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
let opts = ListingOptions::new(Arc::new(ParquetFormat::default()));
|
// Check cache first — cached batches are bucket-agnostic since
|
||||||
|
// they're already materialized in memory.
|
||||||
|
if let Some(cached) = self.cache.get(&dataset.name).await {
|
||||||
|
let delta_batches = delta::load_deltas(&self.store, &dataset.name).await.unwrap_or_default();
|
||||||
|
|
||||||
|
let mut all_batches = cached.batches;
|
||||||
|
all_batches.extend(delta_batches);
|
||||||
|
|
||||||
|
let mem_table = MemTable::try_new(cached.schema, vec![all_batches])
|
||||||
|
.map_err(|e| format!("MemTable error for {}: {e}", dataset.name))?;
|
||||||
|
ctx.register_table(&dataset.name, Arc::new(mem_table))
|
||||||
|
.map_err(|e| format!("register cached {}: {e}", dataset.name))?;
|
||||||
|
|
||||||
|
tracing::debug!("using cached table: {}", dataset.name);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build ListingTable URLs that include the per-object bucket
|
||||||
|
// scheme. DataFusion routes each scan to the right store
|
||||||
|
// automatically based on the URL.
|
||||||
|
let opts = ListingOptions::new(Arc::new(ParquetFormat::default()));
|
||||||
let table_paths: Vec<ListingTableUrl> = dataset.objects.iter()
|
let table_paths: Vec<ListingTableUrl> = dataset.objects.iter()
|
||||||
.filter_map(|o| {
|
.filter_map(|o| {
|
||||||
let url_str = format!("{STORE_SCHEME}://data/{}", o.key);
|
let bucket = if o.bucket.is_empty() { "primary" } else { &o.bucket };
|
||||||
ListingTableUrl::parse(&url_str).ok()
|
ListingTableUrl::parse(&bucket_object_url(bucket, &o.key)).ok()
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
if table_paths.is_empty() {
|
if table_paths.is_empty() {
|
||||||
tracing::warn!("dataset {} has no valid paths, skipping", dataset.name);
|
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Infer schema from the first file
|
|
||||||
let schema = opts.infer_schema(&ctx.state(), &table_paths[0]).await
|
let schema = opts.infer_schema(&ctx.state(), &table_paths[0]).await
|
||||||
.map_err(|e| format!("schema inference failed for {}: {e}", dataset.name))?;
|
.map_err(|e| format!("schema inference failed for {}: {e}", dataset.name))?;
|
||||||
|
|
||||||
let config = ListingTableConfig::new_with_multi_paths(table_paths)
|
let config = ListingTableConfig::new_with_multi_paths(table_paths)
|
||||||
.with_listing_options(opts)
|
.with_listing_options(opts)
|
||||||
.with_schema(schema);
|
.with_schema(schema);
|
||||||
|
|
||||||
let table = ListingTable::try_new(config)
|
let table = ListingTable::try_new(config)
|
||||||
.map_err(|e| format!("table creation failed for {}: {e}", dataset.name))?;
|
.map_err(|e| format!("table creation failed for {}: {e}", dataset.name))?;
|
||||||
|
|
||||||
ctx.register_table(&dataset.name, Arc::new(table))
|
// Decide the registration name: if this dataset has any
|
||||||
.map_err(|e| format!("table registration failed for {}: {e}", dataset.name))?;
|
// tombstones, the raw table gets an internal name and the
|
||||||
|
// public name becomes a filtered view.
|
||||||
|
let tombstone_entry = tombstones_by_dataset.get(&dataset.name);
|
||||||
|
let register_name = if tombstone_entry.is_some() {
|
||||||
|
format!("__raw__{}", dataset.name)
|
||||||
|
} else {
|
||||||
|
dataset.name.clone()
|
||||||
|
};
|
||||||
|
|
||||||
tracing::debug!("registered table: {}", dataset.name);
|
if let Err(e) = ctx.register_table(register_name.as_str(), Arc::new(table)) {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if msg.contains("already exists") {
|
||||||
|
tracing::debug!("skip duplicate manifest registration: {}", register_name);
|
||||||
|
continue;
|
||||||
|
} else {
|
||||||
|
return Err(format!("table registration failed for {}: {}", register_name, msg));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// If there are tombstones, register the public name as a
|
||||||
|
// filtered view that excludes tombstoned row_key_values.
|
||||||
|
if let Some((key_col, values)) = tombstone_entry {
|
||||||
|
// Build WHERE NOT IN (...) — quote values to be SQL-safe.
|
||||||
|
// For string keys this is a literal list; for integer keys
|
||||||
|
// CAST(col AS VARCHAR) makes the comparison unambiguous.
|
||||||
|
let quoted: Vec<String> = values.iter()
|
||||||
|
.map(|v| format!("'{}'", v.replace('\'', "''")))
|
||||||
|
.collect();
|
||||||
|
let sql = format!(
|
||||||
|
"SELECT * FROM \"{}\" WHERE CAST(\"{}\" AS VARCHAR) NOT IN ({})",
|
||||||
|
register_name, key_col, quoted.join(", "),
|
||||||
|
);
|
||||||
|
tracing::debug!(
|
||||||
|
"tombstone filter for '{}': {} row_keys excluded",
|
||||||
|
dataset.name, values.len(),
|
||||||
|
);
|
||||||
|
match ctx.sql(&sql).await {
|
||||||
|
Ok(df) => {
|
||||||
|
if let Err(e) = ctx.register_table(dataset.name.as_str(), df.into_view()) {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if msg.contains("already exists") {
|
||||||
|
tracing::debug!("skip duplicate tombstone view: {}", dataset.name);
|
||||||
|
} else {
|
||||||
|
tracing::warn!("tombstone view registration failed for {}: {msg}", dataset.name);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!("tombstone view SQL failed for '{}': {e}", dataset.name);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Phase D: register every AiView as a DataFusion view that wraps
|
||||||
|
// the base table with the safe projection + filter + redactions.
|
||||||
|
for view in self.registry.list_views().await {
|
||||||
|
// Build the SELECT clause with column whitelist and any
|
||||||
|
// redaction expressions. Each column is either the raw column
|
||||||
|
// name or an expression aliased back to the column name so
|
||||||
|
// downstream queries see the same shape.
|
||||||
|
let select_cols: Vec<String> = view.columns.iter().map(|col| {
|
||||||
|
match view.column_redactions.get(col) {
|
||||||
|
None => format!("\"{}\"", col),
|
||||||
|
Some(shared::types::Redaction::Null) => {
|
||||||
|
format!("CAST(NULL AS VARCHAR) AS \"{}\"", col)
|
||||||
|
}
|
||||||
|
Some(shared::types::Redaction::Hash) => {
|
||||||
|
// DataFusion has digest('sha256', value) — but column
|
||||||
|
// type might not be string. Cast first to be safe.
|
||||||
|
format!("digest(CAST(\"{}\" AS VARCHAR), 'sha256') AS \"{}\"", col, col)
|
||||||
|
}
|
||||||
|
Some(shared::types::Redaction::Mask { keep_prefix, keep_suffix }) => {
|
||||||
|
// Keep first N + last M chars, replace middle with stars.
|
||||||
|
// CASE NULL through with COALESCE for null-safe behavior.
|
||||||
|
format!(
|
||||||
|
"CASE WHEN \"{c}\" IS NULL THEN NULL ELSE \
|
||||||
|
concat(\
|
||||||
|
substr(CAST(\"{c}\" AS VARCHAR), 1, {p}), \
|
||||||
|
repeat('*', greatest(0, length(CAST(\"{c}\" AS VARCHAR)) - {p} - {s})), \
|
||||||
|
substr(CAST(\"{c}\" AS VARCHAR), greatest(1, length(CAST(\"{c}\" AS VARCHAR)) - {s} + 1))\
|
||||||
|
) END AS \"{c}\"",
|
||||||
|
c = col, p = keep_prefix, s = keep_suffix,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
|
||||||
|
let where_clause = view.row_filter
|
||||||
|
.as_deref()
|
||||||
|
.map(|f| format!(" WHERE {}", f))
|
||||||
|
.unwrap_or_default();
|
||||||
|
let view_sql = format!(
|
||||||
|
"SELECT {} FROM \"{}\"{}",
|
||||||
|
select_cols.join(", "),
|
||||||
|
view.base_dataset,
|
||||||
|
where_clause,
|
||||||
|
);
|
||||||
|
|
||||||
|
tracing::debug!("registering AiView '{}': {}", view.name, view_sql);
|
||||||
|
|
||||||
|
let df = match ctx.sql(&view_sql).await {
|
||||||
|
Ok(df) => df,
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!("view '{}' SQL invalid, skipping: {e}", view.name);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if let Err(e) = ctx.register_table(view.name.as_str(), df.into_view()) {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if msg.contains("already exists") {
|
||||||
|
tracing::debug!("skip duplicate view registration: {}", view.name);
|
||||||
|
} else {
|
||||||
|
tracing::warn!("view '{}' registration failed: {e}", view.name);
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(ctx)
|
Ok(ctx)
|
||||||
|
|||||||
@@ -0,0 +1,269 @@
|
|||||||
|
/// Delta store for incremental updates.
|
||||||
|
/// Instead of rewriting an entire Parquet file to change one row,
|
||||||
|
/// we write small delta files. At query time, deltas are merged with the base.
|
||||||
|
/// Periodic compaction merges deltas into the base file.
|
||||||
|
|
||||||
|
use arrow::array::{Array, BooleanArray, Int32Array, Int64Array, RecordBatch, StringArray};
|
||||||
|
use arrow::compute::filter_record_batch;
|
||||||
|
use bytes::Bytes;
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use parquet::arrow::ArrowWriter;
|
||||||
|
use parquet::basic::Compression;
|
||||||
|
use parquet::file::properties::WriterProperties;
|
||||||
|
use shared::types::Tombstone;
|
||||||
|
use std::collections::HashSet;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use shared::arrow_helpers::{parquet_to_record_batches, record_batch_to_parquet};
|
||||||
|
use storaged::ops;
|
||||||
|
|
||||||
|
/// Write a delta file for a dataset (new/updated rows).
|
||||||
|
pub async fn write_delta(
|
||||||
|
store: &Arc<dyn ObjectStore>,
|
||||||
|
dataset_name: &str,
|
||||||
|
batch: &RecordBatch,
|
||||||
|
) -> Result<String, String> {
|
||||||
|
let ts = chrono::Utc::now().timestamp_millis();
|
||||||
|
let key = format!("deltas/{dataset_name}/{ts}.parquet");
|
||||||
|
let parquet = record_batch_to_parquet(batch)?;
|
||||||
|
ops::put(store, &key, parquet).await?;
|
||||||
|
tracing::info!("wrote delta for '{}': {} rows at {}", dataset_name, batch.num_rows(), key);
|
||||||
|
Ok(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all delta files for a dataset.
|
||||||
|
pub async fn list_deltas(
|
||||||
|
store: &Arc<dyn ObjectStore>,
|
||||||
|
dataset_name: &str,
|
||||||
|
) -> Result<Vec<String>, String> {
|
||||||
|
let prefix = format!("deltas/{dataset_name}/");
|
||||||
|
ops::list(store, Some(&prefix)).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Load all delta batches for a dataset.
|
||||||
|
pub async fn load_deltas(
|
||||||
|
store: &Arc<dyn ObjectStore>,
|
||||||
|
dataset_name: &str,
|
||||||
|
) -> Result<Vec<RecordBatch>, String> {
|
||||||
|
let keys = list_deltas(store, dataset_name).await?;
|
||||||
|
let mut all_batches = Vec::new();
|
||||||
|
|
||||||
|
for key in &keys {
|
||||||
|
let data = ops::get(store, key).await?;
|
||||||
|
let (_, batches) = parquet_to_record_batches(&data)?;
|
||||||
|
all_batches.extend(batches);
|
||||||
|
}
|
||||||
|
|
||||||
|
if !all_batches.is_empty() {
|
||||||
|
let total_rows: usize = all_batches.iter().map(|b| b.num_rows()).sum();
|
||||||
|
tracing::debug!("loaded {} delta files ({} rows) for '{}'", keys.len(), total_rows, dataset_name);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(all_batches)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Compact: merge base Parquet + all deltas into a single new base file.
|
||||||
|
/// Optionally filters out tombstoned rows — this is the physical-deletion
|
||||||
|
/// half of Phase E (the query-time filter handles the "immediate hide",
|
||||||
|
/// compaction handles the "actually remove from disk" requirement).
|
||||||
|
///
|
||||||
|
/// Tombstones must share one `row_key_column` (enforced at tombstone write).
|
||||||
|
/// If non-empty, every row in the merged output whose row_key_column value
|
||||||
|
/// matches a tombstone is dropped before the new base is written.
|
||||||
|
pub async fn compact(
|
||||||
|
store: &Arc<dyn ObjectStore>,
|
||||||
|
dataset_name: &str,
|
||||||
|
base_key: &str,
|
||||||
|
primary_key_col: Option<&str>,
|
||||||
|
tombstones: &[Tombstone],
|
||||||
|
) -> Result<CompactResult, String> {
|
||||||
|
// Load base
|
||||||
|
let base_data = ops::get(store, base_key).await?;
|
||||||
|
let (_schema, mut base_batches) = parquet_to_record_batches(&base_data)?;
|
||||||
|
|
||||||
|
// Load deltas
|
||||||
|
let delta_batches = load_deltas(store, dataset_name).await?;
|
||||||
|
let delta_count = delta_batches.len();
|
||||||
|
|
||||||
|
let has_tombstones = !tombstones.is_empty();
|
||||||
|
let nothing_to_do = delta_batches.is_empty() && !has_tombstones;
|
||||||
|
if nothing_to_do {
|
||||||
|
return Ok(CompactResult {
|
||||||
|
base_rows: base_batches.iter().map(|b| b.num_rows()).sum(),
|
||||||
|
delta_rows: 0,
|
||||||
|
final_rows: base_batches.iter().map(|b| b.num_rows()).sum(),
|
||||||
|
deltas_merged: 0,
|
||||||
|
tombstones_applied: 0,
|
||||||
|
rows_dropped_by_tombstones: 0,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
base_batches.extend(delta_batches);
|
||||||
|
let pre_filter_rows: usize = base_batches.iter().map(|b| b.num_rows()).sum();
|
||||||
|
|
||||||
|
// If primary key specified, deduplicate (keep last occurrence)
|
||||||
|
let merged_batches = if let Some(_pk) = primary_key_col {
|
||||||
|
// Current simplification — full dedup requires sort by PK; leave
|
||||||
|
// that work for a follow-up. Tombstone filter still runs below.
|
||||||
|
base_batches
|
||||||
|
} else {
|
||||||
|
base_batches
|
||||||
|
};
|
||||||
|
|
||||||
|
// Tombstone filter: drop rows whose row_key matches.
|
||||||
|
let (final_batches, dropped) = if has_tombstones {
|
||||||
|
apply_tombstone_filter(merged_batches, tombstones)?
|
||||||
|
} else {
|
||||||
|
(merged_batches, 0)
|
||||||
|
};
|
||||||
|
|
||||||
|
let final_rows: usize = final_batches.iter().map(|b| b.num_rows()).sum();
|
||||||
|
|
||||||
|
// Write merged base atomically: build ONE Parquet file (single
|
||||||
|
// footer!) via ArrowWriter, stage it under a temporary key, read
|
||||||
|
// back to verify row count matches what we intended, THEN swap
|
||||||
|
// into base_key. If anything goes wrong between build and swap,
|
||||||
|
// the original base_key is untouched.
|
||||||
|
//
|
||||||
|
// Why this matters: concatenating N record_batch_to_parquet()
|
||||||
|
// outputs produces N independent Parquet files glued together —
|
||||||
|
// readers see only the first footer. That silently corrupted the
|
||||||
|
// candidates table before this code was fixed.
|
||||||
|
let schema_ref = final_batches
|
||||||
|
.first()
|
||||||
|
.map(|b| b.schema())
|
||||||
|
.ok_or_else(|| "no batches to write".to_string())?;
|
||||||
|
|
||||||
|
// Use Snappy to match the compression of files written by ingest —
|
||||||
|
// otherwise compaction would inflate file size 3× on each rewrite.
|
||||||
|
let props = WriterProperties::builder()
|
||||||
|
.set_compression(Compression::SNAPPY)
|
||||||
|
.build();
|
||||||
|
|
||||||
|
let mut buf: Vec<u8> = Vec::with_capacity(16 * 1024 * 1024);
|
||||||
|
{
|
||||||
|
let mut writer = ArrowWriter::try_new(&mut buf, schema_ref.clone(), Some(props))
|
||||||
|
.map_err(|e| format!("ArrowWriter init: {e}"))?;
|
||||||
|
for batch in &final_batches {
|
||||||
|
writer.write(batch).map_err(|e| format!("write batch: {e}"))?;
|
||||||
|
}
|
||||||
|
writer.close().map_err(|e| format!("close writer: {e}"))?;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify before committing — if the written bytes don't parse as
|
||||||
|
// the row count we expect, DON'T overwrite the base file.
|
||||||
|
let (_, verify_batches) = parquet_to_record_batches(&buf)
|
||||||
|
.map_err(|e| format!("verify parse: {e}"))?;
|
||||||
|
let verify_rows: usize = verify_batches.iter().map(|b| b.num_rows()).sum();
|
||||||
|
if verify_rows != final_rows {
|
||||||
|
return Err(format!(
|
||||||
|
"compact verification failed: wrote {} bytes that parse as {} rows (expected {}); base_key '{}' untouched",
|
||||||
|
buf.len(), verify_rows, final_rows, base_key,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stage under a temp key first, then promote via a second put
|
||||||
|
// over base_key. Object stores don't have atomic rename across
|
||||||
|
// arbitrary backends; this is "write twice + delete temp" which
|
||||||
|
// is the widely-adopted pattern.
|
||||||
|
let temp_key = format!("{base_key}.compact-{}.tmp", chrono::Utc::now().timestamp_millis());
|
||||||
|
ops::put(store, &temp_key, Bytes::from(buf.clone())).await?;
|
||||||
|
ops::put(store, base_key, Bytes::from(buf)).await?;
|
||||||
|
let _ = ops::delete(store, &temp_key).await;
|
||||||
|
|
||||||
|
// Only now that base is durably updated, delete delta files.
|
||||||
|
let delta_keys = list_deltas(store, dataset_name).await?;
|
||||||
|
for key in &delta_keys {
|
||||||
|
let _ = ops::delete(store, key).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
"compacted '{}': {} deltas merged, {} tombstones applied ({} rows dropped), {} → {} rows",
|
||||||
|
dataset_name, delta_count, tombstones.len(), dropped, pre_filter_rows, final_rows,
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(CompactResult {
|
||||||
|
base_rows: pre_filter_rows - delta_count, // rough base-before-deltas
|
||||||
|
delta_rows: delta_count,
|
||||||
|
final_rows,
|
||||||
|
deltas_merged: delta_count,
|
||||||
|
tombstones_applied: tombstones.len(),
|
||||||
|
rows_dropped_by_tombstones: dropped,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Filter every batch in `batches` to exclude rows whose `row_key_column`
|
||||||
|
/// value appears in the tombstone set. All tombstones must share one
|
||||||
|
/// row_key_column — enforced at tombstone write time.
|
||||||
|
fn apply_tombstone_filter(
|
||||||
|
batches: Vec<RecordBatch>,
|
||||||
|
tombstones: &[Tombstone],
|
||||||
|
) -> Result<(Vec<RecordBatch>, usize), String> {
|
||||||
|
let col_name = &tombstones[0].row_key_column;
|
||||||
|
let bad_values: HashSet<String> = tombstones
|
||||||
|
.iter()
|
||||||
|
.map(|t| t.row_key_value.clone())
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let mut out = Vec::with_capacity(batches.len());
|
||||||
|
let mut total_dropped = 0usize;
|
||||||
|
|
||||||
|
for batch in batches {
|
||||||
|
let col_idx = match batch.schema().index_of(col_name) {
|
||||||
|
Ok(i) => i,
|
||||||
|
Err(_) => {
|
||||||
|
// Column not in this batch — nothing to drop, pass through.
|
||||||
|
out.push(batch);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let col = batch.column(col_idx);
|
||||||
|
|
||||||
|
// Build a "keep" mask: true where row is NOT tombstoned.
|
||||||
|
let mask: Vec<bool> = if let Some(s) = col.as_any().downcast_ref::<StringArray>() {
|
||||||
|
(0..s.len()).map(|i| {
|
||||||
|
s.is_null(i) || !bad_values.contains(s.value(i))
|
||||||
|
}).collect()
|
||||||
|
} else if let Some(a) = col.as_any().downcast_ref::<Int32Array>() {
|
||||||
|
(0..a.len()).map(|i| {
|
||||||
|
a.is_null(i) || !bad_values.contains(&a.value(i).to_string())
|
||||||
|
}).collect()
|
||||||
|
} else if let Some(a) = col.as_any().downcast_ref::<Int64Array>() {
|
||||||
|
(0..a.len()).map(|i| {
|
||||||
|
a.is_null(i) || !bad_values.contains(&a.value(i).to_string())
|
||||||
|
}).collect()
|
||||||
|
} else {
|
||||||
|
// Key column of an unsupported type — log and pass through
|
||||||
|
// rather than panic. Operator has to handle this manually.
|
||||||
|
tracing::warn!(
|
||||||
|
"tombstone filter: column '{}' has unsupported type {:?}; passing rows through",
|
||||||
|
col_name, col.data_type(),
|
||||||
|
);
|
||||||
|
out.push(batch);
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
|
||||||
|
let kept = mask.iter().filter(|b| **b).count();
|
||||||
|
let dropped = mask.len() - kept;
|
||||||
|
total_dropped += dropped;
|
||||||
|
if dropped == 0 {
|
||||||
|
out.push(batch);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let bool_arr = BooleanArray::from(mask);
|
||||||
|
let filtered = filter_record_batch(&batch, &bool_arr)
|
||||||
|
.map_err(|e| format!("filter batch: {e}"))?;
|
||||||
|
out.push(filtered);
|
||||||
|
}
|
||||||
|
Ok((out, total_dropped))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct CompactResult {
|
||||||
|
pub base_rows: usize,
|
||||||
|
pub delta_rows: usize,
|
||||||
|
pub final_rows: usize,
|
||||||
|
pub deltas_merged: usize,
|
||||||
|
pub tombstones_applied: usize,
|
||||||
|
pub rows_dropped_by_tombstones: usize,
|
||||||
|
}
|
||||||
@@ -1,2 +1,7 @@
|
|||||||
|
pub mod cache;
|
||||||
pub mod context;
|
pub mod context;
|
||||||
|
pub mod delta;
|
||||||
|
pub mod paged;
|
||||||
pub mod service;
|
pub mod service;
|
||||||
|
pub mod workspace;
|
||||||
|
pub mod workspace_service;
|
||||||
|
|||||||
@@ -0,0 +1,222 @@
|
|||||||
|
/// Paged query results with server-side cursor.
|
||||||
|
/// Large results are stored server-side, client fetches pages on demand.
|
||||||
|
/// Pattern matches the embedding supervisor: chunk → process → retry.
|
||||||
|
|
||||||
|
use arrow::array::RecordBatch;
|
||||||
|
use arrow::json::writer::{JsonArray, Writer as JsonWriter};
|
||||||
|
use serde::Serialize;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
use crate::context::QueryEngine;
|
||||||
|
|
||||||
|
/// A stored query result that clients page through.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct StoredResult {
|
||||||
|
pub query_id: String,
|
||||||
|
pub sql: String,
|
||||||
|
pub columns: Vec<ColumnInfo>,
|
||||||
|
pub total_rows: usize,
|
||||||
|
pub batches: Vec<RecordBatch>,
|
||||||
|
pub created_at: std::time::Instant,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Serialize)]
|
||||||
|
pub struct ColumnInfo {
|
||||||
|
pub name: String,
|
||||||
|
pub data_type: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Response for initial query — metadata only, no rows.
|
||||||
|
#[derive(Serialize)]
|
||||||
|
pub struct QueryHandle {
|
||||||
|
pub query_id: String,
|
||||||
|
pub total_rows: usize,
|
||||||
|
pub columns: Vec<ColumnInfo>,
|
||||||
|
pub total_pages: usize,
|
||||||
|
pub page_size: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A page of results.
|
||||||
|
#[derive(Serialize)]
|
||||||
|
pub struct PageResult {
|
||||||
|
pub query_id: String,
|
||||||
|
pub page: usize,
|
||||||
|
pub page_size: usize,
|
||||||
|
pub total_rows: usize,
|
||||||
|
pub total_pages: usize,
|
||||||
|
pub rows: serde_json::Value,
|
||||||
|
pub row_count: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Server-side result store with TTL.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct ResultStore {
|
||||||
|
results: Arc<RwLock<HashMap<String, StoredResult>>>,
|
||||||
|
default_page_size: usize,
|
||||||
|
max_results: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ResultStore {
|
||||||
|
pub fn new(default_page_size: usize, max_results: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
results: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
default_page_size,
|
||||||
|
max_results,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute query and store results. Returns handle with metadata.
|
||||||
|
pub async fn execute_and_store(
|
||||||
|
&self,
|
||||||
|
engine: &QueryEngine,
|
||||||
|
sql: &str,
|
||||||
|
) -> Result<QueryHandle, String> {
|
||||||
|
let batches = engine.query(sql).await?;
|
||||||
|
|
||||||
|
if batches.is_empty() {
|
||||||
|
let qid = format!("qr-{}", chrono::Utc::now().timestamp_millis());
|
||||||
|
return Ok(QueryHandle {
|
||||||
|
query_id: qid,
|
||||||
|
total_rows: 0,
|
||||||
|
columns: vec![],
|
||||||
|
total_pages: 0,
|
||||||
|
page_size: self.default_page_size,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let schema = batches[0].schema();
|
||||||
|
let columns: Vec<ColumnInfo> = schema.fields().iter().map(|f| ColumnInfo {
|
||||||
|
name: f.name().clone(),
|
||||||
|
data_type: f.data_type().to_string(),
|
||||||
|
}).collect();
|
||||||
|
|
||||||
|
let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
|
||||||
|
let total_pages = (total_rows + self.default_page_size - 1) / self.default_page_size;
|
||||||
|
|
||||||
|
let qid = format!("qr-{}", chrono::Utc::now().timestamp_millis());
|
||||||
|
|
||||||
|
let result = StoredResult {
|
||||||
|
query_id: qid.clone(),
|
||||||
|
sql: sql.to_string(),
|
||||||
|
columns: columns.clone(),
|
||||||
|
total_rows,
|
||||||
|
batches,
|
||||||
|
created_at: std::time::Instant::now(),
|
||||||
|
};
|
||||||
|
|
||||||
|
// Store (evict old results if at capacity)
|
||||||
|
let mut store = self.results.write().await;
|
||||||
|
if store.len() >= self.max_results {
|
||||||
|
// Evict oldest
|
||||||
|
if let Some(oldest_key) = store.iter()
|
||||||
|
.min_by_key(|(_, v)| v.created_at)
|
||||||
|
.map(|(k, _)| k.clone())
|
||||||
|
{
|
||||||
|
store.remove(&oldest_key);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
store.insert(qid.clone(), result);
|
||||||
|
|
||||||
|
Ok(QueryHandle {
|
||||||
|
query_id: qid,
|
||||||
|
total_rows,
|
||||||
|
columns,
|
||||||
|
total_pages,
|
||||||
|
page_size: self.default_page_size,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get a page of results. Handles chunking and retry-safe.
|
||||||
|
pub async fn get_page(
|
||||||
|
&self,
|
||||||
|
query_id: &str,
|
||||||
|
page: usize,
|
||||||
|
page_size: Option<usize>,
|
||||||
|
) -> Result<PageResult, String> {
|
||||||
|
let store = self.results.read().await;
|
||||||
|
let result = store.get(query_id)
|
||||||
|
.ok_or_else(|| format!("query result not found: {query_id} (may have expired)"))?;
|
||||||
|
|
||||||
|
let ps = page_size.unwrap_or(self.default_page_size);
|
||||||
|
let start_row = page * ps;
|
||||||
|
|
||||||
|
if start_row >= result.total_rows {
|
||||||
|
return Ok(PageResult {
|
||||||
|
query_id: query_id.to_string(),
|
||||||
|
page,
|
||||||
|
page_size: ps,
|
||||||
|
total_rows: result.total_rows,
|
||||||
|
total_pages: (result.total_rows + ps - 1) / ps,
|
||||||
|
rows: serde_json::Value::Array(vec![]),
|
||||||
|
row_count: 0,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let end_row = (start_row + ps).min(result.total_rows);
|
||||||
|
|
||||||
|
// Extract the right rows from batches
|
||||||
|
let page_batches = slice_batches(&result.batches, start_row, end_row);
|
||||||
|
let rows = batches_to_json(&page_batches)?;
|
||||||
|
let row_count = rows.as_array().map(|a| a.len()).unwrap_or(0);
|
||||||
|
|
||||||
|
Ok(PageResult {
|
||||||
|
query_id: query_id.to_string(),
|
||||||
|
page,
|
||||||
|
page_size: ps,
|
||||||
|
total_rows: result.total_rows,
|
||||||
|
total_pages: (result.total_rows + ps - 1) / ps,
|
||||||
|
rows,
|
||||||
|
row_count,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn default_page_size(&self) -> usize {
|
||||||
|
self.default_page_size
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Slice record batches to extract rows [start, end).
|
||||||
|
fn slice_batches(batches: &[RecordBatch], start: usize, end: usize) -> Vec<RecordBatch> {
|
||||||
|
let mut result = Vec::new();
|
||||||
|
let mut current_offset = 0;
|
||||||
|
|
||||||
|
for batch in batches {
|
||||||
|
let batch_end = current_offset + batch.num_rows();
|
||||||
|
|
||||||
|
if batch_end <= start {
|
||||||
|
current_offset = batch_end;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if current_offset >= end {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
let local_start = if start > current_offset { start - current_offset } else { 0 };
|
||||||
|
let local_end = if end < batch_end { end - current_offset } else { batch.num_rows() };
|
||||||
|
let length = local_end - local_start;
|
||||||
|
|
||||||
|
if length > 0 {
|
||||||
|
result.push(batch.slice(local_start, length));
|
||||||
|
}
|
||||||
|
|
||||||
|
current_offset = batch_end;
|
||||||
|
}
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
fn batches_to_json(batches: &[RecordBatch]) -> Result<serde_json::Value, String> {
|
||||||
|
if batches.is_empty() {
|
||||||
|
return Ok(serde_json::Value::Array(vec![]));
|
||||||
|
}
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
let mut writer = JsonWriter::<_, JsonArray>::new(&mut buf);
|
||||||
|
for batch in batches {
|
||||||
|
writer.write(batch).map_err(|e| format!("JSON write: {e}"))?;
|
||||||
|
}
|
||||||
|
writer.finish().map_err(|e| format!("JSON finish: {e}"))?;
|
||||||
|
drop(writer);
|
||||||
|
serde_json::from_slice(&buf).map_err(|e| format!("JSON parse: {e}"))
|
||||||
|
}
|
||||||
@@ -2,26 +2,47 @@ use arrow::array::RecordBatch;
|
|||||||
use arrow::json::writer::{JsonArray, Writer as JsonWriter};
|
use arrow::json::writer::{JsonArray, Writer as JsonWriter};
|
||||||
use axum::{
|
use axum::{
|
||||||
Json, Router,
|
Json, Router,
|
||||||
extract::State,
|
extract::{Path, Query, State},
|
||||||
http::StatusCode,
|
http::StatusCode,
|
||||||
response::IntoResponse,
|
response::IntoResponse,
|
||||||
routing::{get, post},
|
routing::{get, post},
|
||||||
};
|
};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
use crate::cache::CacheStats;
|
||||||
use crate::context::QueryEngine;
|
use crate::context::QueryEngine;
|
||||||
|
use crate::delta;
|
||||||
|
use crate::paged::ResultStore;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct QueryState {
|
||||||
|
pub engine: QueryEngine,
|
||||||
|
pub result_store: ResultStore,
|
||||||
|
}
|
||||||
|
|
||||||
pub fn router(engine: QueryEngine) -> Router {
|
pub fn router(engine: QueryEngine) -> Router {
|
||||||
|
let state = QueryState {
|
||||||
|
engine: engine.clone(),
|
||||||
|
result_store: ResultStore::new(100, 50), // 100 rows/page, keep 50 results
|
||||||
|
};
|
||||||
Router::new()
|
Router::new()
|
||||||
.route("/health", get(health))
|
.route("/health", get(health))
|
||||||
.route("/sql", post(execute_query))
|
.route("/sql", post(execute_query))
|
||||||
.with_state(engine)
|
.route("/paged", post(paged_query))
|
||||||
|
.route("/page/{query_id}/{page}", get(get_page))
|
||||||
|
.route("/cache/pin", post(pin_dataset))
|
||||||
|
.route("/cache/evict", post(evict_dataset))
|
||||||
|
.route("/cache/stats", get(cache_stats))
|
||||||
|
.route("/compact", post(compact_dataset))
|
||||||
|
.with_state(state)
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn health() -> &'static str {
|
async fn health() -> &'static str {
|
||||||
"queryd ok"
|
"queryd ok"
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// --- SQL Query ---
|
||||||
|
|
||||||
#[derive(Deserialize)]
|
#[derive(Deserialize)]
|
||||||
struct QueryRequest {
|
struct QueryRequest {
|
||||||
sql: String,
|
sql: String,
|
||||||
@@ -52,12 +73,12 @@ fn batches_to_json(batches: &[RecordBatch]) -> Result<serde_json::Value, String>
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn execute_query(
|
async fn execute_query(
|
||||||
State(engine): State<QueryEngine>,
|
State(state): State<QueryState>,
|
||||||
Json(req): Json<QueryRequest>,
|
Json(req): Json<QueryRequest>,
|
||||||
) -> impl IntoResponse {
|
) -> impl IntoResponse {
|
||||||
tracing::info!("executing query: {}", req.sql);
|
tracing::info!("executing query: {}", req.sql);
|
||||||
|
|
||||||
match engine.query(&req.sql).await {
|
match state.engine.query(&req.sql).await {
|
||||||
Ok(batches) => {
|
Ok(batches) => {
|
||||||
if batches.is_empty() {
|
if batches.is_empty() {
|
||||||
return Ok(Json(QueryResponse {
|
return Ok(Json(QueryResponse {
|
||||||
@@ -87,3 +108,107 @@ async fn execute_query(
|
|||||||
Err(e) => Err((StatusCode::BAD_REQUEST, e)),
|
Err(e) => Err((StatusCode::BAD_REQUEST, e)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// --- Paged Queries (large result sets) ---
|
||||||
|
|
||||||
|
async fn paged_query(
|
||||||
|
State(state): State<QueryState>,
|
||||||
|
Json(req): Json<QueryRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
tracing::info!("paged query: {}", req.sql);
|
||||||
|
match state.result_store.execute_and_store(&state.engine, &req.sql).await {
|
||||||
|
Ok(handle) => Ok(Json(handle)),
|
||||||
|
Err(e) => Err((StatusCode::BAD_REQUEST, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct PageQuery {
|
||||||
|
size: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_page(
|
||||||
|
State(state): State<QueryState>,
|
||||||
|
Path((query_id, page)): Path<(String, usize)>,
|
||||||
|
Query(q): Query<PageQuery>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match state.result_store.get_page(&query_id, page, q.size).await {
|
||||||
|
Ok(result) => Ok(Json(result)),
|
||||||
|
Err(e) => Err((StatusCode::NOT_FOUND, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Cache Management ---
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct CacheRequest {
|
||||||
|
dataset: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn pin_dataset(
|
||||||
|
State(state): State<QueryState>,
|
||||||
|
Json(req): Json<CacheRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match state.engine.pin_dataset(&req.dataset).await {
|
||||||
|
Ok(()) => Ok((StatusCode::OK, format!("pinned: {}", req.dataset))),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn evict_dataset(
|
||||||
|
State(state): State<QueryState>,
|
||||||
|
Json(req): Json<CacheRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
if state.engine.cache().evict(&req.dataset).await {
|
||||||
|
(StatusCode::OK, format!("evicted: {}", req.dataset))
|
||||||
|
} else {
|
||||||
|
(StatusCode::NOT_FOUND, format!("not cached: {}", req.dataset))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn cache_stats(State(state): State<QueryState>) -> impl IntoResponse {
|
||||||
|
let stats = state.engine.cache().stats().await;
|
||||||
|
Json(stats)
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Compaction ---
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct CompactRequest {
|
||||||
|
dataset: String,
|
||||||
|
base_key: String,
|
||||||
|
primary_key: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn compact_dataset(
|
||||||
|
State(state): State<QueryState>,
|
||||||
|
Json(req): Json<CompactRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
// Phase E: pull tombstones for this dataset and let compact physically
|
||||||
|
// drop those rows. After a successful rewrite, clear the tombstone log
|
||||||
|
// — the rows are gone from disk, the tombstones have done their job.
|
||||||
|
let tombstones = state
|
||||||
|
.engine
|
||||||
|
.catalog()
|
||||||
|
.list_tombstones(&req.dataset)
|
||||||
|
.await
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
match delta::compact(
|
||||||
|
state.engine.store(),
|
||||||
|
&req.dataset,
|
||||||
|
&req.base_key,
|
||||||
|
req.primary_key.as_deref(),
|
||||||
|
&tombstones,
|
||||||
|
).await {
|
||||||
|
Ok(result) => {
|
||||||
|
if result.rows_dropped_by_tombstones > 0 {
|
||||||
|
if let Err(e) = state.engine.catalog().tombstones().clear(&req.dataset).await {
|
||||||
|
tracing::warn!("post-compact tombstone clear failed: {e}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(Json(result))
|
||||||
|
}
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,272 @@
|
|||||||
|
/// Agent workspaces — named overlays for contract/search-specific work.
|
||||||
|
/// Each workspace tracks an agent's activity on a specific contract or search,
|
||||||
|
/// with daily/weekly/monthly tiers and instant handoff capability.
|
||||||
|
|
||||||
|
use arrow::array::{ArrayRef, RecordBatch, StringArray, Int64Array};
|
||||||
|
use arrow::datatypes::{DataType, Field, Schema};
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
use crate::delta;
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
|
||||||
|
/// Retention tier for workspace data.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||||
|
#[serde(rename_all = "lowercase")]
|
||||||
|
pub enum Tier {
|
||||||
|
Daily, // expires end of day, active search scratch
|
||||||
|
Weekly, // expires end of week, active contract work
|
||||||
|
Monthly, // expires end of month, contract lifecycle
|
||||||
|
Pinned, // never expires, manually managed
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A saved query/filter within a workspace.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct SavedSearch {
|
||||||
|
pub name: String,
|
||||||
|
pub sql: String,
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A shortlisted candidate or record.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ShortlistEntry {
|
||||||
|
pub dataset: String,
|
||||||
|
pub record_id: String,
|
||||||
|
pub notes: String,
|
||||||
|
pub added_at: DateTime<Utc>,
|
||||||
|
pub added_by: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Activity log entry.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ActivityEntry {
|
||||||
|
pub action: String, // "search", "shortlist", "call", "email", "update", "ingest"
|
||||||
|
pub detail: String,
|
||||||
|
pub timestamp: DateTime<Utc>,
|
||||||
|
pub agent: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A workspace — an agent's working context for a contract or search.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct Workspace {
|
||||||
|
pub id: String,
|
||||||
|
pub name: String, // "Apex Corp .NET Developers - Chicago"
|
||||||
|
pub description: String,
|
||||||
|
pub tier: Tier,
|
||||||
|
pub owner: String, // current agent
|
||||||
|
pub previous_owners: Vec<HandoffRecord>,
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
pub updated_at: DateTime<Utc>,
|
||||||
|
|
||||||
|
// Work content
|
||||||
|
pub saved_searches: Vec<SavedSearch>,
|
||||||
|
pub shortlist: Vec<ShortlistEntry>,
|
||||||
|
pub activity: Vec<ActivityEntry>,
|
||||||
|
pub ingested_datasets: Vec<String>, // datasets this workspace created
|
||||||
|
pub delta_keys: Vec<String>, // delta files specific to this workspace
|
||||||
|
pub tags: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Record of a handoff between agents.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct HandoffRecord {
|
||||||
|
pub from_agent: String,
|
||||||
|
pub to_agent: String,
|
||||||
|
pub reason: String,
|
||||||
|
pub timestamp: DateTime<Utc>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Workspace manager — in-memory registry with persistence.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct WorkspaceManager {
|
||||||
|
workspaces: Arc<RwLock<HashMap<String, Workspace>>>,
|
||||||
|
store: Arc<dyn ObjectStore>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl WorkspaceManager {
|
||||||
|
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||||
|
Self {
|
||||||
|
workspaces: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
store,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Rebuild from persisted workspace files on startup.
|
||||||
|
pub async fn rebuild(&self) -> Result<usize, String> {
|
||||||
|
let keys = storaged::ops::list(&self.store, Some("workspaces/")).await?;
|
||||||
|
let mut ws_map = self.workspaces.write().await;
|
||||||
|
ws_map.clear();
|
||||||
|
|
||||||
|
for key in &keys {
|
||||||
|
if !key.ends_with(".json") { continue; }
|
||||||
|
let data = storaged::ops::get(&self.store, key).await?;
|
||||||
|
match serde_json::from_slice::<Workspace>(&data) {
|
||||||
|
Ok(ws) => { ws_map.insert(ws.id.clone(), ws); }
|
||||||
|
Err(e) => tracing::warn!("failed to load workspace {key}: {e}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let count = ws_map.len();
|
||||||
|
if count > 0 {
|
||||||
|
tracing::info!("loaded {count} workspaces");
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create a new workspace.
|
||||||
|
pub async fn create(&self, name: String, description: String, owner: String, tier: Tier) -> Result<Workspace, String> {
|
||||||
|
let now = Utc::now();
|
||||||
|
let id = format!("ws-{}", now.timestamp_millis());
|
||||||
|
|
||||||
|
let ws = Workspace {
|
||||||
|
id: id.clone(),
|
||||||
|
name,
|
||||||
|
description,
|
||||||
|
tier,
|
||||||
|
owner,
|
||||||
|
previous_owners: vec![],
|
||||||
|
created_at: now,
|
||||||
|
updated_at: now,
|
||||||
|
saved_searches: vec![],
|
||||||
|
shortlist: vec![],
|
||||||
|
activity: vec![],
|
||||||
|
ingested_datasets: vec![],
|
||||||
|
delta_keys: vec![],
|
||||||
|
tags: vec![],
|
||||||
|
};
|
||||||
|
|
||||||
|
self.persist(&ws).await?;
|
||||||
|
self.workspaces.write().await.insert(id.clone(), ws.clone());
|
||||||
|
tracing::info!("created workspace: {} ({})", ws.name, ws.id);
|
||||||
|
Ok(ws)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handoff workspace to another agent. Instant — no data copy.
|
||||||
|
pub async fn handoff(&self, workspace_id: &str, to_agent: &str, reason: &str) -> Result<Workspace, String> {
|
||||||
|
let mut ws_map = self.workspaces.write().await;
|
||||||
|
let ws = ws_map.get_mut(workspace_id)
|
||||||
|
.ok_or_else(|| format!("workspace not found: {workspace_id}"))?;
|
||||||
|
|
||||||
|
let record = HandoffRecord {
|
||||||
|
from_agent: ws.owner.clone(),
|
||||||
|
to_agent: to_agent.to_string(),
|
||||||
|
reason: reason.to_string(),
|
||||||
|
timestamp: Utc::now(),
|
||||||
|
};
|
||||||
|
|
||||||
|
ws.previous_owners.push(record);
|
||||||
|
ws.owner = to_agent.to_string();
|
||||||
|
ws.updated_at = Utc::now();
|
||||||
|
|
||||||
|
ws.activity.push(ActivityEntry {
|
||||||
|
action: "handoff".to_string(),
|
||||||
|
detail: format!("handed off to {} — {}", to_agent, reason),
|
||||||
|
timestamp: Utc::now(),
|
||||||
|
agent: to_agent.to_string(),
|
||||||
|
});
|
||||||
|
|
||||||
|
let ws_clone = ws.clone();
|
||||||
|
drop(ws_map);
|
||||||
|
self.persist(&ws_clone).await?;
|
||||||
|
|
||||||
|
tracing::info!("workspace '{}' handed off to {}", ws_clone.name, to_agent);
|
||||||
|
Ok(ws_clone)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Add a saved search to a workspace.
|
||||||
|
pub async fn add_search(&self, workspace_id: &str, name: String, sql: String, agent: &str) -> Result<(), String> {
|
||||||
|
let mut ws_map = self.workspaces.write().await;
|
||||||
|
let ws = ws_map.get_mut(workspace_id)
|
||||||
|
.ok_or_else(|| format!("workspace not found: {workspace_id}"))?;
|
||||||
|
|
||||||
|
ws.saved_searches.push(SavedSearch {
|
||||||
|
name: name.clone(),
|
||||||
|
sql,
|
||||||
|
created_at: Utc::now(),
|
||||||
|
});
|
||||||
|
ws.activity.push(ActivityEntry {
|
||||||
|
action: "search".into(),
|
||||||
|
detail: format!("saved search: {name}"),
|
||||||
|
timestamp: Utc::now(),
|
||||||
|
agent: agent.to_string(),
|
||||||
|
});
|
||||||
|
ws.updated_at = Utc::now();
|
||||||
|
|
||||||
|
let ws_clone = ws.clone();
|
||||||
|
drop(ws_map);
|
||||||
|
self.persist(&ws_clone).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Add a candidate/record to the shortlist.
|
||||||
|
pub async fn add_to_shortlist(&self, workspace_id: &str, dataset: String, record_id: String, notes: String, agent: &str) -> Result<(), String> {
|
||||||
|
let mut ws_map = self.workspaces.write().await;
|
||||||
|
let ws = ws_map.get_mut(workspace_id)
|
||||||
|
.ok_or_else(|| format!("workspace not found: {workspace_id}"))?;
|
||||||
|
|
||||||
|
ws.shortlist.push(ShortlistEntry {
|
||||||
|
dataset: dataset.clone(),
|
||||||
|
record_id: record_id.clone(),
|
||||||
|
notes,
|
||||||
|
added_at: Utc::now(),
|
||||||
|
added_by: agent.to_string(),
|
||||||
|
});
|
||||||
|
ws.activity.push(ActivityEntry {
|
||||||
|
action: "shortlist".into(),
|
||||||
|
detail: format!("added {record_id} from {dataset}"),
|
||||||
|
timestamp: Utc::now(),
|
||||||
|
agent: agent.to_string(),
|
||||||
|
});
|
||||||
|
ws.updated_at = Utc::now();
|
||||||
|
|
||||||
|
let ws_clone = ws.clone();
|
||||||
|
drop(ws_map);
|
||||||
|
self.persist(&ws_clone).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Log an activity.
|
||||||
|
pub async fn log_activity(&self, workspace_id: &str, action: String, detail: String, agent: &str) -> Result<(), String> {
|
||||||
|
let mut ws_map = self.workspaces.write().await;
|
||||||
|
let ws = ws_map.get_mut(workspace_id)
|
||||||
|
.ok_or_else(|| format!("workspace not found: {workspace_id}"))?;
|
||||||
|
|
||||||
|
ws.activity.push(ActivityEntry {
|
||||||
|
action,
|
||||||
|
detail,
|
||||||
|
timestamp: Utc::now(),
|
||||||
|
agent: agent.to_string(),
|
||||||
|
});
|
||||||
|
ws.updated_at = Utc::now();
|
||||||
|
|
||||||
|
let ws_clone = ws.clone();
|
||||||
|
drop(ws_map);
|
||||||
|
self.persist(&ws_clone).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get a workspace.
|
||||||
|
pub async fn get(&self, workspace_id: &str) -> Option<Workspace> {
|
||||||
|
self.workspaces.read().await.get(workspace_id).cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all workspaces, optionally filtered by owner or tier.
|
||||||
|
pub async fn list(&self, owner: Option<&str>, tier: Option<&Tier>) -> Vec<Workspace> {
|
||||||
|
let ws_map = self.workspaces.read().await;
|
||||||
|
ws_map.values()
|
||||||
|
.filter(|ws| {
|
||||||
|
owner.map_or(true, |o| ws.owner == o) &&
|
||||||
|
tier.map_or(true, |t| ws.tier == *t)
|
||||||
|
})
|
||||||
|
.cloned()
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Persist workspace to object storage.
|
||||||
|
async fn persist(&self, ws: &Workspace) -> Result<(), String> {
|
||||||
|
let key = format!("workspaces/{}.json", ws.id);
|
||||||
|
let json = serde_json::to_vec_pretty(ws).map_err(|e| e.to_string())?;
|
||||||
|
storaged::ops::put(&self.store, &key, json.into()).await
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,136 @@
|
|||||||
|
use axum::{
|
||||||
|
Json, Router,
|
||||||
|
extract::{Path, Query, State},
|
||||||
|
http::StatusCode,
|
||||||
|
response::IntoResponse,
|
||||||
|
routing::{get, post},
|
||||||
|
};
|
||||||
|
use serde::Deserialize;
|
||||||
|
|
||||||
|
use crate::workspace::{Tier, WorkspaceManager};
|
||||||
|
|
||||||
|
pub fn router(manager: WorkspaceManager) -> Router {
|
||||||
|
Router::new()
|
||||||
|
.route("/", get(list_workspaces))
|
||||||
|
.route("/create", post(create_workspace))
|
||||||
|
.route("/{id}", get(get_workspace))
|
||||||
|
.route("/{id}/handoff", post(handoff))
|
||||||
|
.route("/{id}/search", post(add_search))
|
||||||
|
.route("/{id}/shortlist", post(add_to_shortlist))
|
||||||
|
.route("/{id}/activity", post(log_activity))
|
||||||
|
.with_state(manager)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct ListQuery {
|
||||||
|
owner: Option<String>,
|
||||||
|
tier: Option<Tier>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_workspaces(
|
||||||
|
State(mgr): State<WorkspaceManager>,
|
||||||
|
Query(q): Query<ListQuery>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let workspaces = mgr.list(q.owner.as_deref(), q.tier.as_ref()).await;
|
||||||
|
Json(workspaces)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct CreateRequest {
|
||||||
|
name: String,
|
||||||
|
description: String,
|
||||||
|
owner: String,
|
||||||
|
tier: Tier,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn create_workspace(
|
||||||
|
State(mgr): State<WorkspaceManager>,
|
||||||
|
Json(req): Json<CreateRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match mgr.create(req.name, req.description, req.owner, req.tier).await {
|
||||||
|
Ok(ws) => Ok((StatusCode::CREATED, Json(ws))),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_workspace(
|
||||||
|
State(mgr): State<WorkspaceManager>,
|
||||||
|
Path(id): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match mgr.get(&id).await {
|
||||||
|
Some(ws) => Ok(Json(ws)),
|
||||||
|
None => Err((StatusCode::NOT_FOUND, format!("workspace not found: {id}"))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct HandoffRequest {
|
||||||
|
to_agent: String,
|
||||||
|
reason: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn handoff(
|
||||||
|
State(mgr): State<WorkspaceManager>,
|
||||||
|
Path(id): Path<String>,
|
||||||
|
Json(req): Json<HandoffRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match mgr.handoff(&id, &req.to_agent, &req.reason).await {
|
||||||
|
Ok(ws) => Ok(Json(ws)),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct SearchRequest {
|
||||||
|
name: String,
|
||||||
|
sql: String,
|
||||||
|
agent: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn add_search(
|
||||||
|
State(mgr): State<WorkspaceManager>,
|
||||||
|
Path(id): Path<String>,
|
||||||
|
Json(req): Json<SearchRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match mgr.add_search(&id, req.name, req.sql, &req.agent).await {
|
||||||
|
Ok(()) => Ok((StatusCode::OK, "search saved")),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct ShortlistRequest {
|
||||||
|
dataset: String,
|
||||||
|
record_id: String,
|
||||||
|
notes: String,
|
||||||
|
agent: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn add_to_shortlist(
|
||||||
|
State(mgr): State<WorkspaceManager>,
|
||||||
|
Path(id): Path<String>,
|
||||||
|
Json(req): Json<ShortlistRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match mgr.add_to_shortlist(&id, req.dataset, req.record_id, req.notes, &req.agent).await {
|
||||||
|
Ok(()) => Ok((StatusCode::OK, "added to shortlist")),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct ActivityRequest {
|
||||||
|
action: String,
|
||||||
|
detail: String,
|
||||||
|
agent: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn log_activity(
|
||||||
|
State(mgr): State<WorkspaceManager>,
|
||||||
|
Path(id): Path<String>,
|
||||||
|
Json(req): Json<ActivityRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match mgr.log_activity(&id, req.action, req.detail, &req.agent).await {
|
||||||
|
Ok(()) => Ok((StatusCode::OK, "activity logged")),
|
||||||
|
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -15,3 +15,5 @@ bytes = { workspace = true }
|
|||||||
sha2 = { workspace = true }
|
sha2 = { workspace = true }
|
||||||
toml = { workspace = true }
|
toml = { workspace = true }
|
||||||
tracing = { workspace = true }
|
tracing = { workspace = true }
|
||||||
|
tokio = { workspace = true }
|
||||||
|
async-trait = "0.1"
|
||||||
|
|||||||
@@ -16,8 +16,46 @@ pub struct Config {
|
|||||||
pub auth: AuthConfig,
|
pub auth: AuthConfig,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub observability: ObservabilityConfig,
|
pub observability: ObservabilityConfig,
|
||||||
|
#[serde(default)]
|
||||||
|
pub agent: AgentSettings,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Phase 16.2 — background autotune agent settings.
|
||||||
|
///
|
||||||
|
/// Duplicated from `vectord::agent::AgentConfig` because `shared` can't
|
||||||
|
/// depend on `vectord` (vectord already depends on shared). The gateway
|
||||||
|
/// copies these into the vectord config at startup.
|
||||||
|
#[derive(Debug, Clone, Deserialize)]
|
||||||
|
pub struct AgentSettings {
|
||||||
|
#[serde(default)]
|
||||||
|
pub enabled: bool,
|
||||||
|
#[serde(default = "default_cycle_interval_secs")]
|
||||||
|
pub cycle_interval_secs: u64,
|
||||||
|
#[serde(default = "default_cooldown_secs")]
|
||||||
|
pub cooldown_between_trials_secs: u64,
|
||||||
|
#[serde(default = "default_min_recall")]
|
||||||
|
pub min_recall: f32,
|
||||||
|
#[serde(default = "default_max_trials_per_hour")]
|
||||||
|
pub max_trials_per_hour: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for AgentSettings {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
enabled: false,
|
||||||
|
cycle_interval_secs: default_cycle_interval_secs(),
|
||||||
|
cooldown_between_trials_secs: default_cooldown_secs(),
|
||||||
|
min_recall: default_min_recall(),
|
||||||
|
max_trials_per_hour: default_max_trials_per_hour(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_cycle_interval_secs() -> u64 { 60 }
|
||||||
|
fn default_cooldown_secs() -> u64 { 30 }
|
||||||
|
fn default_min_recall() -> f32 { 0.9 }
|
||||||
|
fn default_max_trials_per_hour() -> u32 { 30 }
|
||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize)]
|
#[derive(Debug, Clone, Deserialize)]
|
||||||
pub struct GatewayConfig {
|
pub struct GatewayConfig {
|
||||||
#[serde(default = "default_host")]
|
#[serde(default = "default_host")]
|
||||||
@@ -28,8 +66,39 @@ pub struct GatewayConfig {
|
|||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize)]
|
#[derive(Debug, Clone, Deserialize)]
|
||||||
pub struct StorageConfig {
|
pub struct StorageConfig {
|
||||||
|
/// Legacy single-backend root. If `buckets` is empty, this is used to
|
||||||
|
/// create an implicit `primary` bucket at this path — preserves the
|
||||||
|
/// pre-federation config shape.
|
||||||
#[serde(default = "default_storage_root")]
|
#[serde(default = "default_storage_root")]
|
||||||
pub root: String,
|
pub root: String,
|
||||||
|
|
||||||
|
/// Where profile buckets are rooted when auto-provisioned.
|
||||||
|
#[serde(default = "default_profile_root")]
|
||||||
|
pub profile_root: String,
|
||||||
|
|
||||||
|
/// Name of the bucket used for read fallback when a target bucket is
|
||||||
|
/// unreachable. If `None`, no fallback — reads fail hard.
|
||||||
|
#[serde(default)]
|
||||||
|
pub rescue_bucket: Option<String>,
|
||||||
|
|
||||||
|
/// Explicitly configured buckets. Empty = backward-compat single-bucket
|
||||||
|
/// mode driven by `root`.
|
||||||
|
#[serde(default)]
|
||||||
|
pub buckets: Vec<BucketConfig>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Deserialize)]
|
||||||
|
pub struct BucketConfig {
|
||||||
|
pub name: String,
|
||||||
|
pub backend: String, // "local" | "s3"
|
||||||
|
/// Local filesystem root (for backend = "local")
|
||||||
|
pub root: Option<String>,
|
||||||
|
/// S3 bucket name (for backend = "s3")
|
||||||
|
pub bucket: Option<String>,
|
||||||
|
pub region: Option<String>,
|
||||||
|
pub endpoint: Option<String>,
|
||||||
|
/// Handle for the secrets provider — never the literal credential.
|
||||||
|
pub secret_ref: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize, Default)]
|
#[derive(Debug, Clone, Deserialize, Default)]
|
||||||
@@ -78,6 +147,7 @@ pub struct ObservabilityConfig {
|
|||||||
fn default_host() -> String { "0.0.0.0".to_string() }
|
fn default_host() -> String { "0.0.0.0".to_string() }
|
||||||
fn default_gateway_port() -> u16 { 3100 }
|
fn default_gateway_port() -> u16 { 3100 }
|
||||||
fn default_storage_root() -> String { "./data".to_string() }
|
fn default_storage_root() -> String { "./data".to_string() }
|
||||||
|
fn default_profile_root() -> String { "./data/_profiles".to_string() }
|
||||||
fn default_manifest_prefix() -> String { "_catalog/manifests".to_string() }
|
fn default_manifest_prefix() -> String { "_catalog/manifests".to_string() }
|
||||||
fn default_sidecar_url() -> String { "http://localhost:3200".to_string() }
|
fn default_sidecar_url() -> String { "http://localhost:3200".to_string() }
|
||||||
fn default_embed_model() -> String { "nomic-embed-text".to_string() }
|
fn default_embed_model() -> String { "nomic-embed-text".to_string() }
|
||||||
@@ -115,13 +185,19 @@ impl Default for Config {
|
|||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
Self {
|
Self {
|
||||||
gateway: GatewayConfig { host: default_host(), port: default_gateway_port() },
|
gateway: GatewayConfig { host: default_host(), port: default_gateway_port() },
|
||||||
storage: StorageConfig { root: default_storage_root() },
|
storage: StorageConfig {
|
||||||
|
root: default_storage_root(),
|
||||||
|
profile_root: default_profile_root(),
|
||||||
|
rescue_bucket: None,
|
||||||
|
buckets: Vec::new(),
|
||||||
|
},
|
||||||
catalog: CatalogConfig::default(),
|
catalog: CatalogConfig::default(),
|
||||||
query: QueryConfig::default(),
|
query: QueryConfig::default(),
|
||||||
sidecar: SidecarConfig { url: default_sidecar_url() },
|
sidecar: SidecarConfig { url: default_sidecar_url() },
|
||||||
ai: AiConfig::default(),
|
ai: AiConfig::default(),
|
||||||
auth: AuthConfig::default(),
|
auth: AuthConfig::default(),
|
||||||
observability: ObservabilityConfig::default(),
|
observability: ObservabilityConfig::default(),
|
||||||
|
agent: AgentSettings::default(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,3 +2,5 @@ pub mod types;
|
|||||||
pub mod errors;
|
pub mod errors;
|
||||||
pub mod arrow_helpers;
|
pub mod arrow_helpers;
|
||||||
pub mod config;
|
pub mod config;
|
||||||
|
pub mod pii;
|
||||||
|
pub mod secrets;
|
||||||
|
|||||||
@@ -0,0 +1,116 @@
|
|||||||
|
/// Auto-detect PII columns by name patterns.
|
||||||
|
/// Conservative: flags likely PII, doesn't miss obvious cases.
|
||||||
|
|
||||||
|
use crate::types::Sensitivity;
|
||||||
|
|
||||||
|
/// Check if a column name suggests PII content.
|
||||||
|
pub fn detect_sensitivity(column_name: &str) -> Option<Sensitivity> {
|
||||||
|
let lower = column_name.to_lowercase();
|
||||||
|
|
||||||
|
// Direct PII identifiers
|
||||||
|
if matches!(lower.as_str(),
|
||||||
|
"ssn" | "social_security" | "social_security_number" |
|
||||||
|
"sin" | "national_id" | "passport" | "passport_number" |
|
||||||
|
"drivers_license" | "driver_license" | "dl_number"
|
||||||
|
) {
|
||||||
|
return Some(Sensitivity::Pii);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Names
|
||||||
|
if lower.contains("first_name") || lower.contains("last_name") ||
|
||||||
|
lower.contains("full_name") || lower.contains("middle_name") ||
|
||||||
|
lower == "name" || lower == "fname" || lower == "lname" {
|
||||||
|
return Some(Sensitivity::Pii);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Contact info
|
||||||
|
if lower.contains("email") || lower.contains("e_mail") ||
|
||||||
|
lower.contains("phone") || lower.contains("mobile") || lower.contains("cell") ||
|
||||||
|
lower.contains("fax") || lower == "tel" || lower == "telephone" {
|
||||||
|
return Some(Sensitivity::Pii);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Address
|
||||||
|
if lower.contains("address") || lower.contains("street") ||
|
||||||
|
(lower.contains("zip") && !lower.contains("unzip")) ||
|
||||||
|
lower == "postal_code" || lower == "postcode" {
|
||||||
|
return Some(Sensitivity::Pii);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Financial
|
||||||
|
if lower.contains("salary") || lower.contains("wage") ||
|
||||||
|
lower.contains("pay_rate") || lower.contains("bill_rate") ||
|
||||||
|
lower.contains("compensation") || lower.contains("revenue") ||
|
||||||
|
lower.contains("bank_account") || lower.contains("routing_number") ||
|
||||||
|
lower.contains("credit_card") {
|
||||||
|
return Some(Sensitivity::Financial);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Health
|
||||||
|
if lower.contains("diagnosis") || lower.contains("medication") ||
|
||||||
|
lower.contains("medical") || lower.contains("health") ||
|
||||||
|
lower.contains("patient_id") || lower.contains("mrn") {
|
||||||
|
return Some(Sensitivity::Phi);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Date of birth
|
||||||
|
if lower == "dob" || lower == "date_of_birth" || lower == "birthdate" || lower == "birth_date" {
|
||||||
|
return Some(Sensitivity::Pii);
|
||||||
|
}
|
||||||
|
|
||||||
|
None
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Classify all columns and return the highest sensitivity found.
|
||||||
|
pub fn detect_dataset_sensitivity(column_names: &[&str]) -> Option<Sensitivity> {
|
||||||
|
let mut highest: Option<Sensitivity> = None;
|
||||||
|
|
||||||
|
for name in column_names {
|
||||||
|
if let Some(sens) = detect_sensitivity(name) {
|
||||||
|
highest = Some(match (&highest, &sens) {
|
||||||
|
(None, s) => s.clone(),
|
||||||
|
(Some(Sensitivity::Public), s) => s.clone(),
|
||||||
|
(Some(Sensitivity::Internal), s) if !matches!(s, Sensitivity::Public | Sensitivity::Internal) => s.clone(),
|
||||||
|
(Some(Sensitivity::Financial), Sensitivity::Pii | Sensitivity::Phi) => sens.clone(),
|
||||||
|
(Some(Sensitivity::Pii), Sensitivity::Phi) => Sensitivity::Phi,
|
||||||
|
(existing, _) => existing.clone().unwrap(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
highest
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn detects_email_as_pii() {
|
||||||
|
assert_eq!(detect_sensitivity("email"), Some(Sensitivity::Pii));
|
||||||
|
assert_eq!(detect_sensitivity("contact_email"), Some(Sensitivity::Pii));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn detects_salary_as_financial() {
|
||||||
|
assert_eq!(detect_sensitivity("salary"), Some(Sensitivity::Financial));
|
||||||
|
assert_eq!(detect_sensitivity("bill_rate"), Some(Sensitivity::Financial));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn detects_ssn() {
|
||||||
|
assert_eq!(detect_sensitivity("ssn"), Some(Sensitivity::Pii));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn non_sensitive_returns_none() {
|
||||||
|
assert_eq!(detect_sensitivity("status"), None);
|
||||||
|
assert_eq!(detect_sensitivity("created_at"), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn dataset_sensitivity_picks_highest() {
|
||||||
|
let cols = vec!["id", "name", "email", "status"];
|
||||||
|
assert_eq!(detect_dataset_sensitivity(&cols), Some(Sensitivity::Pii));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,144 @@
|
|||||||
|
/// Secrets layer for bucket credentials.
|
||||||
|
///
|
||||||
|
/// Credentials never appear in `lakehouse.toml`. Each bucket entry has an
|
||||||
|
/// opaque `secret_ref` handle, resolved at startup by a `SecretsProvider`.
|
||||||
|
///
|
||||||
|
/// MVP ships `FileSecretsProvider` — reads a TOML file at a path given by
|
||||||
|
/// `LAKEHOUSE_SECRETS` env var, defaulting to `/etc/lakehouse/secrets.toml`.
|
||||||
|
/// The file should be root-owned, mode 0600, never in git.
|
||||||
|
///
|
||||||
|
/// Future providers (Vault, SOPS, OS keyring) implement the same trait.
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde::Deserialize;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct BucketCredentials {
|
||||||
|
pub access_key: String,
|
||||||
|
pub secret_key: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
// Redact credentials when debug-printing — NEVER let these end up in logs.
|
||||||
|
impl BucketCredentials {
|
||||||
|
pub fn redacted(&self) -> String {
|
||||||
|
let prefix: String = self.access_key.chars().take(4).collect();
|
||||||
|
format!("access_key={}... (redacted)", prefix)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
pub trait SecretsProvider: Send + Sync {
|
||||||
|
/// Resolve a handle to credentials. Returns Err if the handle is unknown.
|
||||||
|
async fn resolve(&self, handle: &str) -> Result<BucketCredentials, String>;
|
||||||
|
|
||||||
|
/// List known handles (for diagnostics only, never values).
|
||||||
|
async fn list_handles(&self) -> Vec<String>;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct RawSecretEntry {
|
||||||
|
access_key: String,
|
||||||
|
secret_key: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// File-backed secrets provider.
|
||||||
|
pub struct FileSecretsProvider {
|
||||||
|
path: PathBuf,
|
||||||
|
cache: tokio::sync::RwLock<HashMap<String, BucketCredentials>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl FileSecretsProvider {
|
||||||
|
pub fn new(path: impl Into<PathBuf>) -> Self {
|
||||||
|
Self {
|
||||||
|
path: path.into(),
|
||||||
|
cache: tokio::sync::RwLock::new(HashMap::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Default path: env var `LAKEHOUSE_SECRETS` or `/etc/lakehouse/secrets.toml`.
|
||||||
|
pub fn from_env() -> Self {
|
||||||
|
let path = std::env::var("LAKEHOUSE_SECRETS")
|
||||||
|
.unwrap_or_else(|_| "/etc/lakehouse/secrets.toml".to_string());
|
||||||
|
Self::new(path)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Arc wrapper for convenient injection.
|
||||||
|
pub fn shared(path: impl Into<PathBuf>) -> Arc<dyn SecretsProvider> {
|
||||||
|
Arc::new(Self::new(path))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Load and populate the cache. Call once at startup.
|
||||||
|
pub async fn load(&self) -> Result<usize, String> {
|
||||||
|
if !self.path.exists() {
|
||||||
|
tracing::info!(
|
||||||
|
"secrets file {} not present — running with zero credentials configured",
|
||||||
|
self.path.display()
|
||||||
|
);
|
||||||
|
return Ok(0);
|
||||||
|
}
|
||||||
|
check_permissions(&self.path)?;
|
||||||
|
let content = std::fs::read_to_string(&self.path)
|
||||||
|
.map_err(|e| format!("read secrets file: {e}"))?;
|
||||||
|
let raw: HashMap<String, RawSecretEntry> = toml::from_str(&content)
|
||||||
|
.map_err(|e| format!("parse secrets file: {e}"))?;
|
||||||
|
|
||||||
|
let mut cache = self.cache.write().await;
|
||||||
|
for (handle, entry) in raw {
|
||||||
|
cache.insert(handle.clone(), BucketCredentials {
|
||||||
|
access_key: entry.access_key,
|
||||||
|
secret_key: entry.secret_key,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
let count = cache.len();
|
||||||
|
tracing::info!("secrets: loaded {count} handles from {}", self.path.display());
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl SecretsProvider for FileSecretsProvider {
|
||||||
|
async fn resolve(&self, handle: &str) -> Result<BucketCredentials, String> {
|
||||||
|
let cache = self.cache.read().await;
|
||||||
|
cache
|
||||||
|
.get(handle)
|
||||||
|
.cloned()
|
||||||
|
.ok_or_else(|| format!("secret handle not found: {handle}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_handles(&self) -> Vec<String> {
|
||||||
|
self.cache.read().await.keys().cloned().collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Reject world-readable secrets files. Fail closed.
|
||||||
|
fn check_permissions(path: &Path) -> Result<(), String> {
|
||||||
|
#[cfg(unix)]
|
||||||
|
{
|
||||||
|
use std::os::unix::fs::PermissionsExt;
|
||||||
|
let meta = std::fs::metadata(path)
|
||||||
|
.map_err(|e| format!("stat secrets file: {e}"))?;
|
||||||
|
let mode = meta.permissions().mode() & 0o777;
|
||||||
|
if mode & 0o044 != 0 {
|
||||||
|
return Err(format!(
|
||||||
|
"secrets file {} is group/world-readable (mode {:o}); chmod 600 to proceed",
|
||||||
|
path.display(), mode
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let _ = path;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A zero-handles provider for tests / purely-local setups.
|
||||||
|
pub struct EmptySecretsProvider;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl SecretsProvider for EmptySecretsProvider {
|
||||||
|
async fn resolve(&self, handle: &str) -> Result<BucketCredentials, String> {
|
||||||
|
Err(format!("no secrets provider configured (wanted handle '{handle}')"))
|
||||||
|
}
|
||||||
|
async fn list_handles(&self) -> Vec<String> { Vec::new() }
|
||||||
|
}
|
||||||
+311
-1
@@ -30,7 +30,50 @@ pub struct ObjectRef {
|
|||||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||||
pub struct SchemaFingerprint(pub String);
|
pub struct SchemaFingerprint(pub String);
|
||||||
|
|
||||||
/// Dataset manifest — the catalog's view of a dataset.
|
/// Sensitivity classification for a field or dataset.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||||
|
#[serde(rename_all = "lowercase")]
|
||||||
|
pub enum Sensitivity {
|
||||||
|
Public, // anyone can see
|
||||||
|
Internal, // employees only
|
||||||
|
Pii, // personally identifiable (name, email, phone, SSN)
|
||||||
|
Phi, // protected health info (HIPAA)
|
||||||
|
Financial, // compensation, billing, revenue
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Column-level metadata.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ColumnMeta {
|
||||||
|
pub name: String,
|
||||||
|
pub data_type: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub sensitivity: Option<Sensitivity>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub description: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub is_pii: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Data lineage — where this dataset came from.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct Lineage {
|
||||||
|
pub source_system: String, // "csv_upload", "mysql_import", "api_ingest"
|
||||||
|
pub source_file: String, // original filename or URI
|
||||||
|
pub ingest_job: String, // job ID that created this dataset
|
||||||
|
pub ingest_timestamp: chrono::DateTime<chrono::Utc>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub parent_datasets: Vec<String>, // derived from these datasets
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Freshness contract — how often this data should be updated.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct FreshnessContract {
|
||||||
|
pub expected_frequency: String, // "daily", "weekly", "hourly", "manual"
|
||||||
|
pub stale_after_hours: u32, // alert if not updated within this many hours
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Dataset manifest — the catalog's complete view of a dataset.
|
||||||
|
/// All new fields use #[serde(default)] for backward compatibility with existing manifests.
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct DatasetManifest {
|
pub struct DatasetManifest {
|
||||||
pub id: DatasetId,
|
pub id: DatasetId,
|
||||||
@@ -39,4 +82,271 @@ pub struct DatasetManifest {
|
|||||||
pub objects: Vec<ObjectRef>,
|
pub objects: Vec<ObjectRef>,
|
||||||
pub created_at: chrono::DateTime<chrono::Utc>,
|
pub created_at: chrono::DateTime<chrono::Utc>,
|
||||||
pub updated_at: chrono::DateTime<chrono::Utc>,
|
pub updated_at: chrono::DateTime<chrono::Utc>,
|
||||||
|
|
||||||
|
// --- Rich metadata (Phase 10) ---
|
||||||
|
|
||||||
|
/// Human-readable description (auto-generated on ingest or manual)
|
||||||
|
#[serde(default)]
|
||||||
|
pub description: String,
|
||||||
|
|
||||||
|
/// Who owns this dataset
|
||||||
|
#[serde(default)]
|
||||||
|
pub owner: String,
|
||||||
|
|
||||||
|
/// Overall sensitivity classification
|
||||||
|
#[serde(default)]
|
||||||
|
pub sensitivity: Option<Sensitivity>,
|
||||||
|
|
||||||
|
/// Column-level metadata including sensitivity and descriptions
|
||||||
|
#[serde(default)]
|
||||||
|
pub columns: Vec<ColumnMeta>,
|
||||||
|
|
||||||
|
/// Where this data came from
|
||||||
|
#[serde(default)]
|
||||||
|
pub lineage: Option<Lineage>,
|
||||||
|
|
||||||
|
/// How fresh the data should be
|
||||||
|
#[serde(default)]
|
||||||
|
pub freshness: Option<FreshnessContract>,
|
||||||
|
|
||||||
|
/// Free-form tags for discovery
|
||||||
|
#[serde(default)]
|
||||||
|
pub tags: Vec<String>,
|
||||||
|
|
||||||
|
/// Row count (updated on ingest/compact)
|
||||||
|
#[serde(default)]
|
||||||
|
pub row_count: Option<u64>,
|
||||||
|
|
||||||
|
// --- Embedding freshness (Phase C) ---
|
||||||
|
|
||||||
|
/// When the attached vector index was last refreshed. `None` means this
|
||||||
|
/// dataset has no vector index yet, or has never been embedded.
|
||||||
|
#[serde(default)]
|
||||||
|
pub last_embedded_at: Option<chrono::DateTime<chrono::Utc>>,
|
||||||
|
|
||||||
|
/// When data was written that hasn't yet been embedded. `Some(t)` means
|
||||||
|
/// the vector index is out-of-date as of timestamp `t`. Cleared on refresh.
|
||||||
|
#[serde(default)]
|
||||||
|
pub embedding_stale_since: Option<chrono::DateTime<chrono::Utc>>,
|
||||||
|
|
||||||
|
/// How this dataset wants stale embeddings handled.
|
||||||
|
#[serde(default)]
|
||||||
|
pub embedding_refresh_policy: Option<RefreshPolicy>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Controls what happens when new data is written to a dataset with an
|
||||||
|
/// attached vector index.
|
||||||
|
///
|
||||||
|
/// - `Manual` (default): data writes set `embedding_stale_since`; nothing
|
||||||
|
/// embeds until an operator or agent calls `/vectors/refresh/{dataset}`.
|
||||||
|
/// - `OnAppend`: ingest fires a background refresh immediately after writing.
|
||||||
|
/// Suitable for datasets where vector freshness matters more than ingest
|
||||||
|
/// latency.
|
||||||
|
/// - `Scheduled(cron)`: a timer or external scheduler triggers refresh at
|
||||||
|
/// the named cadence. The scheduler itself is not in this ADR scope —
|
||||||
|
/// this just declares the intent so operators can see which policy a
|
||||||
|
/// dataset expects.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||||
|
pub enum RefreshPolicy {
|
||||||
|
Manual,
|
||||||
|
OnAppend,
|
||||||
|
Scheduled { cron: String },
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for RefreshPolicy {
|
||||||
|
fn default() -> Self { Self::Manual }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// AI-safe view of a dataset (Phase D).
|
||||||
|
///
|
||||||
|
/// A view is a named projection over a base dataset that:
|
||||||
|
/// - whitelists the columns visible (everything else is invisible),
|
||||||
|
/// - optionally applies a row filter (a SQL WHERE clause fragment),
|
||||||
|
/// - optionally redacts column values (e.g. mask the last 4 of an SSN).
|
||||||
|
///
|
||||||
|
/// Queries that reference the view name see only the safe projection.
|
||||||
|
/// The base table remains separately accessible by its raw name — this
|
||||||
|
/// is a *capability surface*, not a data classification fence. Use
|
||||||
|
/// access control (Phase 13) on top to enforce who can query which view.
|
||||||
|
///
|
||||||
|
/// Stored as JSON at `_catalog/views/{name}.json` in the primary bucket
|
||||||
|
/// alongside dataset manifests. The view itself doesn't store data —
|
||||||
|
/// queryd materializes it as a DataFusion view at SessionContext build
|
||||||
|
/// time.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct AiView {
|
||||||
|
pub name: String,
|
||||||
|
/// Source dataset name. Must exist in the catalog.
|
||||||
|
pub base_dataset: String,
|
||||||
|
/// Whitelisted columns. Empty list is invalid (would expose nothing).
|
||||||
|
pub columns: Vec<String>,
|
||||||
|
/// Optional WHERE clause fragment (without the `WHERE` keyword), e.g.
|
||||||
|
/// `status != 'blocked' AND city = 'Chicago'`. SQL injection here is
|
||||||
|
/// the operator's problem — only trust this from authenticated admins.
|
||||||
|
#[serde(default)]
|
||||||
|
pub row_filter: Option<String>,
|
||||||
|
/// Per-column redactions: `{column_name -> Redaction}`.
|
||||||
|
#[serde(default)]
|
||||||
|
pub column_redactions: std::collections::HashMap<String, Redaction>,
|
||||||
|
pub created_at: chrono::DateTime<chrono::Utc>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub created_by: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub description: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Per-model configuration for HNSW builds. Declared inline so the PRD
|
||||||
|
/// invariant 9 ("every reader gets its own profile") can carry tuning
|
||||||
|
/// preferences without forcing a round-trip to vectord's trial system.
|
||||||
|
/// Mirrors `vectord::trial::HnswConfig` to avoid a cross-crate dep cycle.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct ProfileHnswConfig {
|
||||||
|
pub ef_construction: usize,
|
||||||
|
pub ef_search: usize,
|
||||||
|
#[serde(default)]
|
||||||
|
pub seed: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for ProfileHnswConfig {
|
||||||
|
/// Matches the production HNSW default from Phase 15 trial sweep
|
||||||
|
/// (see docs/ADR and vectord::trial::HnswConfig::default).
|
||||||
|
fn default() -> Self {
|
||||||
|
Self { ef_construction: 80, ef_search: 30, seed: None }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A named client of the substrate — a local LLM, an AI agent, or a
|
||||||
|
/// human role that should see a specific scoped view of the data.
|
||||||
|
/// Profiles are the first-class "every reader gets its own profile"
|
||||||
|
/// concept from PRD invariant 9.
|
||||||
|
///
|
||||||
|
/// Bound datasets can be raw catalog tables OR `AiView` names — either
|
||||||
|
/// works at query time since both register as DataFusion tables with
|
||||||
|
/// the same name. Mixing raw and views lets you compose access rules:
|
||||||
|
/// bind a model to `candidates_safe` (PII-redacted view) and `placements`
|
||||||
|
/// (raw table, no PII), and the model gets exactly what it needs.
|
||||||
|
///
|
||||||
|
/// Stored at `_catalog/profiles/{id}.json`. Activation (`/profile/{id}/
|
||||||
|
/// activate`) pre-loads the embedding cache and builds HNSW with the
|
||||||
|
/// profile's config — first search after activate is warm.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ModelProfile {
|
||||||
|
/// Unique id — slug form (alphanumeric + `-`, `_`). Used in URLs,
|
||||||
|
/// bucket names (`profile:{id}`), and trial journal prefixes.
|
||||||
|
pub id: String,
|
||||||
|
/// Ollama model tag (e.g. "qwen2.5:7b", "nomic-embed-text:latest").
|
||||||
|
/// Used when this profile is the "generate" or "embed" caller.
|
||||||
|
pub ollama_name: String,
|
||||||
|
/// Free-text description surfaced to operators. Who/what is this
|
||||||
|
/// profile for?
|
||||||
|
#[serde(default)]
|
||||||
|
pub description: String,
|
||||||
|
/// Names of datasets or AiViews this profile can read. A search
|
||||||
|
/// against an index whose `source` isn't in this list is refused.
|
||||||
|
pub bound_datasets: Vec<String>,
|
||||||
|
/// Per-profile HNSW config, used when activate builds the graph.
|
||||||
|
#[serde(default)]
|
||||||
|
pub hnsw_config: ProfileHnswConfig,
|
||||||
|
/// Which embedding model this profile expects. Defaults to
|
||||||
|
/// `nomic-embed-text` (the only one currently installed).
|
||||||
|
#[serde(default = "default_embed_model")]
|
||||||
|
pub embed_model: String,
|
||||||
|
pub created_at: chrono::DateTime<chrono::Utc>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub created_by: String,
|
||||||
|
/// Federation layer 2: which bucket this profile's artifacts
|
||||||
|
/// (trial journals, promotions for its indexes) live in. If set and
|
||||||
|
/// not yet registered, `POST /vectors/profile/{id}/activate` auto-
|
||||||
|
/// provisions a local bucket under `storage.profile_root` and
|
||||||
|
/// registers it before warming indexes. `None` = share the primary
|
||||||
|
/// bucket (pre-federation behavior).
|
||||||
|
#[serde(default)]
|
||||||
|
pub bucket: Option<String>,
|
||||||
|
/// ADR-019 hybrid: which storage backend serves this profile's
|
||||||
|
/// indexes. Default Parquet so pre-existing profiles unchanged.
|
||||||
|
/// New profiles can opt into Lance for append-heavy / random-access
|
||||||
|
/// / 5M+ vector workloads.
|
||||||
|
#[serde(default)]
|
||||||
|
pub vector_backend: VectorBackend,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_embed_model() -> String { "nomic-embed-text".to_string() }
|
||||||
|
|
||||||
|
/// Per-profile / per-index storage backend choice.
|
||||||
|
///
|
||||||
|
/// `Parquet` (default) — vectors as binary-blob in Parquet + in-RAM HNSW
|
||||||
|
/// for ANN. Best when:
|
||||||
|
/// - Index ≤ ~5M vectors (fits in RAM)
|
||||||
|
/// - Search latency dominates (HNSW p50 < 1ms at 100K)
|
||||||
|
/// - Workload is mostly read, not append
|
||||||
|
///
|
||||||
|
/// `Lance` — vectors as native FixedSizeList in a Lance dataset, IVF_PQ
|
||||||
|
/// on-disk index. Best when:
|
||||||
|
/// - Index will exceed RAM ceiling (5M+ vectors per index)
|
||||||
|
/// - Heavy random-row access for RAG (Lance: 311us; Parquet: ~35ms full
|
||||||
|
/// scan per fetch)
|
||||||
|
/// - Append-heavy ingest (Lance: 0.08s for 10K rows; Parquet: full
|
||||||
|
/// ~330MB rewrite per change)
|
||||||
|
/// - Need to rebuild index frequently (Lance IVF_PQ: 16s; HNSW: 230s
|
||||||
|
/// at the same scale)
|
||||||
|
///
|
||||||
|
/// Defaults to Parquet for backward compat with pre-ADR-019 indexes.
|
||||||
|
/// See `docs/ADR-019-vector-storage.md` for the full scorecard.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||||
|
#[serde(rename_all = "lowercase")]
|
||||||
|
pub enum VectorBackend {
|
||||||
|
Parquet,
|
||||||
|
Lance,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for VectorBackend {
|
||||||
|
fn default() -> Self { VectorBackend::Parquet }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Soft-delete marker (Phase E).
|
||||||
|
///
|
||||||
|
/// Tombstones live beside the dataset in `_catalog/tombstones/{dataset}/`
|
||||||
|
/// as append-log JSONL. Query-time filter in queryd reads all tombstones
|
||||||
|
/// for each dataset and wraps the base table in a DataFusion view that
|
||||||
|
/// excludes tombstoned rows. Physical deletion happens later (compaction),
|
||||||
|
/// so the row count immediately reflects the delete but data is still on
|
||||||
|
/// disk until compact runs. That's deliberate — it gives a reversal
|
||||||
|
/// window and keeps the event journal audit trail intact.
|
||||||
|
///
|
||||||
|
/// All tombstones for a given dataset must use the same `row_key_column`
|
||||||
|
/// (enforced on write); otherwise the query-time filter can't be built
|
||||||
|
/// as a single WHERE clause.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct Tombstone {
|
||||||
|
pub dataset: String,
|
||||||
|
/// Column name that identifies the row (e.g. "candidate_id").
|
||||||
|
pub row_key_column: String,
|
||||||
|
/// Value of that column for the tombstoned row.
|
||||||
|
pub row_key_value: String,
|
||||||
|
pub deleted_at: chrono::DateTime<chrono::Utc>,
|
||||||
|
/// Human / system actor responsible for the delete (audit).
|
||||||
|
#[serde(default)]
|
||||||
|
pub actor: String,
|
||||||
|
/// Why (e.g. "GDPR request #1234", "user-requested erasure").
|
||||||
|
#[serde(default)]
|
||||||
|
pub reason: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// How a column's values should be transformed before being returned.
|
||||||
|
/// `Mask` is the most common — keeps a few visible chars, replaces the
|
||||||
|
/// rest with `*`. `Hash` returns SHA-256 of the value for join keys you
|
||||||
|
/// want to preserve referential integrity on without exposing the value.
|
||||||
|
/// `Null` returns NULL unconditionally.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||||
|
pub enum Redaction {
|
||||||
|
Null,
|
||||||
|
Hash,
|
||||||
|
Mask {
|
||||||
|
/// Number of leading characters left visible.
|
||||||
|
keep_prefix: usize,
|
||||||
|
/// Number of trailing characters left visible.
|
||||||
|
keep_suffix: usize,
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -13,3 +13,4 @@ tracing = { workspace = true }
|
|||||||
object_store = { workspace = true }
|
object_store = { workspace = true }
|
||||||
bytes = { workspace = true }
|
bytes = { workspace = true }
|
||||||
futures = { workspace = true }
|
futures = { workspace = true }
|
||||||
|
chrono = { workspace = true }
|
||||||
|
|||||||
@@ -0,0 +1,239 @@
|
|||||||
|
/// Write-once batched append log.
|
||||||
|
///
|
||||||
|
/// Problem we're fixing: the error journal and HNSW trial journal both
|
||||||
|
/// previously did read-modify-write of their whole JSONL file on every
|
||||||
|
/// event. That's O(N²) cumulative work and generates huge churn at scale.
|
||||||
|
/// It's exactly the pattern llms3.com flags as the "small-file /
|
||||||
|
/// rewrite-amplification" anti-pattern.
|
||||||
|
///
|
||||||
|
/// This helper implements the pattern object storage actually wants:
|
||||||
|
///
|
||||||
|
/// - Events accumulate in an **in-memory buffer** (reads see them immediately).
|
||||||
|
/// - When the buffer hits a threshold, or `flush()` is called, the buffer is
|
||||||
|
/// written **as one new object** with a timestamp-sorted key.
|
||||||
|
/// - Existing objects are never rewritten.
|
||||||
|
/// - Reads enumerate all batch files, sort by key, and concat in order.
|
||||||
|
/// - An explicit `compact()` reads every batch file, writes one consolidated
|
||||||
|
/// file, and deletes the originals — the LSM-tree compaction idea applied
|
||||||
|
/// to small JSONL events.
|
||||||
|
///
|
||||||
|
/// Storage layout:
|
||||||
|
/// ```
|
||||||
|
/// {prefix}/
|
||||||
|
/// batch_0001776319628000123.jsonl
|
||||||
|
/// batch_0001776319745987654.jsonl
|
||||||
|
/// batch_compacted_00001776319800000000.jsonl (after compact())
|
||||||
|
/// ```
|
||||||
|
/// Key format: the zero-padded epoch microsecond of the write, so lexical
|
||||||
|
/// sort == chronological sort.
|
||||||
|
|
||||||
|
use bytes::Bytes;
|
||||||
|
use chrono::Utc;
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::Mutex;
|
||||||
|
|
||||||
|
use crate::ops;
|
||||||
|
|
||||||
|
const DEFAULT_FLUSH_THRESHOLD: usize = 32;
|
||||||
|
|
||||||
|
pub struct AppendLog {
|
||||||
|
store: Arc<dyn ObjectStore>,
|
||||||
|
prefix: String,
|
||||||
|
buffer: Mutex<Vec<Vec<u8>>>,
|
||||||
|
flush_threshold: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AppendLog {
|
||||||
|
/// Create a new append log rooted at `prefix` in the given object store.
|
||||||
|
/// Events auto-flush when buffer reaches `flush_threshold` (default 32).
|
||||||
|
pub fn new(store: Arc<dyn ObjectStore>, prefix: impl Into<String>) -> Self {
|
||||||
|
Self {
|
||||||
|
store,
|
||||||
|
prefix: prefix.into(),
|
||||||
|
buffer: Mutex::new(Vec::new()),
|
||||||
|
flush_threshold: DEFAULT_FLUSH_THRESHOLD,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn with_flush_threshold(mut self, threshold: usize) -> Self {
|
||||||
|
self.flush_threshold = threshold.max(1);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Add an event. The returned future completes either immediately
|
||||||
|
/// (buffered) or after a flush, depending on whether the buffer hit the
|
||||||
|
/// threshold. Callers don't need to care either way.
|
||||||
|
pub async fn append(&self, line: Vec<u8>) -> Result<(), String> {
|
||||||
|
let should_flush = {
|
||||||
|
let mut buf = self.buffer.lock().await;
|
||||||
|
buf.push(line);
|
||||||
|
buf.len() >= self.flush_threshold
|
||||||
|
};
|
||||||
|
if should_flush {
|
||||||
|
self.flush().await?;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Force-flush the in-memory buffer to object storage as a single new
|
||||||
|
/// batch file. Safe to call anytime (idempotent no-op when buffer empty).
|
||||||
|
pub async fn flush(&self) -> Result<(), String> {
|
||||||
|
let batch = {
|
||||||
|
let mut buf = self.buffer.lock().await;
|
||||||
|
if buf.is_empty() { return Ok(()); }
|
||||||
|
std::mem::take(&mut *buf)
|
||||||
|
};
|
||||||
|
|
||||||
|
let ts_us = Utc::now().timestamp_micros().max(0) as u128;
|
||||||
|
let key = format!("{}/batch_{:019}.jsonl", self.prefix, ts_us);
|
||||||
|
|
||||||
|
let mut body = Vec::with_capacity(batch.iter().map(|b| b.len() + 1).sum());
|
||||||
|
for line in batch {
|
||||||
|
body.extend_from_slice(&line);
|
||||||
|
if !body.ends_with(b"\n") { body.push(b'\n'); }
|
||||||
|
}
|
||||||
|
ops::put(&self.store, &key, Bytes::from(body)).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Read every event across all batch files + unflushed in-memory buffer.
|
||||||
|
/// Events are returned in chronological order.
|
||||||
|
pub async fn read_all(&self) -> Result<Vec<Vec<u8>>, String> {
|
||||||
|
let mut keys = self.list_batch_keys().await?;
|
||||||
|
keys.sort();
|
||||||
|
|
||||||
|
let mut out = Vec::new();
|
||||||
|
for key in keys {
|
||||||
|
let bytes = ops::get(&self.store, &key).await?;
|
||||||
|
for line in bytes.split(|b| *b == b'\n') {
|
||||||
|
if !line.is_empty() {
|
||||||
|
out.push(line.to_vec());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Include unflushed events so callers see the latest state
|
||||||
|
// whether or not someone ran flush() recently.
|
||||||
|
let buf = self.buffer.lock().await;
|
||||||
|
for line in buf.iter() {
|
||||||
|
out.push(line.clone());
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(out)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Consolidate all current batch files into one compacted file, then
|
||||||
|
/// delete the originals. Safe to call while appends are in flight:
|
||||||
|
/// new batches written during compaction get a higher timestamp and
|
||||||
|
/// survive. Fails closed — if anything goes wrong mid-delete, the
|
||||||
|
/// compacted file coexists with originals and next read sees duplicates
|
||||||
|
/// (which the dedup caller must handle) rather than data loss.
|
||||||
|
pub async fn compact(&self) -> Result<CompactStats, String> {
|
||||||
|
// Snapshot which files to compact BEFORE we write the new one.
|
||||||
|
let mut originals = self.list_batch_keys().await?;
|
||||||
|
originals.sort();
|
||||||
|
|
||||||
|
if originals.len() < 2 {
|
||||||
|
return Ok(CompactStats { merged_files: originals.len(), events: 0, new_key: None });
|
||||||
|
}
|
||||||
|
|
||||||
|
// Gather all existing events.
|
||||||
|
let mut events = Vec::new();
|
||||||
|
for key in &originals {
|
||||||
|
let bytes = ops::get(&self.store, key).await?;
|
||||||
|
for line in bytes.split(|b| *b == b'\n') {
|
||||||
|
if !line.is_empty() {
|
||||||
|
events.push(line.to_vec());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let total_events = events.len();
|
||||||
|
if total_events == 0 {
|
||||||
|
// Clean up empty files without writing a new one.
|
||||||
|
for key in &originals {
|
||||||
|
let _ = ops::delete(&self.store, key).await;
|
||||||
|
}
|
||||||
|
return Ok(CompactStats { merged_files: originals.len(), events: 0, new_key: None });
|
||||||
|
}
|
||||||
|
|
||||||
|
// Name the compacted file with the SAME `batch_{ts}` format so it
|
||||||
|
// sorts chronologically with future batches. Using a distinct prefix
|
||||||
|
// ("batch_compacted_") would break lex ordering: later `batch_N`
|
||||||
|
// files would sort BEFORE the compacted file because 'c' > digits.
|
||||||
|
// Timestamp = now, so any appends arriving during compaction (which
|
||||||
|
// get the current wall-clock time) sort AFTER this file.
|
||||||
|
let ts_us = Utc::now().timestamp_micros().max(0) as u128;
|
||||||
|
let new_key = format!("{}/batch_{:019}.jsonl", self.prefix, ts_us);
|
||||||
|
|
||||||
|
let mut body = Vec::new();
|
||||||
|
for line in &events {
|
||||||
|
body.extend_from_slice(line);
|
||||||
|
body.push(b'\n');
|
||||||
|
}
|
||||||
|
ops::put(&self.store, &new_key, Bytes::from(body)).await?;
|
||||||
|
|
||||||
|
// Only delete originals once the consolidated file is persisted.
|
||||||
|
let mut failures = 0;
|
||||||
|
for key in &originals {
|
||||||
|
if ops::delete(&self.store, key).await.is_err() {
|
||||||
|
failures += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if failures > 0 {
|
||||||
|
tracing::warn!(
|
||||||
|
"compact '{}': {} original files failed to delete — consolidated file {} has the data",
|
||||||
|
self.prefix, failures, new_key,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(CompactStats {
|
||||||
|
merged_files: originals.len(),
|
||||||
|
events: total_events,
|
||||||
|
new_key: Some(new_key),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// How many batch files exist on disk right now.
|
||||||
|
pub async fn file_count(&self) -> Result<usize, String> {
|
||||||
|
Ok(self.list_batch_keys().await?.len())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_batch_keys(&self) -> Result<Vec<String>, String> {
|
||||||
|
let prefix_with_slash = format!("{}/", self.prefix);
|
||||||
|
// list the prefix then filter for keys that match our naming scheme;
|
||||||
|
// unrelated files at the same prefix won't be touched.
|
||||||
|
let raw = ops::list(&self.store, Some(&prefix_with_slash)).await?;
|
||||||
|
Ok(raw
|
||||||
|
.into_iter()
|
||||||
|
.filter(|k| {
|
||||||
|
let basename = k.rsplit('/').next().unwrap_or(k);
|
||||||
|
basename.starts_with("batch_") && basename.ends_with(".jsonl")
|
||||||
|
})
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct CompactStats {
|
||||||
|
pub merged_files: usize,
|
||||||
|
pub events: usize,
|
||||||
|
pub new_key: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
// Proactively flush on drop, best-effort.
|
||||||
|
// We can't `.await` in Drop; we spawn the flush on the tokio runtime if one
|
||||||
|
// is available. If the runtime is already shutting down the flush is lost —
|
||||||
|
// acceptable for these observability journals, which are hints not records.
|
||||||
|
impl Drop for AppendLog {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
let buf_len = self.buffer.try_lock().map(|b| b.len()).unwrap_or(0);
|
||||||
|
if buf_len == 0 { return; }
|
||||||
|
// Can't spawn from sync Drop on every runtime shape; log + move on.
|
||||||
|
tracing::debug!(
|
||||||
|
"append_log '{}' dropping with {} unflushed events",
|
||||||
|
self.prefix, buf_len,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
@@ -0,0 +1,200 @@
|
|||||||
|
/// Bucket operation error journal.
|
||||||
|
///
|
||||||
|
/// Every bucket op failure and every rescue fallback lands here. Goal:
|
||||||
|
/// answering "is anything broken?" with one HTTP call.
|
||||||
|
///
|
||||||
|
/// Storage: batched write-once files under `_errors/bucket_errors/` in the
|
||||||
|
/// primary bucket. Uses the shared `AppendLog` helper so we never rewrite
|
||||||
|
/// existing files — see `append_log.rs` for the full pattern rationale.
|
||||||
|
///
|
||||||
|
/// In-memory ring buffer holds the last N events for fast response; on
|
||||||
|
/// startup `load_recent` hydrates it from all batch files.
|
||||||
|
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::VecDeque;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
use crate::append_log::{AppendLog, CompactStats};
|
||||||
|
|
||||||
|
const JOURNAL_PREFIX: &str = "_errors/bucket_errors";
|
||||||
|
const RING_CAPACITY: usize = 2000;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
#[serde(rename_all = "lowercase")]
|
||||||
|
pub enum BucketOp {
|
||||||
|
Read,
|
||||||
|
Write,
|
||||||
|
Delete,
|
||||||
|
List,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct BucketErrorEvent {
|
||||||
|
pub ts: DateTime<Utc>,
|
||||||
|
pub op: BucketOp,
|
||||||
|
pub target: String,
|
||||||
|
pub key: String,
|
||||||
|
pub error: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub rescued: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl BucketErrorEvent {
|
||||||
|
pub fn new_read(target: &str, key: &str, error: &str) -> Self {
|
||||||
|
Self { ts: Utc::now(), op: BucketOp::Read, target: target.into(), key: key.into(), error: error.into(), rescued: false }
|
||||||
|
}
|
||||||
|
pub fn new_write(target: &str, key: &str, error: &str) -> Self {
|
||||||
|
Self { ts: Utc::now(), op: BucketOp::Write, target: target.into(), key: key.into(), error: error.into(), rescued: false }
|
||||||
|
}
|
||||||
|
pub fn new_delete(target: &str, key: &str, error: &str) -> Self {
|
||||||
|
Self { ts: Utc::now(), op: BucketOp::Delete, target: target.into(), key: key.into(), error: error.into(), rescued: false }
|
||||||
|
}
|
||||||
|
pub fn new_list(target: &str, prefix: &str, error: &str) -> Self {
|
||||||
|
Self { ts: Utc::now(), op: BucketOp::List, target: target.into(), key: prefix.into(), error: error.into(), rescued: false }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct ErrorJournal {
|
||||||
|
log: Arc<AppendLog>,
|
||||||
|
ring: Arc<RwLock<VecDeque<BucketErrorEvent>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct HealthReport {
|
||||||
|
pub period_minutes: i64,
|
||||||
|
pub total_errors: usize,
|
||||||
|
pub per_bucket: std::collections::HashMap<String, usize>,
|
||||||
|
pub unhealthy_buckets: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ErrorJournal {
|
||||||
|
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||||
|
// Keep flush threshold lowish — operators checking /storage/errors
|
||||||
|
// after a recent incident want to see fresh rows even without an
|
||||||
|
// explicit flush.
|
||||||
|
let log = Arc::new(
|
||||||
|
AppendLog::new(store, JOURNAL_PREFIX).with_flush_threshold(8),
|
||||||
|
);
|
||||||
|
Self {
|
||||||
|
log,
|
||||||
|
ring: Arc::new(RwLock::new(VecDeque::with_capacity(RING_CAPACITY))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Hydrate the ring buffer from existing batch files. Call once at
|
||||||
|
/// startup. Tolerates malformed lines (skipped with a warning) and
|
||||||
|
/// missing files (returns 0).
|
||||||
|
pub async fn load_recent(&self) -> Result<usize, String> {
|
||||||
|
let lines = self.log.read_all().await.unwrap_or_default();
|
||||||
|
let mut ring = self.ring.write().await;
|
||||||
|
for line in lines {
|
||||||
|
match serde_json::from_slice::<BucketErrorEvent>(&line) {
|
||||||
|
Ok(ev) => {
|
||||||
|
if ring.len() >= RING_CAPACITY { ring.pop_front(); }
|
||||||
|
ring.push_back(ev);
|
||||||
|
}
|
||||||
|
Err(e) => tracing::warn!("error journal: skip malformed line ({e})"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(ring.len())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append an event. In-memory ring updated immediately; persistence
|
||||||
|
/// happens in batches via the underlying AppendLog.
|
||||||
|
pub async fn append(&self, event: BucketErrorEvent) {
|
||||||
|
{
|
||||||
|
let mut ring = self.ring.write().await;
|
||||||
|
if ring.len() >= RING_CAPACITY { ring.pop_front(); }
|
||||||
|
ring.push_back(event.clone());
|
||||||
|
}
|
||||||
|
|
||||||
|
match serde_json::to_vec(&event) {
|
||||||
|
Ok(line) => {
|
||||||
|
if let Err(e) = self.log.append(line).await {
|
||||||
|
tracing::error!("error journal persist failed: {e}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => tracing::error!("error journal serialize failed: {e}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mark the most recent matching in-memory event as rescued.
|
||||||
|
/// Only updates the ring buffer — the JSONL line already persisted
|
||||||
|
/// records the failure fact; rescue status travels in response headers.
|
||||||
|
pub async fn mark_rescued_last(&self, target: &str, key: &str) {
|
||||||
|
let mut ring = self.ring.write().await;
|
||||||
|
for ev in ring.iter_mut().rev() {
|
||||||
|
if matches!(ev.op, BucketOp::Read) && ev.target == target && ev.key == key && !ev.rescued {
|
||||||
|
ev.rescued = true;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn recent(&self, limit: usize) -> Vec<BucketErrorEvent> {
|
||||||
|
let ring = self.ring.read().await;
|
||||||
|
let start = ring.len().saturating_sub(limit);
|
||||||
|
ring.iter().skip(start).cloned().collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn filter(
|
||||||
|
&self,
|
||||||
|
bucket: Option<&str>,
|
||||||
|
since: Option<DateTime<Utc>>,
|
||||||
|
limit: usize,
|
||||||
|
) -> Vec<BucketErrorEvent> {
|
||||||
|
let ring = self.ring.read().await;
|
||||||
|
ring.iter()
|
||||||
|
.rev()
|
||||||
|
.filter(|ev| bucket.map_or(true, |b| ev.target == b))
|
||||||
|
.filter(|ev| since.map_or(true, |s| ev.ts >= s))
|
||||||
|
.take(limit)
|
||||||
|
.cloned()
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.into_iter()
|
||||||
|
.rev()
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Summarize errors in the last N minutes.
|
||||||
|
pub async fn health(&self, period_minutes: i64) -> HealthReport {
|
||||||
|
use std::collections::HashMap;
|
||||||
|
let cutoff = Utc::now() - chrono::Duration::minutes(period_minutes);
|
||||||
|
let ring = self.ring.read().await;
|
||||||
|
let recent: Vec<_> = ring.iter().filter(|ev| ev.ts >= cutoff).collect();
|
||||||
|
let mut per_bucket: HashMap<String, usize> = HashMap::new();
|
||||||
|
for ev in &recent {
|
||||||
|
*per_bucket.entry(ev.target.clone()).or_insert(0) += 1;
|
||||||
|
}
|
||||||
|
let unhealthy_buckets: Vec<String> = per_bucket
|
||||||
|
.iter()
|
||||||
|
.filter(|(_, c)| **c >= 3)
|
||||||
|
.map(|(k, _)| k.clone())
|
||||||
|
.collect();
|
||||||
|
HealthReport {
|
||||||
|
period_minutes,
|
||||||
|
total_errors: recent.len(),
|
||||||
|
per_bucket,
|
||||||
|
unhealthy_buckets,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Force an immediate flush of buffered events to object storage.
|
||||||
|
pub async fn flush(&self) -> Result<(), String> {
|
||||||
|
self.log.flush().await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Consolidate all batch files into one. Operator cleanup.
|
||||||
|
pub async fn compact(&self) -> Result<CompactStats, String> {
|
||||||
|
self.log.compact().await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// How many JSONL batch files currently exist.
|
||||||
|
pub async fn file_count(&self) -> Result<usize, String> {
|
||||||
|
self.log.file_count().await
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,179 @@
|
|||||||
|
/// HTTP surface for federation operator endpoints.
|
||||||
|
///
|
||||||
|
/// - `GET /buckets` — configured buckets with reachability status
|
||||||
|
/// - `GET /errors` — recent bucket op failures (filterable)
|
||||||
|
/// - `GET /bucket-health` — aggregated errors-per-bucket in the last 5 minutes
|
||||||
|
///
|
||||||
|
/// Mounted under `/storage` by the gateway. `/storage/health` is already
|
||||||
|
/// claimed by the existing storage router for service liveness, so the
|
||||||
|
/// aggregated bucket health lives at `/storage/bucket-health`.
|
||||||
|
|
||||||
|
use axum::{
|
||||||
|
Json, Router,
|
||||||
|
extract::{Path, Query, State},
|
||||||
|
http::StatusCode,
|
||||||
|
response::IntoResponse,
|
||||||
|
routing::{delete, get, post, put},
|
||||||
|
};
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use serde::Deserialize;
|
||||||
|
use shared::config::BucketConfig;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use crate::error_journal::BucketErrorEvent;
|
||||||
|
use crate::registry::{BucketInfo, BucketRegistry};
|
||||||
|
|
||||||
|
pub fn router(registry: Arc<BucketRegistry>) -> Router {
|
||||||
|
Router::new()
|
||||||
|
.route("/buckets", get(list_buckets).post(add_bucket))
|
||||||
|
.route("/buckets/{name}", delete(remove_bucket))
|
||||||
|
.route("/errors", get(list_errors))
|
||||||
|
.route("/errors/flush", post(flush_errors))
|
||||||
|
.route("/errors/compact", post(compact_errors))
|
||||||
|
.route("/bucket-health", get(get_health))
|
||||||
|
.route("/buckets/{bucket}/objects/{*key}", put(put_bucket_object))
|
||||||
|
.route("/buckets/{bucket}/objects/{*key}", get(get_bucket_object))
|
||||||
|
.with_state(registry)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Provision + register a bucket at runtime. Body is a `BucketConfig`.
|
||||||
|
///
|
||||||
|
/// Federation layer 2: profile buckets (`profile:alice`) and tenant
|
||||||
|
/// buckets can be added after startup without a service restart.
|
||||||
|
async fn add_bucket(
|
||||||
|
State(reg): State<Arc<BucketRegistry>>,
|
||||||
|
Json(bc): Json<BucketConfig>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
// Safety net: local backends must land somewhere under the configured
|
||||||
|
// profile_root (for profile:*) or be otherwise explicitly rooted.
|
||||||
|
// Refuse empty / missing root to avoid "bucket created at ./" surprises.
|
||||||
|
if bc.backend == "local" {
|
||||||
|
match bc.root.as_deref() {
|
||||||
|
Some(r) if !r.trim().is_empty() => {}
|
||||||
|
_ => return Err((StatusCode::BAD_REQUEST,
|
||||||
|
"local bucket requires a non-empty 'root' path".to_string())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
match reg.add_bucket(bc).await {
|
||||||
|
Ok(info) => Ok((StatusCode::CREATED, Json(info))),
|
||||||
|
Err(e) => {
|
||||||
|
let code = if e.contains("already registered") {
|
||||||
|
StatusCode::CONFLICT
|
||||||
|
} else {
|
||||||
|
StatusCode::BAD_REQUEST
|
||||||
|
};
|
||||||
|
Err((code, e))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Unregister a bucket. Refused for primary / rescue / unknown names.
|
||||||
|
async fn remove_bucket(
|
||||||
|
State(reg): State<Arc<BucketRegistry>>,
|
||||||
|
Path(name): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match reg.remove_bucket(&name) {
|
||||||
|
Ok(()) => Ok((StatusCode::OK, format!("unregistered: {name}"))),
|
||||||
|
Err(e) => {
|
||||||
|
let code = if e.contains("cannot remove") {
|
||||||
|
StatusCode::FORBIDDEN
|
||||||
|
} else if e.contains("not registered") {
|
||||||
|
StatusCode::NOT_FOUND
|
||||||
|
} else {
|
||||||
|
StatusCode::BAD_REQUEST
|
||||||
|
};
|
||||||
|
Err((code, e))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_buckets(State(reg): State<Arc<BucketRegistry>>) -> Json<Vec<BucketInfo>> {
|
||||||
|
Json(reg.list().await)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct ErrorQuery {
|
||||||
|
bucket: Option<String>,
|
||||||
|
since: Option<DateTime<Utc>>,
|
||||||
|
#[serde(default = "default_limit")]
|
||||||
|
limit: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_limit() -> usize { 50 }
|
||||||
|
|
||||||
|
async fn list_errors(
|
||||||
|
State(reg): State<Arc<BucketRegistry>>,
|
||||||
|
Query(q): Query<ErrorQuery>,
|
||||||
|
) -> Json<Vec<BucketErrorEvent>> {
|
||||||
|
Json(reg.journal().filter(q.bucket.as_deref(), q.since, q.limit).await)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct HealthQuery {
|
||||||
|
#[serde(default = "default_period")]
|
||||||
|
minutes: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_period() -> i64 { 5 }
|
||||||
|
|
||||||
|
async fn get_health(
|
||||||
|
State(reg): State<Arc<BucketRegistry>>,
|
||||||
|
Query(q): Query<HealthQuery>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
Json(reg.journal().health(q.minutes).await)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Bucket-aware write. Hard-fails if the target bucket is unreachable;
|
||||||
|
/// never falls back to rescue for writes.
|
||||||
|
async fn put_bucket_object(
|
||||||
|
State(reg): State<Arc<BucketRegistry>>,
|
||||||
|
Path((bucket, key)): Path<(String, String)>,
|
||||||
|
body: bytes::Bytes,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match reg.write_smart(&bucket, &key, body).await {
|
||||||
|
Ok(()) => (StatusCode::CREATED, format!("stored: {bucket}/{key}")).into_response(),
|
||||||
|
Err(e) => (StatusCode::SERVICE_UNAVAILABLE, e).into_response(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Bucket-aware read. Falls through to the rescue bucket if the target
|
||||||
|
/// is unreachable or the key is missing; emits observability headers so
|
||||||
|
/// callers can detect the fallback.
|
||||||
|
async fn flush_errors(State(reg): State<Arc<BucketRegistry>>) -> impl IntoResponse {
|
||||||
|
match reg.journal().flush().await {
|
||||||
|
Ok(()) => (StatusCode::OK, "flushed").into_response(),
|
||||||
|
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e).into_response(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn compact_errors(State(reg): State<Arc<BucketRegistry>>) -> impl IntoResponse {
|
||||||
|
match reg.journal().compact().await {
|
||||||
|
Ok(stats) => Json(stats).into_response(),
|
||||||
|
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e).into_response(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_bucket_object(
|
||||||
|
State(reg): State<Arc<BucketRegistry>>,
|
||||||
|
Path((bucket, key)): Path<(String, String)>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match reg.read_smart(&bucket, &key).await {
|
||||||
|
Ok(outcome) => {
|
||||||
|
let mut resp = outcome.data.into_response();
|
||||||
|
if outcome.rescued {
|
||||||
|
let headers = resp.headers_mut();
|
||||||
|
headers.insert("x-lakehouse-rescue-used", "true".parse().unwrap());
|
||||||
|
headers.insert(
|
||||||
|
"x-lakehouse-original-bucket",
|
||||||
|
outcome.original_bucket.parse().unwrap(),
|
||||||
|
);
|
||||||
|
headers.insert(
|
||||||
|
"x-lakehouse-served-by",
|
||||||
|
outcome.served_by.parse().unwrap(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
resp
|
||||||
|
}
|
||||||
|
Err(e) => (StatusCode::NOT_FOUND, e).into_response(),
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,3 +1,7 @@
|
|||||||
|
pub mod append_log;
|
||||||
pub mod backend;
|
pub mod backend;
|
||||||
|
pub mod error_journal;
|
||||||
|
pub mod federation_service;
|
||||||
pub mod ops;
|
pub mod ops;
|
||||||
|
pub mod registry;
|
||||||
pub mod service;
|
pub mod service;
|
||||||
|
|||||||
@@ -0,0 +1,377 @@
|
|||||||
|
/// Multi-backend bucket registry — the federation foundation.
|
||||||
|
///
|
||||||
|
/// Federation rule: every `ObjectRef` belongs to exactly one named bucket.
|
||||||
|
/// The registry resolves bucket names to `object_store` backends, handles
|
||||||
|
/// rescue-bucket fallback on read failure, writes every failure to the
|
||||||
|
/// error journal, and exposes a health summary for operators.
|
||||||
|
///
|
||||||
|
/// Existing call sites can keep using `ops::*` with `registry.get(name)`.
|
||||||
|
/// New bucket-aware call sites use `registry.read_smart` / `write_smart`
|
||||||
|
/// which handle fallback + journaling automatically.
|
||||||
|
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use object_store::local::LocalFileSystem;
|
||||||
|
use serde::Serialize;
|
||||||
|
use shared::config::{BucketConfig, StorageConfig};
|
||||||
|
use shared::secrets::{BucketCredentials, SecretsProvider};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::{Arc, RwLock};
|
||||||
|
|
||||||
|
use crate::error_journal::{BucketErrorEvent, ErrorJournal};
|
||||||
|
|
||||||
|
/// A registered bucket — the store handle + its configuration.
|
||||||
|
pub struct BucketEntry {
|
||||||
|
pub name: String,
|
||||||
|
pub backend: String,
|
||||||
|
pub store: Arc<dyn ObjectStore>,
|
||||||
|
pub config: BucketConfig,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Read outcome — may have been rescued.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct ReadOutcome {
|
||||||
|
pub data: bytes::Bytes,
|
||||||
|
pub rescued: bool,
|
||||||
|
pub original_bucket: String,
|
||||||
|
pub served_by: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Summary entry for GET /storage/buckets.
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct BucketInfo {
|
||||||
|
pub name: String,
|
||||||
|
pub backend: String,
|
||||||
|
pub reachable: bool,
|
||||||
|
pub role: BucketRole,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, PartialEq)]
|
||||||
|
#[serde(rename_all = "lowercase")]
|
||||||
|
pub enum BucketRole {
|
||||||
|
Primary,
|
||||||
|
Rescue,
|
||||||
|
Profile,
|
||||||
|
Tenant,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct BucketRegistry {
|
||||||
|
/// RwLock because federation layer 2 adds runtime bucket lifecycle
|
||||||
|
/// (`POST /storage/buckets` / `DELETE /storage/buckets/{name}`).
|
||||||
|
/// Almost all accesses are reads — writes happen at provision /
|
||||||
|
/// deactivate time only.
|
||||||
|
buckets: RwLock<HashMap<String, Arc<BucketEntry>>>,
|
||||||
|
default: String,
|
||||||
|
rescue: Option<String>,
|
||||||
|
profile_root: String,
|
||||||
|
/// Held so runtime `add_bucket` can resolve secret_ref handles.
|
||||||
|
secrets: Arc<dyn SecretsProvider>,
|
||||||
|
journal: ErrorJournal,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl BucketRegistry {
|
||||||
|
/// Build the registry from storage config + secrets provider.
|
||||||
|
/// Back-compat: if `buckets` is empty, synthesize a `primary` bucket from
|
||||||
|
/// the legacy `root` field so pre-federation configs keep working.
|
||||||
|
pub async fn from_config(
|
||||||
|
cfg: &StorageConfig,
|
||||||
|
secrets: Arc<dyn SecretsProvider>,
|
||||||
|
) -> Result<Self, String> {
|
||||||
|
let mut buckets: HashMap<String, Arc<BucketEntry>> = HashMap::new();
|
||||||
|
|
||||||
|
let bucket_configs: Vec<BucketConfig> = if cfg.buckets.is_empty() {
|
||||||
|
vec![BucketConfig {
|
||||||
|
name: "primary".to_string(),
|
||||||
|
backend: "local".to_string(),
|
||||||
|
root: Some(cfg.root.clone()),
|
||||||
|
bucket: None,
|
||||||
|
region: None,
|
||||||
|
endpoint: None,
|
||||||
|
secret_ref: None,
|
||||||
|
}]
|
||||||
|
} else {
|
||||||
|
cfg.buckets.clone()
|
||||||
|
};
|
||||||
|
|
||||||
|
for bc in bucket_configs {
|
||||||
|
let store = build_store(&bc, secrets.as_ref()).await?;
|
||||||
|
let entry = Arc::new(BucketEntry {
|
||||||
|
name: bc.name.clone(),
|
||||||
|
backend: bc.backend.clone(),
|
||||||
|
store,
|
||||||
|
config: bc.clone(),
|
||||||
|
});
|
||||||
|
buckets.insert(bc.name.clone(), entry);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Ensure `primary` always exists — it's where error journals live.
|
||||||
|
if !buckets.contains_key("primary") {
|
||||||
|
return Err("no bucket named 'primary' configured — required as error-journal home".into());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Rescue bucket is optional but, if named, must exist.
|
||||||
|
if let Some(r) = &cfg.rescue_bucket {
|
||||||
|
if !buckets.contains_key(r) {
|
||||||
|
return Err(format!("rescue_bucket '{r}' not found among configured buckets"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let journal = ErrorJournal::new(buckets.get("primary").unwrap().store.clone());
|
||||||
|
let _ = journal.load_recent().await;
|
||||||
|
|
||||||
|
Ok(Self {
|
||||||
|
buckets: RwLock::new(buckets),
|
||||||
|
default: "primary".to_string(),
|
||||||
|
rescue: cfg.rescue_bucket.clone(),
|
||||||
|
profile_root: cfg.profile_root.clone(),
|
||||||
|
secrets,
|
||||||
|
journal,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn default_name(&self) -> &str { &self.default }
|
||||||
|
|
||||||
|
pub fn rescue_name(&self) -> Option<&str> { self.rescue.as_deref() }
|
||||||
|
|
||||||
|
pub fn profile_root(&self) -> &str { &self.profile_root }
|
||||||
|
|
||||||
|
pub fn journal(&self) -> &ErrorJournal { &self.journal }
|
||||||
|
|
||||||
|
/// Resolve a bucket name to its object store. Existing call sites use
|
||||||
|
/// this as a drop-in replacement for the old single-store pattern.
|
||||||
|
///
|
||||||
|
/// Uses `std::sync::RwLock` — every caller holds the guard for
|
||||||
|
/// microseconds (clones an Arc and drops), so there's no async-
|
||||||
|
/// blocking concern. Never hold across an await.
|
||||||
|
pub fn get(&self, bucket: &str) -> Result<Arc<dyn ObjectStore>, String> {
|
||||||
|
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||||
|
guard
|
||||||
|
.get(bucket)
|
||||||
|
.map(|e| e.store.clone())
|
||||||
|
.ok_or_else(|| format!("unknown bucket: {bucket}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The default bucket's store — use for code paths that don't yet know
|
||||||
|
/// about buckets.
|
||||||
|
pub fn default_store(&self) -> Arc<dyn ObjectStore> {
|
||||||
|
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||||
|
guard.get(&self.default).unwrap().store.clone()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// True if this bucket name is registered.
|
||||||
|
pub fn contains(&self, bucket: &str) -> bool {
|
||||||
|
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||||
|
guard.contains_key(bucket)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all registered buckets. Checks reachability by doing a trivial
|
||||||
|
/// `list` with limit 1 on each.
|
||||||
|
pub async fn list(&self) -> Vec<BucketInfo> {
|
||||||
|
// Snapshot (name, backend, store_clone) under the lock, then probe
|
||||||
|
// outside — probing can be slow (network), don't hold the lock.
|
||||||
|
let snapshot: Vec<(String, String, Arc<dyn ObjectStore>)> = {
|
||||||
|
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||||
|
guard
|
||||||
|
.iter()
|
||||||
|
.map(|(n, e)| (n.clone(), e.backend.clone(), e.store.clone()))
|
||||||
|
.collect()
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut out = Vec::with_capacity(snapshot.len());
|
||||||
|
for (name, backend, store) in snapshot {
|
||||||
|
let reachable = probe(&store).await;
|
||||||
|
let role = self.classify(&name);
|
||||||
|
out.push(BucketInfo { name, backend, reachable, role });
|
||||||
|
}
|
||||||
|
out.sort_by(|a, b| a.name.cmp(&b.name));
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|
||||||
|
fn classify(&self, name: &str) -> BucketRole {
|
||||||
|
if name == self.default { BucketRole::Primary }
|
||||||
|
else if Some(name) == self.rescue.as_deref() { BucketRole::Rescue }
|
||||||
|
else if name.starts_with("profile:") { BucketRole::Profile }
|
||||||
|
else { BucketRole::Tenant }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Provision + register a bucket at runtime (federation layer 2).
|
||||||
|
/// Returns Err if the name is already registered. Local backends get
|
||||||
|
/// their root directory created automatically.
|
||||||
|
pub async fn add_bucket(&self, bc: BucketConfig) -> Result<BucketInfo, String> {
|
||||||
|
if self.contains(&bc.name) {
|
||||||
|
return Err(format!("bucket '{}' already registered", bc.name));
|
||||||
|
}
|
||||||
|
let store = build_store(&bc, self.secrets.as_ref()).await?;
|
||||||
|
let reachable = probe(&store).await;
|
||||||
|
let entry = Arc::new(BucketEntry {
|
||||||
|
name: bc.name.clone(),
|
||||||
|
backend: bc.backend.clone(),
|
||||||
|
store,
|
||||||
|
config: bc.clone(),
|
||||||
|
});
|
||||||
|
let role = self.classify(&bc.name);
|
||||||
|
let info = BucketInfo {
|
||||||
|
name: bc.name.clone(),
|
||||||
|
backend: bc.backend.clone(),
|
||||||
|
reachable,
|
||||||
|
role,
|
||||||
|
};
|
||||||
|
{
|
||||||
|
let mut guard = self.buckets.write().expect("bucket registry lock poisoned");
|
||||||
|
guard.insert(bc.name.clone(), entry);
|
||||||
|
}
|
||||||
|
tracing::info!("registered bucket '{}' backend={} reachable={}",
|
||||||
|
bc.name, bc.backend, reachable);
|
||||||
|
Ok(info)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Unregister a bucket. Refuses to remove `primary` or the configured
|
||||||
|
/// rescue bucket — those are load-bearing. Caller is responsible for
|
||||||
|
/// any final flush of bucket-local state before calling this.
|
||||||
|
pub fn remove_bucket(&self, name: &str) -> Result<(), String> {
|
||||||
|
if name == self.default {
|
||||||
|
return Err("cannot remove primary bucket".into());
|
||||||
|
}
|
||||||
|
if Some(name) == self.rescue.as_deref() {
|
||||||
|
return Err(format!("cannot remove rescue bucket '{name}'"));
|
||||||
|
}
|
||||||
|
let removed = {
|
||||||
|
let mut guard = self.buckets.write().expect("bucket registry lock poisoned");
|
||||||
|
guard.remove(name)
|
||||||
|
};
|
||||||
|
if removed.is_none() {
|
||||||
|
return Err(format!("bucket '{name}' not registered"));
|
||||||
|
}
|
||||||
|
tracing::info!("unregistered bucket '{}'", name);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Read with rescue-bucket fallback. If the target bucket fails and a
|
||||||
|
/// rescue is configured, retries against rescue. Records every failure
|
||||||
|
/// in the error journal.
|
||||||
|
pub async fn read_smart(&self, bucket: &str, key: &str) -> Result<ReadOutcome, String> {
|
||||||
|
let target_store = {
|
||||||
|
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||||
|
guard.get(bucket).map(|e| e.store.clone())
|
||||||
|
.ok_or_else(|| format!("unknown bucket: {bucket}"))?
|
||||||
|
};
|
||||||
|
|
||||||
|
match crate::ops::get(&target_store, key).await {
|
||||||
|
Ok(data) => Ok(ReadOutcome {
|
||||||
|
data, rescued: false,
|
||||||
|
original_bucket: bucket.to_string(),
|
||||||
|
served_by: bucket.to_string(),
|
||||||
|
}),
|
||||||
|
Err(err) => {
|
||||||
|
// Record failure regardless of what happens next.
|
||||||
|
self.journal.append(BucketErrorEvent::new_read(bucket, key, &err)).await;
|
||||||
|
|
||||||
|
// Try rescue, if any.
|
||||||
|
if let Some(rescue_name) = &self.rescue {
|
||||||
|
if rescue_name != bucket {
|
||||||
|
let rescue_store = {
|
||||||
|
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||||
|
guard.get(rescue_name).map(|e| e.store.clone())
|
||||||
|
};
|
||||||
|
if let Some(rescue_store) = rescue_store {
|
||||||
|
match crate::ops::get(&rescue_store, key).await {
|
||||||
|
Ok(data) => {
|
||||||
|
self.journal.mark_rescued_last(bucket, key).await;
|
||||||
|
return Ok(ReadOutcome {
|
||||||
|
data, rescued: true,
|
||||||
|
original_bucket: bucket.to_string(),
|
||||||
|
served_by: rescue_name.clone(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Err(rescue_err) => {
|
||||||
|
return Err(format!(
|
||||||
|
"read '{key}' failed in '{bucket}' ({err}); rescue '{rescue_name}' also failed ({rescue_err})"
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(format!("read '{key}' failed in '{bucket}': {err}"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Write always goes to target. No rescue fallback for writes — writes
|
||||||
|
/// that silently vanish are the worst possible failure.
|
||||||
|
pub async fn write_smart(
|
||||||
|
&self,
|
||||||
|
bucket: &str,
|
||||||
|
key: &str,
|
||||||
|
data: bytes::Bytes,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let target_store = {
|
||||||
|
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||||
|
guard.get(bucket).map(|e| e.store.clone())
|
||||||
|
.ok_or_else(|| format!("unknown bucket: {bucket}"))?
|
||||||
|
};
|
||||||
|
match crate::ops::put(&target_store, key, data).await {
|
||||||
|
Ok(()) => Ok(()),
|
||||||
|
Err(err) => {
|
||||||
|
self.journal.append(BucketErrorEvent::new_write(bucket, key, &err)).await;
|
||||||
|
Err(format!("write '{key}' failed in '{bucket}': {err}"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Trivial reachability check — try to list with limit 0.
|
||||||
|
async fn probe(store: &Arc<dyn ObjectStore>) -> bool {
|
||||||
|
use futures::StreamExt;
|
||||||
|
let mut stream = store.list(None);
|
||||||
|
// Pulling the first item confirms the store responds. Empty bucket = ok.
|
||||||
|
match stream.next().await {
|
||||||
|
Some(Ok(_)) => true,
|
||||||
|
None => true, // empty but reachable
|
||||||
|
Some(Err(_)) => false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a concrete ObjectStore from a BucketConfig.
|
||||||
|
async fn build_store(
|
||||||
|
bc: &BucketConfig,
|
||||||
|
secrets: &dyn SecretsProvider,
|
||||||
|
) -> Result<Arc<dyn ObjectStore>, String> {
|
||||||
|
match bc.backend.as_str() {
|
||||||
|
"local" => {
|
||||||
|
let root = bc.root.as_deref()
|
||||||
|
.ok_or_else(|| format!("bucket '{}' is backend=local but has no root", bc.name))?;
|
||||||
|
std::fs::create_dir_all(root)
|
||||||
|
.map_err(|e| format!("create bucket dir '{root}': {e}"))?;
|
||||||
|
let fs = LocalFileSystem::new_with_prefix(root)
|
||||||
|
.map_err(|e| format!("init local bucket '{}': {e}", bc.name))?;
|
||||||
|
Ok(Arc::new(fs))
|
||||||
|
}
|
||||||
|
"s3" => {
|
||||||
|
let handle = bc.secret_ref.as_deref()
|
||||||
|
.ok_or_else(|| format!("s3 bucket '{}' has no secret_ref", bc.name))?;
|
||||||
|
let creds: BucketCredentials = secrets.resolve(handle).await?;
|
||||||
|
let s3_bucket = bc.bucket.as_deref()
|
||||||
|
.ok_or_else(|| format!("s3 bucket '{}' has no `bucket` name", bc.name))?;
|
||||||
|
let region = bc.region.as_deref().unwrap_or("us-east-1");
|
||||||
|
|
||||||
|
let mut builder = object_store::aws::AmazonS3Builder::new()
|
||||||
|
.with_bucket_name(s3_bucket)
|
||||||
|
.with_region(region)
|
||||||
|
.with_access_key_id(&creds.access_key)
|
||||||
|
.with_secret_access_key(&creds.secret_key);
|
||||||
|
if let Some(endpoint) = &bc.endpoint {
|
||||||
|
builder = builder.with_endpoint(endpoint);
|
||||||
|
// MinIO and other S3-compatible services often run on plain
|
||||||
|
// HTTP. object_store refuses HTTP by default — opt in when
|
||||||
|
// a custom endpoint is configured (TLS endpoints work either way).
|
||||||
|
if endpoint.starts_with("http://") {
|
||||||
|
builder = builder.with_allow_http(true);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let s3 = builder.build()
|
||||||
|
.map_err(|e| format!("init s3 bucket '{}': {e}", bc.name))?;
|
||||||
|
Ok(Arc::new(s3))
|
||||||
|
}
|
||||||
|
other => Err(format!("unknown backend '{other}' for bucket '{}'", bc.name)),
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -162,3 +162,46 @@ th {
|
|||||||
}
|
}
|
||||||
td { padding: 8px 14px; border-bottom: 1px solid var(--border); }
|
td { padding: 8px 14px; border-bottom: 1px solid var(--border); }
|
||||||
tr:hover td { background: var(--accent-glow); }
|
tr:hover td { background: var(--accent-glow); }
|
||||||
|
|
||||||
|
/* Dashboard */
|
||||||
|
.dashboard-hero { margin-bottom: 28px; }
|
||||||
|
.dashboard-hero h2 { font-size: 32px; font-weight: 700; }
|
||||||
|
|
||||||
|
.stat-grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(160px, 1fr)); gap: 12px; margin-bottom: 32px; }
|
||||||
|
.stat-card {
|
||||||
|
background: var(--surface); border: 1px solid var(--border); border-radius: 10px;
|
||||||
|
padding: 20px 16px; text-align: center;
|
||||||
|
}
|
||||||
|
.stat-card.accent { border-color: var(--accent); background: var(--accent-glow); }
|
||||||
|
.stat-value { font-size: 28px; font-weight: 700; color: var(--text); }
|
||||||
|
.stat-card.accent .stat-value { color: var(--accent); }
|
||||||
|
.stat-label { font-size: 11px; color: var(--text-dim); margin-top: 4px; text-transform: uppercase; letter-spacing: 0.05em; }
|
||||||
|
|
||||||
|
.architecture-section, .phases-section { margin-bottom: 28px; }
|
||||||
|
.architecture-section h3, .phases-section h3 { font-size: 16px; font-weight: 600; color: var(--text-dim); margin-bottom: 12px; text-transform: uppercase; letter-spacing: 0.08em; }
|
||||||
|
|
||||||
|
.arch-grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(240px, 1fr)); gap: 10px; }
|
||||||
|
.arch-card {
|
||||||
|
background: var(--surface); border: 1px solid var(--border); border-radius: 8px;
|
||||||
|
padding: 14px; border-left: 3px solid var(--accent);
|
||||||
|
}
|
||||||
|
.arch-title { font-weight: 600; font-size: 13px; margin-bottom: 4px; color: var(--accent); }
|
||||||
|
.arch-items { font-size: 11px; color: var(--text-dim); line-height: 1.6; }
|
||||||
|
|
||||||
|
.phase-list { display: flex; flex-direction: column; gap: 2px; }
|
||||||
|
.phase-item {
|
||||||
|
display: grid; grid-template-columns: 40px 180px 1fr; gap: 12px; align-items: center;
|
||||||
|
padding: 8px 12px; border-radius: 6px; font-size: 12px;
|
||||||
|
background: var(--surface); border-left: 3px solid var(--success);
|
||||||
|
}
|
||||||
|
.phase-num { font-weight: 700; color: var(--accent); text-align: center; }
|
||||||
|
.phase-name { font-weight: 600; }
|
||||||
|
.phase-detail { color: var(--text-dim); }
|
||||||
|
|
||||||
|
/* Ingest */
|
||||||
|
.table-list { margin-top: 12px; }
|
||||||
|
.table-item {
|
||||||
|
display: flex; justify-content: space-between; align-items: center;
|
||||||
|
padding: 8px 12px; border-bottom: 1px solid var(--border); font-size: 13px;
|
||||||
|
}
|
||||||
|
.table-item:hover { background: var(--accent-glow); }
|
||||||
|
|||||||
@@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"id": "6d0002ef-28de-4b92-abed-0bfb7ab5fb6f",
|
||||||
|
"name": "candidates",
|
||||||
|
"schema_fingerprint": "auto",
|
||||||
|
"objects": [
|
||||||
|
{
|
||||||
|
"bucket": "data",
|
||||||
|
"key": "datasets/candidates.parquet",
|
||||||
|
"size_bytes": 10592165,
|
||||||
|
"created_at": "2026-03-28T01:54:06.898505222Z"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"created_at": "2026-03-28T01:54:06.898506121Z",
|
||||||
|
"updated_at": "2026-03-28T01:54:06.898506121Z",
|
||||||
|
"description": "",
|
||||||
|
"owner": "",
|
||||||
|
"sensitivity": null,
|
||||||
|
"columns": [],
|
||||||
|
"lineage": null,
|
||||||
|
"freshness": null,
|
||||||
|
"tags": [],
|
||||||
|
"row_count": null
|
||||||
|
}
|
||||||
@@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"id": "73a4ef71-5b2a-4124-a2c2-fcefdde99624",
|
||||||
|
"name": "job_orders",
|
||||||
|
"schema_fingerprint": "auto",
|
||||||
|
"objects": [
|
||||||
|
{
|
||||||
|
"bucket": "data",
|
||||||
|
"key": "datasets/job_orders.parquet",
|
||||||
|
"size_bytes": 905534,
|
||||||
|
"created_at": "2026-03-28T01:54:07.007332060Z"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"created_at": "2026-03-28T01:54:07.007332775Z",
|
||||||
|
"updated_at": "2026-03-28T01:54:07.007332775Z",
|
||||||
|
"description": "",
|
||||||
|
"owner": "",
|
||||||
|
"sensitivity": null,
|
||||||
|
"columns": [],
|
||||||
|
"lineage": null,
|
||||||
|
"freshness": null,
|
||||||
|
"tags": [],
|
||||||
|
"row_count": null
|
||||||
|
}
|
||||||
@@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"id": "b56694bb-f4f7-447d-b4d6-ac89103a6a3a",
|
||||||
|
"name": "call_log",
|
||||||
|
"schema_fingerprint": "auto",
|
||||||
|
"objects": [
|
||||||
|
{
|
||||||
|
"bucket": "data",
|
||||||
|
"key": "datasets/call_log.parquet",
|
||||||
|
"size_bytes": 35951077,
|
||||||
|
"created_at": "2026-03-28T01:54:15.457573021Z"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"created_at": "2026-03-28T01:54:15.457573701Z",
|
||||||
|
"updated_at": "2026-03-28T01:54:15.457573701Z",
|
||||||
|
"description": "",
|
||||||
|
"owner": "",
|
||||||
|
"sensitivity": null,
|
||||||
|
"columns": [],
|
||||||
|
"lineage": null,
|
||||||
|
"freshness": null,
|
||||||
|
"tags": [],
|
||||||
|
"row_count": null
|
||||||
|
}
|
||||||
@@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"id": "daa5a11b-6898-40ba-a811-3bd4b6f5b750",
|
||||||
|
"name": "email_log",
|
||||||
|
"schema_fingerprint": "auto",
|
||||||
|
"objects": [
|
||||||
|
{
|
||||||
|
"bucket": "data",
|
||||||
|
"key": "datasets/email_log.parquet",
|
||||||
|
"size_bytes": 16768671,
|
||||||
|
"created_at": "2026-03-28T01:54:17.356056562Z"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"created_at": "2026-03-28T01:54:17.356057529Z",
|
||||||
|
"updated_at": "2026-03-28T01:54:17.356057529Z",
|
||||||
|
"description": "",
|
||||||
|
"owner": "",
|
||||||
|
"sensitivity": null,
|
||||||
|
"columns": [],
|
||||||
|
"lineage": null,
|
||||||
|
"freshness": null,
|
||||||
|
"tags": [],
|
||||||
|
"row_count": null
|
||||||
|
}
|
||||||
@@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"id": "dd396ee4-3a1a-4939-af60-84d409eebed4",
|
||||||
|
"name": "clients",
|
||||||
|
"schema_fingerprint": "auto",
|
||||||
|
"objects": [
|
||||||
|
{
|
||||||
|
"bucket": "data",
|
||||||
|
"key": "datasets/clients.parquet",
|
||||||
|
"size_bytes": 21971,
|
||||||
|
"created_at": "2026-03-28T01:54:06.904379300Z"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"created_at": "2026-03-28T01:54:06.904379793Z",
|
||||||
|
"updated_at": "2026-03-28T01:54:06.904379793Z",
|
||||||
|
"description": "",
|
||||||
|
"owner": "",
|
||||||
|
"sensitivity": null,
|
||||||
|
"columns": [],
|
||||||
|
"lineage": null,
|
||||||
|
"freshness": null,
|
||||||
|
"tags": [],
|
||||||
|
"row_count": null
|
||||||
|
}
|
||||||
@@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"id": "f85285e3-111a-4a55-aa58-05fd740ed862",
|
||||||
|
"name": "timesheets",
|
||||||
|
"schema_fingerprint": "auto",
|
||||||
|
"objects": [
|
||||||
|
{
|
||||||
|
"bucket": "data",
|
||||||
|
"key": "datasets/timesheets.parquet",
|
||||||
|
"size_bytes": 17539932,
|
||||||
|
"created_at": "2026-03-28T01:54:11.947773535Z"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"created_at": "2026-03-28T01:54:11.947778065Z",
|
||||||
|
"updated_at": "2026-03-28T01:54:11.947778065Z",
|
||||||
|
"description": "",
|
||||||
|
"owner": "",
|
||||||
|
"sensitivity": null,
|
||||||
|
"columns": [],
|
||||||
|
"lineage": null,
|
||||||
|
"freshness": null,
|
||||||
|
"tags": [],
|
||||||
|
"row_count": null
|
||||||
|
}
|
||||||
@@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"id": "ffa72582-dc11-4454-a46e-35839a7d04bb",
|
||||||
|
"name": "placements",
|
||||||
|
"schema_fingerprint": "auto",
|
||||||
|
"objects": [
|
||||||
|
{
|
||||||
|
"bucket": "data",
|
||||||
|
"key": "datasets/placements.parquet",
|
||||||
|
"size_bytes": 1213820,
|
||||||
|
"created_at": "2026-03-28T01:54:07.109218264Z"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"created_at": "2026-03-28T01:54:07.109219070Z",
|
||||||
|
"updated_at": "2026-03-28T01:54:07.109219070Z",
|
||||||
|
"description": "",
|
||||||
|
"owner": "",
|
||||||
|
"sensitivity": null,
|
||||||
|
"columns": [],
|
||||||
|
"lineage": null,
|
||||||
|
"freshness": null,
|
||||||
|
"tags": [],
|
||||||
|
"row_count": null
|
||||||
|
}
|
||||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
+461
-17
@@ -109,20 +109,43 @@ async fn fetch_health(path: &str) -> Result<String, String> {
|
|||||||
resp.text().await.map_err(|e| e.to_string())
|
resp.text().await.map_err(|e| e.to_string())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Get schema context for all datasets (used for AI SQL generation)
|
/// Cached schema context — built once, reused across questions.
|
||||||
|
static SCHEMA_CACHE: std::sync::OnceLock<std::sync::Mutex<Option<String>>> = std::sync::OnceLock::new();
|
||||||
|
|
||||||
|
async fn get_schema_context_cached(datasets: &[Dataset]) -> String {
|
||||||
|
// Check cache first
|
||||||
|
{
|
||||||
|
let cache = SCHEMA_CACHE.get_or_init(|| std::sync::Mutex::new(None));
|
||||||
|
if let Ok(guard) = cache.lock() {
|
||||||
|
if let Some(ref cached) = *guard {
|
||||||
|
return cached.clone();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Build and cache
|
||||||
|
let ctx = get_schema_context(datasets).await;
|
||||||
|
if let Some(cache) = SCHEMA_CACHE.get() {
|
||||||
|
if let Ok(mut guard) = cache.lock() {
|
||||||
|
*guard = Some(ctx.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ctx
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get schema context for datasets (used for AI SQL generation).
|
||||||
async fn get_schema_context(datasets: &[Dataset]) -> String {
|
async fn get_schema_context(datasets: &[Dataset]) -> String {
|
||||||
let mut ctx = String::from("Available tables:\n\n");
|
let core_tables = ["candidates", "clients", "job_orders", "placements", "timesheets", "call_log", "email_log"];
|
||||||
for ds in datasets {
|
let mut ctx = String::from("DATABASE SCHEMA:\n\n");
|
||||||
|
for ds in datasets.iter().filter(|d| core_tables.contains(&d.name.as_str())) {
|
||||||
let desc = run_sql(&format!("DESCRIBE {}", ds.name)).await;
|
let desc = run_sql(&format!("DESCRIBE {}", ds.name)).await;
|
||||||
match desc {
|
match desc {
|
||||||
Ok(resp) => {
|
Ok(resp) => {
|
||||||
ctx.push_str(&format!("TABLE: {}\n", ds.name));
|
ctx.push_str(&format!("TABLE: {}\n Columns:\n", ds.name));
|
||||||
if let Some(rows) = resp.rows.as_array() {
|
if let Some(rows) = resp.rows.as_array() {
|
||||||
for row in rows {
|
for row in rows {
|
||||||
let col = row.get("column_name").and_then(|v| v.as_str()).unwrap_or("?");
|
let col = row.get("column_name").and_then(|v| v.as_str()).unwrap_or("?");
|
||||||
let dt = row.get("data_type").and_then(|v| v.as_str()).unwrap_or("?");
|
let dt = row.get("data_type").and_then(|v| v.as_str()).unwrap_or("?");
|
||||||
let nullable = row.get("is_nullable").and_then(|v| v.as_str()).unwrap_or("?");
|
ctx.push_str(&format!(" {}.{} ({})\n", ds.name, col, dt));
|
||||||
ctx.push_str(&format!(" - {} ({}, nullable={})\n", col, dt, nullable));
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
ctx.push('\n');
|
ctx.push('\n');
|
||||||
@@ -132,15 +155,29 @@ async fn get_schema_context(datasets: &[Dataset]) -> String {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
// Add relationship hints so the model knows how to JOIN
|
||||||
|
ctx.push_str("RELATIONSHIPS:\n");
|
||||||
|
ctx.push_str(" candidates.candidate_id = placements.candidate_id\n");
|
||||||
|
ctx.push_str(" candidates.candidate_id = timesheets.candidate_id\n");
|
||||||
|
ctx.push_str(" candidates.candidate_id = call_log.candidate_id\n");
|
||||||
|
ctx.push_str(" candidates.candidate_id = email_log.candidate_id\n");
|
||||||
|
ctx.push_str(" clients.client_id = job_orders.client_id\n");
|
||||||
|
ctx.push_str(" clients.client_id = placements.client_id\n");
|
||||||
|
ctx.push_str(" clients.client_id = timesheets.client_id\n");
|
||||||
|
ctx.push_str(" job_orders.job_order_id = placements.job_order_id\n");
|
||||||
|
ctx.push_str(" placements.placement_id = timesheets.placement_id\n");
|
||||||
|
ctx.push_str("\nNOTE: 'vertical' is only in candidates, clients, and job_orders. To get vertical for timesheets or placements, JOIN to those tables.\n");
|
||||||
ctx
|
ctx
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- Tabs ---
|
// --- Tabs ---
|
||||||
#[derive(Clone, PartialEq)]
|
#[derive(Clone, PartialEq)]
|
||||||
enum Tab {
|
enum Tab {
|
||||||
|
Dashboard,
|
||||||
Ask,
|
Ask,
|
||||||
Explore,
|
Explore,
|
||||||
Sql,
|
Sql,
|
||||||
|
Ingest,
|
||||||
Status,
|
Status,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -177,6 +214,11 @@ fn App() -> Element {
|
|||||||
div { class: "header",
|
div { class: "header",
|
||||||
h1 { "LAKEHOUSE" }
|
h1 { "LAKEHOUSE" }
|
||||||
div { class: "tabs",
|
div { class: "tabs",
|
||||||
|
button {
|
||||||
|
class: if *active_tab.read() == Tab::Dashboard { "tab active" } else { "tab" },
|
||||||
|
onclick: move |_| active_tab.set(Tab::Dashboard),
|
||||||
|
"Dashboard"
|
||||||
|
}
|
||||||
button {
|
button {
|
||||||
class: if *active_tab.read() == Tab::Ask { "tab active" } else { "tab" },
|
class: if *active_tab.read() == Tab::Ask { "tab active" } else { "tab" },
|
||||||
onclick: move |_| active_tab.set(Tab::Ask),
|
onclick: move |_| active_tab.set(Tab::Ask),
|
||||||
@@ -192,6 +234,11 @@ fn App() -> Element {
|
|||||||
onclick: move |_| active_tab.set(Tab::Sql),
|
onclick: move |_| active_tab.set(Tab::Sql),
|
||||||
"SQL"
|
"SQL"
|
||||||
}
|
}
|
||||||
|
button {
|
||||||
|
class: if *active_tab.read() == Tab::Ingest { "tab active" } else { "tab" },
|
||||||
|
onclick: move |_| active_tab.set(Tab::Ingest),
|
||||||
|
"Ingest"
|
||||||
|
}
|
||||||
button {
|
button {
|
||||||
class: if *active_tab.read() == Tab::Status { "tab active" } else { "tab" },
|
class: if *active_tab.read() == Tab::Status { "tab active" } else { "tab" },
|
||||||
onclick: move |_| active_tab.set(Tab::Status),
|
onclick: move |_| active_tab.set(Tab::Status),
|
||||||
@@ -208,9 +255,11 @@ fn App() -> Element {
|
|||||||
|
|
||||||
div { class: "content-full",
|
div { class: "content-full",
|
||||||
match *active_tab.read() {
|
match *active_tab.read() {
|
||||||
|
Tab::Dashboard => rsx! { DashboardPanel {} },
|
||||||
Tab::Ask => rsx! { AskPanel { datasets: datasets.read().clone() } },
|
Tab::Ask => rsx! { AskPanel { datasets: datasets.read().clone() } },
|
||||||
Tab::Explore => rsx! { ExplorePanel { datasets: datasets.read().clone() } },
|
Tab::Explore => rsx! { ExplorePanel { datasets: datasets.read().clone() } },
|
||||||
Tab::Sql => rsx! { SqlPanel {} },
|
Tab::Sql => rsx! { SqlPanel {} },
|
||||||
|
Tab::Ingest => rsx! { IngestPanel {} },
|
||||||
Tab::Status => rsx! { StatusPanel {} },
|
Tab::Status => rsx! { StatusPanel {} },
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -241,26 +290,55 @@ fn AskPanel(datasets: Vec<Dataset>) -> Element {
|
|||||||
|
|
||||||
// Step 1: Get schema context
|
// Step 1: Get schema context
|
||||||
step.set("reading schemas...".into());
|
step.set("reading schemas...".into());
|
||||||
let schema_ctx = get_schema_context(&ds).await;
|
let schema_ctx = get_schema_context_cached(&ds).await;
|
||||||
|
|
||||||
// Step 2: Generate SQL
|
// Step 2: Generate SQL
|
||||||
step.set("writing SQL...".into());
|
step.set("writing SQL...".into());
|
||||||
let prompt = format!(
|
let prompt = format!(
|
||||||
"You are a SQL assistant for a data lakehouse using Apache DataFusion (PostgreSQL-compatible SQL).\n\n\
|
"You are a SQL assistant. You write Apache DataFusion SQL (PostgreSQL-compatible).\n\n\
|
||||||
|
CRITICAL: You MUST only use column names that appear in the schema below. Do NOT invent or guess column names.\n\n\
|
||||||
{schema_ctx}\n\
|
{schema_ctx}\n\
|
||||||
User question: {q}\n\n\
|
User question: {q}\n\n\
|
||||||
Write a SQL query that answers this question. Output ONLY the SQL query, nothing else. No markdown, no explanation, no backticks."
|
Rules:\n\
|
||||||
|
- ONLY use table and column names listed above\n\
|
||||||
|
- If unsure about a column name, pick the closest match from the schema\n\
|
||||||
|
- Output ONLY the SQL query. No markdown, no explanation, no backticks."
|
||||||
);
|
);
|
||||||
|
|
||||||
match ai_generate(&prompt, 512).await {
|
match ai_generate(&prompt, 512).await {
|
||||||
Ok(resp) => {
|
Ok(resp) => {
|
||||||
let sql = resp.text.trim().to_string();
|
let sql = clean_sql(&resp.text);
|
||||||
generated_sql.set(Some(sql.clone()));
|
generated_sql.set(Some(sql.clone()));
|
||||||
|
|
||||||
// Step 3: Execute
|
// Step 3: Execute
|
||||||
step.set("running query...".into());
|
step.set("running query...".into());
|
||||||
let query_result = run_sql(&sql).await;
|
let query_result = run_sql(&sql).await;
|
||||||
result.set(Some(query_result));
|
|
||||||
|
// Step 3b: If schema error, retry with the error as feedback
|
||||||
|
if let Err(ref err) = query_result {
|
||||||
|
if err.contains("error") {
|
||||||
|
step.set("fixing SQL...".into());
|
||||||
|
let retry_prompt = format!(
|
||||||
|
"The SQL you wrote had an error:\n{err}\n\n\
|
||||||
|
{schema_ctx}\n\
|
||||||
|
Original question: {q}\n\n\
|
||||||
|
Write a CORRECTED SQL query using ONLY the columns listed in the schema. Output ONLY SQL."
|
||||||
|
);
|
||||||
|
if let Ok(retry_resp) = ai_generate(&retry_prompt, 512).await {
|
||||||
|
let retry_sql = clean_sql(&retry_resp.text);
|
||||||
|
generated_sql.set(Some(retry_sql.clone()));
|
||||||
|
step.set("running corrected query...".into());
|
||||||
|
let retry_result = run_sql(&retry_sql).await;
|
||||||
|
result.set(Some(retry_result));
|
||||||
|
} else {
|
||||||
|
result.set(Some(query_result));
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
result.set(Some(query_result));
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
result.set(Some(query_result));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
result.set(Some(Err(format!("AI error: {e}"))));
|
result.set(Some(Err(format!("AI error: {e}"))));
|
||||||
@@ -295,21 +373,44 @@ fn AskPanel(datasets: Vec<Dataset>) -> Element {
|
|||||||
generated_sql.set(None);
|
generated_sql.set(None);
|
||||||
result.set(None);
|
result.set(None);
|
||||||
step.set("reading schemas...".into());
|
step.set("reading schemas...".into());
|
||||||
let schema_ctx = get_schema_context(&ds).await;
|
let schema_ctx = get_schema_context_cached(&ds).await;
|
||||||
step.set("writing SQL...".into());
|
step.set("writing SQL...".into());
|
||||||
let prompt = format!(
|
let prompt = format!(
|
||||||
"You are a SQL assistant for a data lakehouse using Apache DataFusion (PostgreSQL-compatible SQL).\n\n\
|
"You are a SQL assistant. You write Apache DataFusion SQL (PostgreSQL-compatible).\n\n\
|
||||||
|
CRITICAL: You MUST only use column names that appear in the schema below. Do NOT invent or guess column names.\n\n\
|
||||||
{schema_ctx}\n\
|
{schema_ctx}\n\
|
||||||
User question: {q}\n\n\
|
User question: {q}\n\n\
|
||||||
Write a SQL query that answers this question. Output ONLY the SQL query, nothing else. No markdown, no explanation, no backticks."
|
Rules:\n\
|
||||||
|
- ONLY use table and column names listed above\n\
|
||||||
|
- If unsure about a column name, pick the closest match from the schema\n\
|
||||||
|
- Output ONLY the SQL query. No markdown, no explanation, no backticks."
|
||||||
);
|
);
|
||||||
match ai_generate(&prompt, 512).await {
|
match ai_generate(&prompt, 512).await {
|
||||||
Ok(resp) => {
|
Ok(resp) => {
|
||||||
let sql = resp.text.trim().to_string();
|
let sql = clean_sql(&resp.text);
|
||||||
generated_sql.set(Some(sql.clone()));
|
generated_sql.set(Some(sql.clone()));
|
||||||
step.set("running query...".into());
|
step.set("running query...".into());
|
||||||
let query_result = run_sql(&sql).await;
|
let query_result = run_sql(&sql).await;
|
||||||
result.set(Some(query_result));
|
if let Err(ref err) = query_result {
|
||||||
|
if err.contains("error") {
|
||||||
|
step.set("fixing SQL...".into());
|
||||||
|
let retry_prompt = format!(
|
||||||
|
"The SQL you wrote had an error:\n{err}\n\n{schema_ctx}\n\nOriginal question: {q}\n\nWrite a CORRECTED SQL query using ONLY the columns listed. Output ONLY SQL."
|
||||||
|
);
|
||||||
|
if let Ok(rr) = ai_generate(&retry_prompt, 512).await {
|
||||||
|
let rsql = clean_sql(&rr.text);
|
||||||
|
generated_sql.set(Some(rsql.clone()));
|
||||||
|
step.set("running corrected query...".into());
|
||||||
|
result.set(Some(run_sql(&rsql).await));
|
||||||
|
} else {
|
||||||
|
result.set(Some(query_result));
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
result.set(Some(query_result));
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
result.set(Some(query_result));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
result.set(Some(Err(format!("AI error: {e}"))));
|
result.set(Some(Err(format!("AI error: {e}"))));
|
||||||
@@ -532,6 +633,303 @@ fn SqlPanel() -> Element {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// === DASHBOARD — System overview ===
|
||||||
|
|
||||||
|
#[component]
|
||||||
|
fn DashboardPanel() -> Element {
|
||||||
|
let mut stats = use_signal(|| None::<serde_json::Value>);
|
||||||
|
let mut loading = use_signal(|| false);
|
||||||
|
|
||||||
|
let do_load = move || {
|
||||||
|
spawn(async move {
|
||||||
|
loading.set(true);
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
|
||||||
|
let datasets = fetch_datasets().await.ok().map(|d| d.len()).unwrap_or(0);
|
||||||
|
let row_count = run_sql("SELECT (SELECT COUNT(*) FROM candidates) + (SELECT COUNT(*) FROM clients) + (SELECT COUNT(*) FROM job_orders) + (SELECT COUNT(*) FROM placements) + (SELECT COUNT(*) FROM timesheets) + (SELECT COUNT(*) FROM call_log) + (SELECT COUNT(*) FROM email_log) as total").await.ok().and_then(|r| r.rows.as_array()?.first()?.get("total")?.as_i64()).unwrap_or(0);
|
||||||
|
|
||||||
|
let cache: Option<serde_json::Value> = match client.get(&format!("{}/query/cache/stats", api_base())).send().await {
|
||||||
|
Ok(r) => r.json().await.ok(),
|
||||||
|
Err(_) => None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let indexes: Vec<serde_json::Value> = match client.get(&format!("{}/vectors/indexes", api_base())).send().await {
|
||||||
|
Ok(r) => r.json().await.unwrap_or_default(),
|
||||||
|
Err(_) => vec![],
|
||||||
|
};
|
||||||
|
|
||||||
|
let tools: Vec<serde_json::Value> = match client.get(&format!("{}/tools", api_base())).send().await {
|
||||||
|
Ok(r) => r.json().await.unwrap_or_default(),
|
||||||
|
Err(_) => vec![],
|
||||||
|
};
|
||||||
|
|
||||||
|
let jobs: Vec<serde_json::Value> = match client.get(&format!("{}/vectors/jobs", api_base())).send().await {
|
||||||
|
Ok(r) => r.json().await.unwrap_or_default(),
|
||||||
|
Err(_) => vec![],
|
||||||
|
};
|
||||||
|
|
||||||
|
let hnsw: Vec<String> = match client.get(&format!("{}/vectors/hnsw/list", api_base())).send().await {
|
||||||
|
Ok(r) => r.json().await.unwrap_or_default(),
|
||||||
|
Err(_) => vec![],
|
||||||
|
};
|
||||||
|
|
||||||
|
let journal: Option<serde_json::Value> = match client.get(&format!("{}/journal/stats", api_base())).send().await {
|
||||||
|
Ok(r) => r.json().await.ok(),
|
||||||
|
Err(_) => None,
|
||||||
|
};
|
||||||
|
|
||||||
|
stats.set(Some(serde_json::json!({
|
||||||
|
"datasets": datasets,
|
||||||
|
"total_rows": row_count,
|
||||||
|
"cache": cache,
|
||||||
|
"vector_indexes": indexes.len(),
|
||||||
|
"vector_total_chunks": indexes.iter().filter_map(|i| i.get("chunk_count")?.as_i64()).sum::<i64>(),
|
||||||
|
"hnsw_loaded": hnsw.len(),
|
||||||
|
"tools": tools.len(),
|
||||||
|
"jobs_total": jobs.len(),
|
||||||
|
"jobs_running": jobs.iter().filter(|j| j.get("status").and_then(|s| s.as_str()) == Some("running")).count(),
|
||||||
|
"journal": journal,
|
||||||
|
})));
|
||||||
|
loading.set(false);
|
||||||
|
});
|
||||||
|
};
|
||||||
|
|
||||||
|
// Auto-load on mount
|
||||||
|
// Load on first render (user clicks Dashboard tab)
|
||||||
|
let mut loaded = use_signal(|| false);
|
||||||
|
use_effect(move || {
|
||||||
|
if !*loaded.read() {
|
||||||
|
loaded.set(true);
|
||||||
|
do_load();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
rsx! {
|
||||||
|
div { class: "panel",
|
||||||
|
div { class: "dashboard-hero",
|
||||||
|
h2 { "Lakehouse" }
|
||||||
|
p { class: "subtitle", "Rust-first data platform — SQL + AI + RAG over object storage" }
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(s) = stats.read().as_ref() {
|
||||||
|
div { class: "stat-grid",
|
||||||
|
div { class: "stat-card",
|
||||||
|
div { class: "stat-value", "{s[\"datasets\"]}" }
|
||||||
|
div { class: "stat-label", "Datasets" }
|
||||||
|
}
|
||||||
|
div { class: "stat-card",
|
||||||
|
div { class: "stat-value", {format!("{:.1}M", s["total_rows"].as_i64().unwrap_or(0) as f64 / 1_000_000.0)} }
|
||||||
|
div { class: "stat-label", "Total Rows" }
|
||||||
|
}
|
||||||
|
div { class: "stat-card",
|
||||||
|
div { class: "stat-value", {format!("{}K", s["vector_total_chunks"].as_i64().unwrap_or(0) / 1000)} }
|
||||||
|
div { class: "stat-label", "Embeddings" }
|
||||||
|
}
|
||||||
|
div { class: "stat-card accent",
|
||||||
|
div { class: "stat-value", "{s[\"hnsw_loaded\"]}" }
|
||||||
|
div { class: "stat-label", "HNSW Indexes (27ms)" }
|
||||||
|
}
|
||||||
|
div { class: "stat-card",
|
||||||
|
div { class: "stat-value", "{s[\"tools\"]}" }
|
||||||
|
div { class: "stat-label", "Governed Tools" }
|
||||||
|
}
|
||||||
|
div { class: "stat-card",
|
||||||
|
div { class: "stat-value", {
|
||||||
|
s.get("cache")
|
||||||
|
.and_then(|c| c.get("datasets"))
|
||||||
|
.and_then(|d| d.as_i64())
|
||||||
|
.map(|n| format!("{n}"))
|
||||||
|
.unwrap_or("0".into())
|
||||||
|
} }
|
||||||
|
div { class: "stat-label", "Cached Datasets" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
div { class: "architecture-section",
|
||||||
|
h3 { "Architecture" }
|
||||||
|
div { class: "arch-grid",
|
||||||
|
div { class: "arch-card",
|
||||||
|
div { class: "arch-title", "Ingest" }
|
||||||
|
div { class: "arch-items", "CSV, JSON, PDF, Text, PostgreSQL, File Watcher" }
|
||||||
|
}
|
||||||
|
div { class: "arch-card",
|
||||||
|
div { class: "arch-title", "Storage" }
|
||||||
|
div { class: "arch-items", "Parquet on Object Storage, Delta Writes, Compaction" }
|
||||||
|
}
|
||||||
|
div { class: "arch-card",
|
||||||
|
div { class: "arch-title", "Query" }
|
||||||
|
div { class: "arch-items", "DataFusion SQL, MemCache (9.8x), Hot/Cold" }
|
||||||
|
}
|
||||||
|
div { class: "arch-card",
|
||||||
|
div { class: "arch-title", "AI" }
|
||||||
|
div { class: "arch-items", "Ollama (local), Embed, Generate, RAG, HNSW" }
|
||||||
|
}
|
||||||
|
div { class: "arch-card",
|
||||||
|
div { class: "arch-title", "Governance" }
|
||||||
|
div { class: "arch-items", "Event Journal, PII Detection, Tool Registry, Access Control" }
|
||||||
|
}
|
||||||
|
div { class: "arch-card",
|
||||||
|
div { class: "arch-title", "Agents" }
|
||||||
|
div { class: "arch-items", "Workspaces, Handoff, Shortlists, Activity Logs" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
div { class: "phases-section",
|
||||||
|
h3 { "Build Progression" }
|
||||||
|
div { class: "phase-list",
|
||||||
|
{rsx! {
|
||||||
|
PhaseItem { num: "0-5", name: "Foundation", detail: "Storage, Catalog, DataFusion, AI, UI, gRPC" }
|
||||||
|
PhaseItem { num: "6", name: "Ingest Pipeline", detail: "CSV/JSON/PDF/Text auto-schema" }
|
||||||
|
PhaseItem { num: "7", name: "Vector + RAG", detail: "Embed, Search, LLM Answers" }
|
||||||
|
PhaseItem { num: "8", name: "Hot Cache", detail: "9.8x speedup, Delta Writes" }
|
||||||
|
PhaseItem { num: "8.5", name: "Agent Workspaces", detail: "Per-contract, Instant Handoff" }
|
||||||
|
PhaseItem { num: "9", name: "Event Journal", detail: "Append-only Mutation History" }
|
||||||
|
PhaseItem { num: "10", name: "Rich Catalog", detail: "PII Detection, Lineage" }
|
||||||
|
PhaseItem { num: "11", name: "Embedding Versioning", detail: "Model-proof Vectors" }
|
||||||
|
PhaseItem { num: "12", name: "Tool Registry", detail: "6 Governed Actions + Audit" }
|
||||||
|
PhaseItem { num: "13", name: "Access Control", detail: "Role-based, Field-level" }
|
||||||
|
PhaseItem { num: "14", name: "Schema Evolution", detail: "Diff Detection, AI Migration" }
|
||||||
|
PhaseItem { num: "15", name: "HNSW Index", detail: "100K Search in 27ms" }
|
||||||
|
PhaseItem { num: "16", name: "File Watcher", detail: "Auto-ingest from Inbox" }
|
||||||
|
PhaseItem { num: "17", name: "DB Connector", detail: "PostgreSQL Import" }
|
||||||
|
}}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else if *loading.read() {
|
||||||
|
div { class: "loading", "loading dashboard..." }
|
||||||
|
}
|
||||||
|
|
||||||
|
button { class: "btn", onclick: move |_| do_load(), "Refresh" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[component]
|
||||||
|
fn PhaseItem(num: String, name: String, detail: String) -> Element {
|
||||||
|
rsx! {
|
||||||
|
div { class: "phase-item",
|
||||||
|
span { class: "phase-num", "{num}" }
|
||||||
|
span { class: "phase-name", "{name}" }
|
||||||
|
span { class: "phase-detail", "{detail}" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// === INGEST — Data on-ramp ===
|
||||||
|
|
||||||
|
#[component]
|
||||||
|
fn IngestPanel() -> Element {
|
||||||
|
let mut pg_host = use_signal(|| "127.0.0.1".to_string());
|
||||||
|
let mut pg_db = use_signal(|| "knowledge_base".to_string());
|
||||||
|
let mut pg_tables = use_signal(|| None::<Vec<String>>);
|
||||||
|
let mut pg_result = use_signal(|| None::<Result<serde_json::Value, String>>);
|
||||||
|
let mut pg_loading = use_signal(|| false);
|
||||||
|
|
||||||
|
let list_tables = move |_| {
|
||||||
|
let host = pg_host.read().clone();
|
||||||
|
let db = pg_db.read().clone();
|
||||||
|
spawn(async move {
|
||||||
|
pg_loading.set(true);
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let resp = client.post(&format!("{}/ingest/postgres/tables", api_base()))
|
||||||
|
.json(&serde_json::json!({"host": host, "port": 5432, "database": db, "user": "postgres", "password": ""}))
|
||||||
|
.send().await;
|
||||||
|
match resp {
|
||||||
|
Ok(r) => {
|
||||||
|
if let Ok(tables) = r.json::<Vec<String>>().await {
|
||||||
|
pg_tables.set(Some(tables));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => pg_tables.set(None),
|
||||||
|
}
|
||||||
|
pg_loading.set(false);
|
||||||
|
});
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut import_table = move |table: String| {
|
||||||
|
let host = pg_host.read().clone();
|
||||||
|
let db = pg_db.read().clone();
|
||||||
|
spawn(async move {
|
||||||
|
pg_result.set(None);
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let resp = client.post(&format!("{}/ingest/postgres/import", api_base()))
|
||||||
|
.json(&serde_json::json!({"host": host, "port": 5432, "database": db, "user": "postgres", "password": "", "table": table}))
|
||||||
|
.send().await;
|
||||||
|
match resp {
|
||||||
|
Ok(r) => {
|
||||||
|
match r.json::<serde_json::Value>().await {
|
||||||
|
Ok(v) => pg_result.set(Some(Ok(v))),
|
||||||
|
Err(e) => pg_result.set(Some(Err(e.to_string()))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => pg_result.set(Some(Err(e.to_string()))),
|
||||||
|
}
|
||||||
|
});
|
||||||
|
};
|
||||||
|
|
||||||
|
rsx! {
|
||||||
|
div { class: "panel",
|
||||||
|
h2 { "Data Ingest" }
|
||||||
|
p { class: "subtitle", "Bring data in from files, databases, or the auto-watch inbox" }
|
||||||
|
|
||||||
|
div { class: "panel-section",
|
||||||
|
h3 { "File Upload" }
|
||||||
|
p { class: "hint", "POST a file to /ingest/file — or drop it in ./inbox/ for auto-ingest" }
|
||||||
|
div { class: "arch-card",
|
||||||
|
div { class: "arch-title", "Supported Formats" }
|
||||||
|
div { class: "arch-items", "CSV (auto-schema) | JSON (nested flattening) | PDF (text extraction) | Text/SMS" }
|
||||||
|
}
|
||||||
|
div { class: "arch-card",
|
||||||
|
div { class: "arch-title", "Auto-Watch Inbox" }
|
||||||
|
div { class: "arch-items", "Drop files in ./inbox/ → auto-detected → Parquet → queryable in <15s" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
div { class: "panel-section",
|
||||||
|
h3 { "PostgreSQL Import" }
|
||||||
|
div { class: "form-row",
|
||||||
|
label { "Host" }
|
||||||
|
input { value: "{pg_host}", oninput: move |e| pg_host.set(e.value()) }
|
||||||
|
}
|
||||||
|
div { class: "form-row",
|
||||||
|
label { "Database" }
|
||||||
|
input { value: "{pg_db}", oninput: move |e| pg_db.set(e.value()) }
|
||||||
|
}
|
||||||
|
button { class: "btn", disabled: *pg_loading.read(), onclick: list_tables, "List Tables" }
|
||||||
|
|
||||||
|
if let Some(tables) = pg_tables.read().as_ref() {
|
||||||
|
div { class: "table-list",
|
||||||
|
for table in tables.iter() {
|
||||||
|
{
|
||||||
|
let t = table.clone();
|
||||||
|
let t2 = table.clone();
|
||||||
|
rsx! {
|
||||||
|
div { class: "table-item",
|
||||||
|
span { "{t}" }
|
||||||
|
button { class: "btn btn-sm", onclick: move |_| import_table(t2.clone()), "Import" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(result) = pg_result.read().as_ref() {
|
||||||
|
match result {
|
||||||
|
Ok(v) => rsx! {
|
||||||
|
div { class: "result-box",
|
||||||
|
pre { {serde_json::to_string_pretty(v).unwrap_or_default()} }
|
||||||
|
}
|
||||||
|
},
|
||||||
|
Err(e) => rsx! { div { class: "error", "{e}" } },
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// === STATUS — System health ===
|
// === STATUS — System health ===
|
||||||
|
|
||||||
#[component]
|
#[component]
|
||||||
@@ -633,6 +1031,9 @@ fn ResultsTable(response: QueryResponse) -> Element {
|
|||||||
if response.row_count == 0 {
|
if response.row_count == 0 {
|
||||||
div { class: "empty-sm", "no rows returned" }
|
div { class: "empty-sm", "no rows returned" }
|
||||||
} else if let Some(rows) = rows {
|
} else if let Some(rows) = rows {
|
||||||
|
if rows.len() > 500 {
|
||||||
|
div { class: "results-info", "Showing first 500 of {response.row_count} rows (use SQL tab with LIMIT for larger)" }
|
||||||
|
}
|
||||||
div { class: "table-wrap",
|
div { class: "table-wrap",
|
||||||
table {
|
table {
|
||||||
thead {
|
thead {
|
||||||
@@ -643,7 +1044,7 @@ fn ResultsTable(response: QueryResponse) -> Element {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
tbody {
|
tbody {
|
||||||
for row in rows.iter() {
|
for row in rows.iter().take(500) {
|
||||||
tr {
|
tr {
|
||||||
for col in response.columns.iter() {
|
for col in response.columns.iter() {
|
||||||
td { {format_cell(row.get(&col.name))} }
|
td { {format_cell(row.get(&col.name))} }
|
||||||
@@ -657,6 +1058,49 @@ fn ResultsTable(response: QueryResponse) -> Element {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Clean AI-generated SQL: extract only the SQL query, strip everything else.
|
||||||
|
fn clean_sql(raw: &str) -> String {
|
||||||
|
let s = raw.trim();
|
||||||
|
|
||||||
|
// Strategy 1: If there's a ```sql...``` block, extract just that
|
||||||
|
if let Some(start) = s.find("```sql") {
|
||||||
|
let after = &s[start + 6..];
|
||||||
|
if let Some(end) = after.find("```") {
|
||||||
|
return after[..end].trim().to_string();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if let Some(start) = s.find("```") {
|
||||||
|
let after = &s[start + 3..];
|
||||||
|
if let Some(end) = after.find("```") {
|
||||||
|
let inner = after[..end].trim();
|
||||||
|
// Skip leading "sql" keyword
|
||||||
|
let inner = inner.strip_prefix("sql").map(|s| s.trim_start()).unwrap_or(inner);
|
||||||
|
return inner.to_string();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Strategy 2: Find the first SELECT/WITH/INSERT/UPDATE/DELETE statement
|
||||||
|
let upper = s.to_uppercase();
|
||||||
|
for keyword in &["SELECT", "WITH", "INSERT", "UPDATE", "DELETE"] {
|
||||||
|
if let Some(pos) = upper.find(keyword) {
|
||||||
|
let sql_part = &s[pos..];
|
||||||
|
// Take up to the first semicolon (or end)
|
||||||
|
let end = sql_part.find(';').map(|p| p + 1).unwrap_or(sql_part.len());
|
||||||
|
return sql_part[..end].trim().to_string();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Strategy 3: Strip leading "sql" and clean up
|
||||||
|
let mut result = s.to_string();
|
||||||
|
let lines: Vec<&str> = result.lines().collect();
|
||||||
|
if let Some(first) = lines.first() {
|
||||||
|
if first.trim().eq_ignore_ascii_case("sql") {
|
||||||
|
result = lines[1..].join("\n").trim().to_string();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
fn format_cell(val: Option<&serde_json::Value>) -> String {
|
fn format_cell(val: Option<&serde_json::Value>) -> String {
|
||||||
match val {
|
match val {
|
||||||
None | Some(serde_json::Value::Null) => "—".to_string(),
|
None | Some(serde_json::Value::Null) => "—".to_string(),
|
||||||
|
|||||||
@@ -0,0 +1,37 @@
|
|||||||
|
[package]
|
||||||
|
name = "vectord-lance"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2024"
|
||||||
|
|
||||||
|
# Firewall crate: the Lance stack (Arrow 57, DataFusion 52) is isolated
|
||||||
|
# from the rest of vectord (Arrow 55, DataFusion 47). Public API uses
|
||||||
|
# only std types — no Arrow types cross the crate boundary — so no
|
||||||
|
# version conflict propagates outward.
|
||||||
|
#
|
||||||
|
# See docs/ADR-019-vector-storage.md for the rationale: "productionizing
|
||||||
|
# will need either workspace-wide upgrade or a firewall via a dedicated
|
||||||
|
# vectord-lance crate." This is that firewall.
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
# S3 support: Lance delegates to its internal object_store crate when
|
||||||
|
# given s3:// URIs. The "dynamodb" feature enables DynamoDB-based
|
||||||
|
# commit locking for multi-writer S3; we don't need that (single-writer)
|
||||||
|
# so just the base AWS/S3 feature is enough.
|
||||||
|
# Lance 4.0 feature "aws" enables S3-compatible storage via its internal
|
||||||
|
# object_store + opendal crates. Reads AWS_* env vars for credentials.
|
||||||
|
lance = { version = "4.0", default-features = false, features = ["aws"] }
|
||||||
|
lance-index = { version = "4.0", default-features = false }
|
||||||
|
lance-linalg = { version = "4.0", default-features = false }
|
||||||
|
|
||||||
|
# These Arrow/Parquet versions MUST match Lance's expectations — Lance
|
||||||
|
# re-exports their types through its API so any mismatch is a hard
|
||||||
|
# compile error. Keep in sync with lance-bench.
|
||||||
|
arrow = "57"
|
||||||
|
arrow-array = "57"
|
||||||
|
arrow-schema = "57"
|
||||||
|
parquet = "57"
|
||||||
|
|
||||||
|
tokio = { version = "1", features = ["full"] }
|
||||||
|
futures = "0.3"
|
||||||
|
serde = { version = "1", features = ["derive"] }
|
||||||
|
bytes = "1"
|
||||||
@@ -0,0 +1,585 @@
|
|||||||
|
//! Production Lance vector backend (ADR-019 — hybrid architecture).
|
||||||
|
//!
|
||||||
|
//! This is the firewall crate. It owns its own Arrow 57 / DataFusion 52
|
||||||
|
//! / Lance 4 dependency tree. The public API uses only std types
|
||||||
|
//! (`Vec<f32>`, `Vec<String>`, `String`, `bool`) so nothing Arrow-shaped
|
||||||
|
//! crosses the crate boundary. That keeps vectord (Arrow 55) from
|
||||||
|
//! picking up an incompatible dep.
|
||||||
|
//!
|
||||||
|
//! Responsibilities:
|
||||||
|
//! - Migrate an existing binary-blob Parquet vector file into a Lance
|
||||||
|
//! dataset with `FixedSizeList<Float32, dims>`. One-time cost.
|
||||||
|
//! - Append new rows natively (no full-file rewrite — Lance's structural win).
|
||||||
|
//! - Build an IVF_PQ ANN index on the vector column.
|
||||||
|
//! - Vector search (`search`) using the IVF_PQ index when present,
|
||||||
|
//! falling back to full scan otherwise.
|
||||||
|
//! - Random-access row fetch by `doc_id` (`get_by_doc_id`) — the O(1)
|
||||||
|
//! lookup that Parquet-on-object-store can't cheaply do.
|
||||||
|
//! - Cheap count + stats introspection.
|
||||||
|
|
||||||
|
use arrow_array::{
|
||||||
|
Array, ArrayRef, BinaryArray, FixedSizeListArray, Float32Array, Int32Array,
|
||||||
|
RecordBatch, RecordBatchIterator, StringArray,
|
||||||
|
};
|
||||||
|
use arrow_schema::{DataType, Field, Schema};
|
||||||
|
use futures::StreamExt;
|
||||||
|
use serde::Serialize;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
// ================= Public types =================
|
||||||
|
|
||||||
|
/// One search result. Mirrors vectord's existing `SearchResult` shape
|
||||||
|
/// structurally but carries simpler types so this crate stays firewalled.
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct Hit {
|
||||||
|
pub doc_id: String,
|
||||||
|
pub chunk_text: String,
|
||||||
|
pub score: f32,
|
||||||
|
/// Optional — set by search_with_vector, not by index-only search.
|
||||||
|
pub distance: Option<f32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A fully-hydrated row fetched by `get_by_doc_id` — includes the vector
|
||||||
|
/// so callers can do downstream work (rerank, cite, etc.) without a
|
||||||
|
/// second round trip.
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct Row {
|
||||||
|
pub doc_id: String,
|
||||||
|
pub chunk_text: String,
|
||||||
|
pub vector: Vec<f32>,
|
||||||
|
pub source: Option<String>,
|
||||||
|
pub chunk_idx: Option<i32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct MigrationStats {
|
||||||
|
pub rows_written: usize,
|
||||||
|
pub dimensions: usize,
|
||||||
|
pub disk_bytes: u64,
|
||||||
|
pub duration_secs: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct AppendStats {
|
||||||
|
pub rows_appended: usize,
|
||||||
|
pub disk_bytes_added: u64,
|
||||||
|
pub duration_secs: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct IndexStats {
|
||||||
|
pub name: String,
|
||||||
|
pub num_partitions: u32,
|
||||||
|
pub num_bits: u32,
|
||||||
|
pub num_sub_vectors: u32,
|
||||||
|
pub build_time_secs: f32,
|
||||||
|
pub disk_bytes_added: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct ScalarIndexStats {
|
||||||
|
pub name: String,
|
||||||
|
pub column: String,
|
||||||
|
pub build_time_secs: f32,
|
||||||
|
pub disk_bytes_added: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct DatasetStats {
|
||||||
|
pub path: String,
|
||||||
|
pub rows: usize,
|
||||||
|
pub disk_bytes: u64,
|
||||||
|
pub has_vector_index: bool,
|
||||||
|
pub has_doc_id_index: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ================= The backend =================
|
||||||
|
|
||||||
|
/// Thin wrapper around a Lance dataset path. Lance handles the heavy
|
||||||
|
/// lifting — we just expose a narrow API.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct LanceVectorStore {
|
||||||
|
/// Local filesystem path or object-store URI (file:///..., s3://...).
|
||||||
|
/// Lance's internal URI parsing handles both.
|
||||||
|
path: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl LanceVectorStore {
|
||||||
|
pub fn new(path: impl Into<String>) -> Self {
|
||||||
|
Self { path: path.into() }
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn path(&self) -> &str { &self.path }
|
||||||
|
|
||||||
|
/// Row count via Lance's fast metadata path (no scan).
|
||||||
|
pub async fn count(&self) -> Result<usize, String> {
|
||||||
|
let dataset = open_or_err(&self.path).await?;
|
||||||
|
let n = dataset.count_rows(None).await.map_err(e)?;
|
||||||
|
Ok(n)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// True if the on-disk dataset exists AND has at least one `vector`
|
||||||
|
/// column index attached.
|
||||||
|
pub async fn has_vector_index(&self) -> Result<bool, String> {
|
||||||
|
use lance_index::DatasetIndexExt;
|
||||||
|
let dataset = match lance::dataset::Dataset::open(&self.path).await {
|
||||||
|
Ok(d) => d,
|
||||||
|
Err(_) => return Ok(false),
|
||||||
|
};
|
||||||
|
let indexes = dataset.load_indices().await.map_err(e)?;
|
||||||
|
Ok(indexes.iter().any(|ix| {
|
||||||
|
ix.fields.iter().any(|fid| {
|
||||||
|
dataset.schema().field_by_id(*fid)
|
||||||
|
.map(|f| f.name == "vector")
|
||||||
|
.unwrap_or(false)
|
||||||
|
})
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn stats(&self) -> Result<DatasetStats, String> {
|
||||||
|
let rows = self.count().await.unwrap_or(0);
|
||||||
|
let disk_bytes = dir_size_bytes(&strip_file_uri(&self.path));
|
||||||
|
let has_vector_index = self.has_vector_index().await.unwrap_or(false);
|
||||||
|
let has_doc_id_index = self.has_scalar_index("doc_id").await.unwrap_or(false);
|
||||||
|
Ok(DatasetStats {
|
||||||
|
path: self.path.clone(),
|
||||||
|
rows,
|
||||||
|
disk_bytes,
|
||||||
|
has_vector_index,
|
||||||
|
has_doc_id_index,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Migrate a vectord-format Parquet file into a Lance dataset.
|
||||||
|
///
|
||||||
|
/// Input schema (vectord's binary-blob format):
|
||||||
|
/// - source : Utf8
|
||||||
|
/// - doc_id : Utf8
|
||||||
|
/// - chunk_idx : Int32
|
||||||
|
/// - chunk_text : Utf8
|
||||||
|
/// - vector : Binary (raw f32 little-endian bytes)
|
||||||
|
///
|
||||||
|
/// Output schema (Lance-friendly):
|
||||||
|
/// - source : Utf8
|
||||||
|
/// - doc_id : Utf8
|
||||||
|
/// - chunk_idx : Int32
|
||||||
|
/// - chunk_text : Utf8
|
||||||
|
/// - vector : FixedSizeList<Float32, dims>
|
||||||
|
///
|
||||||
|
/// Idempotent at the file level — if the target exists, it's
|
||||||
|
/// overwritten. Caller must manage destination paths.
|
||||||
|
pub async fn migrate_from_parquet_bytes(
|
||||||
|
&self,
|
||||||
|
parquet_bytes: &[u8],
|
||||||
|
) -> Result<MigrationStats, String> {
|
||||||
|
let t0 = Instant::now();
|
||||||
|
let (schema, batches, rows) = read_parquet(parquet_bytes)?;
|
||||||
|
let dims = detect_vector_dims(&batches)?;
|
||||||
|
let (new_schema, new_batches) = convert_to_fixed_size_list(&schema, batches, dims)?;
|
||||||
|
|
||||||
|
// Overwrite any prior dataset at this path.
|
||||||
|
let _ = std::fs::remove_dir_all(&strip_file_uri(&self.path));
|
||||||
|
|
||||||
|
write_dataset(&self.path, new_schema, new_batches).await?;
|
||||||
|
let disk_bytes = dir_size_bytes(&strip_file_uri(&self.path));
|
||||||
|
|
||||||
|
Ok(MigrationStats {
|
||||||
|
rows_written: rows,
|
||||||
|
dimensions: dims,
|
||||||
|
disk_bytes,
|
||||||
|
duration_secs: t0.elapsed().as_secs_f32(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Native Lance append — does NOT rewrite existing files. New rows
|
||||||
|
/// land as a separate fragment; readers union across fragments at
|
||||||
|
/// query time. Contrast: our Parquet path requires rewriting the
|
||||||
|
/// entire vector file to add rows.
|
||||||
|
pub async fn append(
|
||||||
|
&self,
|
||||||
|
source: Option<String>,
|
||||||
|
doc_ids: Vec<String>,
|
||||||
|
chunk_idxs: Vec<i32>,
|
||||||
|
chunk_texts: Vec<String>,
|
||||||
|
vectors: Vec<Vec<f32>>,
|
||||||
|
) -> Result<AppendStats, String> {
|
||||||
|
let n = doc_ids.len();
|
||||||
|
if n == 0 {
|
||||||
|
return Ok(AppendStats { rows_appended: 0, disk_bytes_added: 0, duration_secs: 0.0 });
|
||||||
|
}
|
||||||
|
if chunk_idxs.len() != n || chunk_texts.len() != n || vectors.len() != n {
|
||||||
|
return Err(format!(
|
||||||
|
"append: length mismatch (doc_ids={n}, chunk_idxs={}, chunk_texts={}, vectors={})",
|
||||||
|
chunk_idxs.len(), chunk_texts.len(), vectors.len(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let dims = vectors[0].len();
|
||||||
|
for (i, v) in vectors.iter().enumerate() {
|
||||||
|
if v.len() != dims {
|
||||||
|
return Err(format!("append: row {i} has {} dims, expected {}", v.len(), dims));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let t0 = Instant::now();
|
||||||
|
let pre_bytes = dir_size_bytes(&strip_file_uri(&self.path));
|
||||||
|
|
||||||
|
let src_arr = StringArray::from(
|
||||||
|
(0..n).map(|_| source.clone()).collect::<Vec<_>>()
|
||||||
|
);
|
||||||
|
let doc_id_arr = StringArray::from(doc_ids);
|
||||||
|
let chunk_idx_arr = Int32Array::from(chunk_idxs);
|
||||||
|
let chunk_text_arr = StringArray::from(chunk_texts);
|
||||||
|
|
||||||
|
let mut flat: Vec<f32> = Vec::with_capacity(n * dims);
|
||||||
|
for v in vectors { flat.extend(v); }
|
||||||
|
let values = Float32Array::from(flat);
|
||||||
|
let item_field = Arc::new(Field::new("item", DataType::Float32, true));
|
||||||
|
let vec_arr = FixedSizeListArray::try_new(
|
||||||
|
item_field.clone(), dims as i32, Arc::new(values), None,
|
||||||
|
).map_err(e)?;
|
||||||
|
|
||||||
|
let schema = Arc::new(Schema::new(vec![
|
||||||
|
Field::new("source", DataType::Utf8, true),
|
||||||
|
Field::new("doc_id", DataType::Utf8, false),
|
||||||
|
Field::new("chunk_idx", DataType::Int32, true),
|
||||||
|
Field::new("chunk_text", DataType::Utf8, true),
|
||||||
|
Field::new("vector", DataType::FixedSizeList(item_field, dims as i32), false),
|
||||||
|
]));
|
||||||
|
|
||||||
|
let arrays: Vec<ArrayRef> = vec![
|
||||||
|
Arc::new(src_arr), Arc::new(doc_id_arr), Arc::new(chunk_idx_arr),
|
||||||
|
Arc::new(chunk_text_arr), Arc::new(vec_arr),
|
||||||
|
];
|
||||||
|
let batch = RecordBatch::try_new(schema.clone(), arrays).map_err(e)?;
|
||||||
|
let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema);
|
||||||
|
|
||||||
|
use lance::dataset::{Dataset, WriteMode, WriteParams};
|
||||||
|
let params = WriteParams { mode: WriteMode::Append, ..Default::default() };
|
||||||
|
Dataset::write(reader, &self.path, Some(params)).await.map_err(e)?;
|
||||||
|
|
||||||
|
Ok(AppendStats {
|
||||||
|
rows_appended: n,
|
||||||
|
disk_bytes_added: dir_size_bytes(&strip_file_uri(&self.path)).saturating_sub(pre_bytes),
|
||||||
|
duration_secs: t0.elapsed().as_secs_f32(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build an IVF_PQ vector index. Replaces any prior index with the
|
||||||
|
/// same name. Callers pass explicit params — sensible defaults for
|
||||||
|
/// ~100K × 768d: num_partitions=316 (≈√N), num_bits=8, num_sub_vectors=48.
|
||||||
|
pub async fn build_index(
|
||||||
|
&self,
|
||||||
|
num_partitions: u32,
|
||||||
|
num_bits: u32,
|
||||||
|
num_sub_vectors: u32,
|
||||||
|
) -> Result<IndexStats, String> {
|
||||||
|
use lance::dataset::Dataset;
|
||||||
|
use lance::index::vector::VectorIndexParams;
|
||||||
|
use lance_index::{DatasetIndexExt, IndexType};
|
||||||
|
use lance_linalg::distance::MetricType;
|
||||||
|
|
||||||
|
let pre_bytes = dir_size_bytes(&strip_file_uri(&self.path));
|
||||||
|
let t0 = Instant::now();
|
||||||
|
|
||||||
|
let mut dataset = Dataset::open(&self.path).await.map_err(e)?;
|
||||||
|
let params = VectorIndexParams::ivf_pq(
|
||||||
|
num_partitions as usize,
|
||||||
|
num_bits as u8,
|
||||||
|
num_sub_vectors as usize,
|
||||||
|
MetricType::Cosine,
|
||||||
|
50, // max_iterations — same as bench
|
||||||
|
);
|
||||||
|
dataset.create_index(
|
||||||
|
&["vector"],
|
||||||
|
IndexType::Vector,
|
||||||
|
Some("vec_idx".into()),
|
||||||
|
¶ms,
|
||||||
|
true, // replace
|
||||||
|
).await.map_err(e)?;
|
||||||
|
|
||||||
|
Ok(IndexStats {
|
||||||
|
name: "vec_idx".into(),
|
||||||
|
num_partitions,
|
||||||
|
num_bits,
|
||||||
|
num_sub_vectors,
|
||||||
|
build_time_secs: t0.elapsed().as_secs_f32(),
|
||||||
|
disk_bytes_added: dir_size_bytes(&strip_file_uri(&self.path)).saturating_sub(pre_bytes),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a scalar btree index on `doc_id`. Makes `get_by_doc_id`
|
||||||
|
/// O(log N) instead of a filter-scan of every fragment. Small index
|
||||||
|
/// footprint — a few hundred KB on 100K rows.
|
||||||
|
pub async fn build_scalar_index(&self, column: &str) -> Result<ScalarIndexStats, String> {
|
||||||
|
use lance::dataset::Dataset;
|
||||||
|
use lance_index::{DatasetIndexExt, IndexType};
|
||||||
|
use lance_index::scalar::ScalarIndexParams;
|
||||||
|
|
||||||
|
let pre_bytes = dir_size_bytes(&strip_file_uri(&self.path));
|
||||||
|
let t0 = Instant::now();
|
||||||
|
|
||||||
|
let mut dataset = Dataset::open(&self.path).await.map_err(e)?;
|
||||||
|
let idx_name = format!("{column}_btree");
|
||||||
|
dataset.create_index(
|
||||||
|
&[column],
|
||||||
|
IndexType::Scalar,
|
||||||
|
Some(idx_name.clone()),
|
||||||
|
&ScalarIndexParams::default(),
|
||||||
|
true,
|
||||||
|
).await.map_err(e)?;
|
||||||
|
|
||||||
|
Ok(ScalarIndexStats {
|
||||||
|
name: idx_name,
|
||||||
|
column: column.to_string(),
|
||||||
|
build_time_secs: t0.elapsed().as_secs_f32(),
|
||||||
|
disk_bytes_added: dir_size_bytes(&strip_file_uri(&self.path)).saturating_sub(pre_bytes),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// True if a scalar index exists on the named column.
|
||||||
|
pub async fn has_scalar_index(&self, column: &str) -> Result<bool, String> {
|
||||||
|
use lance_index::DatasetIndexExt;
|
||||||
|
let dataset = match lance::dataset::Dataset::open(&self.path).await {
|
||||||
|
Ok(d) => d,
|
||||||
|
Err(_) => return Ok(false),
|
||||||
|
};
|
||||||
|
let indexes = dataset.load_indices().await.map_err(e)?;
|
||||||
|
Ok(indexes.iter().any(|ix| {
|
||||||
|
ix.fields.iter().any(|fid| {
|
||||||
|
dataset.schema().field_by_id(*fid)
|
||||||
|
.map(|f| f.name == column)
|
||||||
|
.unwrap_or(false)
|
||||||
|
})
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Search for top_k nearest neighbors of `query`. Uses the IVF_PQ
|
||||||
|
/// index if one exists; otherwise does a full scan (slow but
|
||||||
|
/// correct — useful during development before index build).
|
||||||
|
pub async fn search(&self, query: &[f32], top_k: usize) -> Result<Vec<Hit>, String> {
|
||||||
|
use lance::dataset::Dataset;
|
||||||
|
|
||||||
|
let dataset = Dataset::open(&self.path).await.map_err(e)?;
|
||||||
|
let qarr = Float32Array::from(query.to_vec());
|
||||||
|
|
||||||
|
let mut scanner = dataset.scan();
|
||||||
|
scanner.nearest("vector", &qarr, top_k as usize).map_err(e)?;
|
||||||
|
scanner.project(&["doc_id", "chunk_text"]).map_err(e)?;
|
||||||
|
|
||||||
|
let mut stream = scanner.try_into_stream().await.map_err(e)?;
|
||||||
|
let mut hits: Vec<Hit> = Vec::with_capacity(top_k);
|
||||||
|
|
||||||
|
while let Some(batch) = stream.next().await {
|
||||||
|
let batch = batch.map_err(e)?;
|
||||||
|
let doc_ids = batch.column_by_name("doc_id")
|
||||||
|
.ok_or_else(|| "no doc_id column in search result".to_string())?
|
||||||
|
.as_any().downcast_ref::<StringArray>()
|
||||||
|
.ok_or_else(|| "doc_id is not StringArray".to_string())?;
|
||||||
|
let chunk_texts = batch.column_by_name("chunk_text")
|
||||||
|
.and_then(|c| c.as_any().downcast_ref::<StringArray>());
|
||||||
|
// Lance tacks on a `_distance` column for nearest() queries.
|
||||||
|
let distances = batch.column_by_name("_distance")
|
||||||
|
.and_then(|c| c.as_any().downcast_ref::<Float32Array>());
|
||||||
|
for row in 0..batch.num_rows() {
|
||||||
|
let d = distances.map(|a| a.value(row));
|
||||||
|
hits.push(Hit {
|
||||||
|
doc_id: doc_ids.value(row).to_string(),
|
||||||
|
chunk_text: chunk_texts.map(|a| a.value(row).to_string()).unwrap_or_default(),
|
||||||
|
score: d.map(|d| 1.0 - d).unwrap_or(0.0), // cosine distance → similarity
|
||||||
|
distance: d,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(hits)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fetch one row by doc_id. Implementation: Lance filter-pushdown
|
||||||
|
/// scan — O(1) with partition pruning on a proper btree index,
|
||||||
|
/// O(N) on vector-only datasets (still far faster than reading
|
||||||
|
/// the whole Parquet file). We don't build a scalar index on
|
||||||
|
/// doc_id yet; that's a future optimization.
|
||||||
|
pub async fn get_by_doc_id(&self, doc_id: &str) -> Result<Option<Row>, String> {
|
||||||
|
use lance::dataset::Dataset;
|
||||||
|
|
||||||
|
let dataset = Dataset::open(&self.path).await.map_err(e)?;
|
||||||
|
let filter = format!("doc_id = '{}'", doc_id.replace('\'', "''"));
|
||||||
|
let mut scanner = dataset.scan();
|
||||||
|
scanner.filter(&filter).map_err(e)?;
|
||||||
|
scanner.limit(Some(1), None).map_err(e)?;
|
||||||
|
let mut stream = scanner.try_into_stream().await.map_err(e)?;
|
||||||
|
|
||||||
|
while let Some(batch) = stream.next().await {
|
||||||
|
let batch = batch.map_err(e)?;
|
||||||
|
if batch.num_rows() == 0 { continue; }
|
||||||
|
return Ok(Some(row_from_batch(&batch, 0)?));
|
||||||
|
}
|
||||||
|
Ok(None)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ================= Internal helpers =================
|
||||||
|
|
||||||
|
fn e<T: std::fmt::Display>(err: T) -> String { err.to_string() }
|
||||||
|
|
||||||
|
/// `file:///abs/path` → `/abs/path`. Leave other URI schemes as-is for
|
||||||
|
/// helpers that only work on local paths (dir_size_bytes, remove_dir_all).
|
||||||
|
fn strip_file_uri(uri: &str) -> String {
|
||||||
|
uri.strip_prefix("file://").unwrap_or(uri).to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn open_or_err(path: &str) -> Result<lance::dataset::Dataset, String> {
|
||||||
|
lance::dataset::Dataset::open(path).await.map_err(e)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn dir_size_bytes(path: &str) -> u64 {
|
||||||
|
fn recurse(p: &std::path::Path) -> u64 {
|
||||||
|
let Ok(meta) = std::fs::metadata(p) else { return 0; };
|
||||||
|
if meta.is_file() { return meta.len(); }
|
||||||
|
let Ok(entries) = std::fs::read_dir(p) else { return 0; };
|
||||||
|
entries.filter_map(|e| e.ok()).map(|e| recurse(&e.path())).sum()
|
||||||
|
}
|
||||||
|
recurse(std::path::Path::new(path))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_parquet(bytes: &[u8]) -> Result<(Arc<Schema>, Vec<RecordBatch>, usize), String> {
|
||||||
|
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
|
||||||
|
let builder = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::copy_from_slice(bytes))
|
||||||
|
.map_err(e)?;
|
||||||
|
let schema = builder.schema().clone();
|
||||||
|
let reader = builder.build().map_err(e)?;
|
||||||
|
let batches: Vec<RecordBatch> = reader.collect::<Result<_, _>>().map_err(e)?;
|
||||||
|
let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
|
||||||
|
Ok((schema, batches, rows))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn detect_vector_dims(batches: &[RecordBatch]) -> Result<usize, String> {
|
||||||
|
for batch in batches {
|
||||||
|
let idx = batch.schema().index_of("vector")
|
||||||
|
.map_err(|_| "no 'vector' column".to_string())?;
|
||||||
|
let col = batch.column(idx);
|
||||||
|
if let Some(binary) = col.as_any().downcast_ref::<BinaryArray>() {
|
||||||
|
for i in 0..binary.len() {
|
||||||
|
if !binary.is_null(i) {
|
||||||
|
return Ok(binary.value(i).len() / 4);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else if let Some(fsl) = col.as_any().downcast_ref::<FixedSizeListArray>() {
|
||||||
|
return Ok(fsl.value_length() as usize);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err("could not determine vector dimensions".into())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn convert_to_fixed_size_list(
|
||||||
|
schema: &Arc<Schema>,
|
||||||
|
batches: Vec<RecordBatch>,
|
||||||
|
dims: usize,
|
||||||
|
) -> Result<(Arc<Schema>, Vec<RecordBatch>), String> {
|
||||||
|
let new_fields: Vec<Arc<Field>> = schema
|
||||||
|
.fields()
|
||||||
|
.iter()
|
||||||
|
.map(|f| {
|
||||||
|
if f.name() == "vector" {
|
||||||
|
Arc::new(Field::new(
|
||||||
|
"vector",
|
||||||
|
DataType::FixedSizeList(
|
||||||
|
Arc::new(Field::new("item", DataType::Float32, true)),
|
||||||
|
dims as i32,
|
||||||
|
),
|
||||||
|
false,
|
||||||
|
))
|
||||||
|
} else {
|
||||||
|
f.clone()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let new_schema = Arc::new(Schema::new(new_fields));
|
||||||
|
|
||||||
|
let mut out = Vec::with_capacity(batches.len());
|
||||||
|
for batch in batches {
|
||||||
|
let vec_idx = batch.schema().index_of("vector").map_err(e)?;
|
||||||
|
let mut new_cols: Vec<ArrayRef> = Vec::with_capacity(batch.num_columns());
|
||||||
|
for (i, col) in batch.columns().iter().enumerate() {
|
||||||
|
if i == vec_idx {
|
||||||
|
if let Some(bin) = col.as_any().downcast_ref::<BinaryArray>() {
|
||||||
|
new_cols.push(Arc::new(binary_to_fsl(bin, dims)?));
|
||||||
|
} else if col.as_any().is::<FixedSizeListArray>() {
|
||||||
|
// Already in the right shape — just clone.
|
||||||
|
new_cols.push(col.clone());
|
||||||
|
} else {
|
||||||
|
return Err("vector column is neither Binary nor FixedSizeList".into());
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
new_cols.push(col.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out.push(RecordBatch::try_new(new_schema.clone(), new_cols).map_err(e)?);
|
||||||
|
}
|
||||||
|
Ok((new_schema, out))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn binary_to_fsl(bin: &BinaryArray, dims: usize) -> Result<FixedSizeListArray, String> {
|
||||||
|
let n = bin.len();
|
||||||
|
let mut flat: Vec<f32> = Vec::with_capacity(n * dims);
|
||||||
|
for i in 0..n {
|
||||||
|
if bin.is_null(i) {
|
||||||
|
flat.extend(std::iter::repeat(0.0).take(dims));
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let b = bin.value(i);
|
||||||
|
if b.len() != dims * 4 {
|
||||||
|
return Err(format!(
|
||||||
|
"row {i}: {} bytes vs expected {} ({} × f32)",
|
||||||
|
b.len(), dims * 4, dims,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
for c in b.chunks_exact(4) {
|
||||||
|
flat.push(f32::from_le_bytes([c[0], c[1], c[2], c[3]]));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let values = Float32Array::from(flat);
|
||||||
|
let field = Arc::new(Field::new("item", DataType::Float32, true));
|
||||||
|
FixedSizeListArray::try_new(field, dims as i32, Arc::new(values), None).map_err(e)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn write_dataset(
|
||||||
|
path: &str,
|
||||||
|
schema: Arc<Schema>,
|
||||||
|
batches: Vec<RecordBatch>,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
use lance::dataset::{Dataset, WriteParams};
|
||||||
|
let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema);
|
||||||
|
Dataset::write(reader, path, Some(WriteParams::default()))
|
||||||
|
.await.map_err(e)?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn row_from_batch(batch: &RecordBatch, row: usize) -> Result<Row, String> {
|
||||||
|
let doc_id = batch.column_by_name("doc_id")
|
||||||
|
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
|
||||||
|
.map(|a| a.value(row).to_string())
|
||||||
|
.ok_or_else(|| "missing doc_id".to_string())?;
|
||||||
|
let chunk_text = batch.column_by_name("chunk_text")
|
||||||
|
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
|
||||||
|
.map(|a| a.value(row).to_string())
|
||||||
|
.unwrap_or_default();
|
||||||
|
let source = batch.column_by_name("source")
|
||||||
|
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
|
||||||
|
.and_then(|a| if a.is_null(row) { None } else { Some(a.value(row).to_string()) });
|
||||||
|
let chunk_idx = batch.column_by_name("chunk_idx")
|
||||||
|
.and_then(|c| c.as_any().downcast_ref::<Int32Array>())
|
||||||
|
.and_then(|a| if a.is_null(row) { None } else { Some(a.value(row)) });
|
||||||
|
|
||||||
|
let vec_col = batch.column_by_name("vector")
|
||||||
|
.ok_or_else(|| "no vector column in row fetch".to_string())?;
|
||||||
|
let fsl = vec_col.as_any().downcast_ref::<FixedSizeListArray>()
|
||||||
|
.ok_or_else(|| "vector column is not FixedSizeList".to_string())?;
|
||||||
|
let inner = fsl.value(row);
|
||||||
|
let floats = inner.as_any().downcast_ref::<Float32Array>()
|
||||||
|
.ok_or_else(|| "vector inner not Float32".to_string())?;
|
||||||
|
let mut v = Vec::with_capacity(floats.len());
|
||||||
|
for i in 0..floats.len() { v.push(floats.value(i)); }
|
||||||
|
|
||||||
|
Ok(Row { doc_id, chunk_text, vector: v, source, chunk_idx })
|
||||||
|
}
|
||||||
@@ -0,0 +1,25 @@
|
|||||||
|
[package]
|
||||||
|
name = "vectord"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2024"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
shared = { path = "../shared" }
|
||||||
|
storaged = { path = "../storaged" }
|
||||||
|
aibridge = { path = "../aibridge" }
|
||||||
|
catalogd = { path = "../catalogd" }
|
||||||
|
# ADR-019 firewall — vectord-lance owns its own Arrow 57 / Lance 4 deps.
|
||||||
|
# Public API uses only std types so no version conflict propagates here.
|
||||||
|
vectord-lance = { path = "../vectord-lance" }
|
||||||
|
tokio = { workspace = true }
|
||||||
|
axum = { workspace = true }
|
||||||
|
serde = { workspace = true }
|
||||||
|
serde_json = { workspace = true }
|
||||||
|
tracing = { workspace = true }
|
||||||
|
bytes = { workspace = true }
|
||||||
|
object_store = { workspace = true }
|
||||||
|
parquet = { workspace = true }
|
||||||
|
arrow = { workspace = true }
|
||||||
|
chrono = { workspace = true }
|
||||||
|
instant-distance = { workspace = true }
|
||||||
|
uuid = { workspace = true }
|
||||||
@@ -0,0 +1,711 @@
|
|||||||
|
//! Phase 16.2 + 16.5 — The autotune agent.
|
||||||
|
//!
|
||||||
|
//! A long-running tokio task that watches the trial journal and
|
||||||
|
//! autonomously proposes + runs new HNSW configs. Distinct from
|
||||||
|
//! `autotune::run_autotune` which is synchronous (one HTTP call, grid
|
||||||
|
//! of trials, done). The agent is the continuous version: it sleeps,
|
||||||
|
//! wakes on triggers, proposes configs based on prior trial history,
|
||||||
|
//! runs them one at a time, and auto-promotes when it finds an
|
||||||
|
//! improvement.
|
||||||
|
//!
|
||||||
|
//! Design invariants:
|
||||||
|
//! - Trials are data (ADR-018). Every proposal reads the journal; every
|
||||||
|
//! attempt appends to it. The journal is the agent's memory.
|
||||||
|
//! - One trial at a time. Bounded Ollama load — the agent never fires
|
||||||
|
//! multiple parallel embeddings and respects `cooldown_between_trials_secs`.
|
||||||
|
//! - Rate-limited. `max_trials_per_hour` is a hard ceiling so a
|
||||||
|
//! misbehaving proposal function can't saturate the system.
|
||||||
|
//! - Never promotes below `min_recall`. Same safety gate as
|
||||||
|
//! `run_autotune` — we will not make the index worse.
|
||||||
|
//! - Triggered OR periodic. Ingest enqueues a `DatasetAppended` event
|
||||||
|
//! when a new batch lands; the agent also wakes periodically to keep
|
||||||
|
//! exploring even when nothing changed externally.
|
||||||
|
//! - Graceful shutdown via the `stop_tx` signal — the handle's Drop
|
||||||
|
//! doesn't force-kill, but `stop()` requests a clean exit after the
|
||||||
|
//! current trial.
|
||||||
|
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::VecDeque;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::{Mutex, RwLock, mpsc, oneshot};
|
||||||
|
|
||||||
|
use aibridge::client::AiClient;
|
||||||
|
use catalogd::registry::Registry as CatalogRegistry;
|
||||||
|
|
||||||
|
use crate::embedding_cache::EmbeddingCache;
|
||||||
|
use crate::harness;
|
||||||
|
use crate::hnsw::HnswStore;
|
||||||
|
use crate::index_registry::IndexRegistry;
|
||||||
|
use crate::promotion::{PromotionEntry, PromotionRegistry};
|
||||||
|
use crate::trial::{HnswConfig, Trial, TrialJournal, TrialMetrics};
|
||||||
|
|
||||||
|
// -------- Public-facing types --------
|
||||||
|
|
||||||
|
/// Runtime configuration for the agent. Mirrored in shared::config under
|
||||||
|
/// `[agent]`. Defaults are conservative — designed to tune slowly in the
|
||||||
|
/// background without fighting real workloads for GPU time.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct AgentConfig {
|
||||||
|
/// Master switch. When false, `spawn` returns a handle but the loop
|
||||||
|
/// doesn't run.
|
||||||
|
pub enabled: bool,
|
||||||
|
/// Periodic wake-up — even if the trigger queue is empty, every N
|
||||||
|
/// seconds the agent picks an index with trials and proposes one
|
||||||
|
/// more config. Keeps exploration alive on idle indexes.
|
||||||
|
pub cycle_interval_secs: u64,
|
||||||
|
/// Minimum gap between two trials on the SAME index. Prevents the
|
||||||
|
/// agent from hammering Ollama when a hot index has many pending
|
||||||
|
/// triggers in a row.
|
||||||
|
pub cooldown_between_trials_secs: u64,
|
||||||
|
/// Below this recall, a proposal is never promoted — even if it
|
||||||
|
/// beats the champion on latency.
|
||||||
|
pub min_recall: f32,
|
||||||
|
/// Budget cap: hard ceiling on trials per hour across all indexes.
|
||||||
|
/// When hit, the agent idles until the hour window rolls.
|
||||||
|
pub max_trials_per_hour: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for AgentConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
enabled: false, // opt-in — don't auto-start until J turns it on
|
||||||
|
cycle_interval_secs: 60,
|
||||||
|
cooldown_between_trials_secs: 30,
|
||||||
|
min_recall: 0.9,
|
||||||
|
max_trials_per_hour: 30,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// What caused the agent to look at a particular index. Recorded in the
|
||||||
|
/// trial's note field so we can tell "new data arrived" trials from
|
||||||
|
/// "periodic exploration" trials in the journal.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub enum TriggerReason {
|
||||||
|
/// Ingest just appended to a dataset that has attached HNSW indexes.
|
||||||
|
DatasetAppended { dataset: String },
|
||||||
|
/// A human or another agent hit `/agent/enqueue/{index}`.
|
||||||
|
Manual,
|
||||||
|
/// Periodic wake — no external event, just keep exploring.
|
||||||
|
Periodic,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One unit of work for the agent.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct TriggerEvent {
|
||||||
|
pub index_name: String,
|
||||||
|
pub reason: TriggerReason,
|
||||||
|
pub enqueued_at: DateTime<Utc>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TriggerEvent {
|
||||||
|
pub fn manual(index_name: impl Into<String>) -> Self {
|
||||||
|
Self { index_name: index_name.into(), reason: TriggerReason::Manual, enqueued_at: Utc::now() }
|
||||||
|
}
|
||||||
|
pub fn dataset_appended(index_name: impl Into<String>, dataset: impl Into<String>) -> Self {
|
||||||
|
Self {
|
||||||
|
index_name: index_name.into(),
|
||||||
|
reason: TriggerReason::DatasetAppended { dataset: dataset.into() },
|
||||||
|
enqueued_at: Utc::now(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
pub fn periodic(index_name: impl Into<String>) -> Self {
|
||||||
|
Self { index_name: index_name.into(), reason: TriggerReason::Periodic, enqueued_at: Utc::now() }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Observable snapshot of the agent's state — what `/agent/status` returns.
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct AgentStatus {
|
||||||
|
pub running: bool,
|
||||||
|
pub config: AgentConfig,
|
||||||
|
pub queue_depth: usize,
|
||||||
|
pub trials_run: u64,
|
||||||
|
pub promotions: u64,
|
||||||
|
pub trials_in_last_hour: u32,
|
||||||
|
pub last_event: Option<AgentEvent>,
|
||||||
|
pub started_at: Option<DateTime<Utc>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Last thing that happened — useful for "why didn't it do anything?" debugging.
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct AgentEvent {
|
||||||
|
pub at: DateTime<Utc>,
|
||||||
|
pub kind: String, // "trial_completed" | "promoted" | "skipped_rate_limit" | etc
|
||||||
|
pub index_name: Option<String>,
|
||||||
|
pub detail: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle returned by `spawn`. Holds the trigger sender + shared status +
|
||||||
|
/// stop signal.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct AgentHandle {
|
||||||
|
trigger_tx: mpsc::Sender<TriggerEvent>,
|
||||||
|
inner: Arc<AgentInner>,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct AgentInner {
|
||||||
|
status: RwLock<AgentStatus>,
|
||||||
|
stop_tx: Mutex<Option<oneshot::Sender<()>>>,
|
||||||
|
queue_len: Mutex<usize>, // mirror of the channel capacity — for status reporting
|
||||||
|
recent_trials: Mutex<VecDeque<DateTime<Utc>>>, // ring of recent trial timestamps for rate limit
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AgentHandle {
|
||||||
|
/// Enqueue a trigger. Returns Err if the agent isn't running or the
|
||||||
|
/// queue is full (backpressure — dropping events is correct here
|
||||||
|
/// since periodic exploration will pick up the slack).
|
||||||
|
pub async fn enqueue(&self, event: TriggerEvent) -> Result<(), String> {
|
||||||
|
self.trigger_tx.try_send(event).map_err(|e| format!("enqueue: {e}"))?;
|
||||||
|
let mut guard = self.inner.queue_len.lock().await;
|
||||||
|
*guard = guard.saturating_add(1);
|
||||||
|
// Update queue_depth in status for observability.
|
||||||
|
let mut s = self.inner.status.write().await;
|
||||||
|
s.queue_depth = *guard;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn status(&self) -> AgentStatus {
|
||||||
|
let mut s = self.inner.status.read().await.clone();
|
||||||
|
// Refresh rate-limit window from ring buffer.
|
||||||
|
let cutoff = Utc::now() - chrono::Duration::hours(1);
|
||||||
|
let ring = self.inner.recent_trials.lock().await;
|
||||||
|
s.trials_in_last_hour = ring.iter().filter(|t| **t >= cutoff).count() as u32;
|
||||||
|
s
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Request a graceful stop. Returns immediately — the loop exits
|
||||||
|
/// after its current trial completes.
|
||||||
|
pub async fn stop(&self) -> bool {
|
||||||
|
let mut guard = self.inner.stop_tx.lock().await;
|
||||||
|
if let Some(tx) = guard.take() {
|
||||||
|
let _ = tx.send(());
|
||||||
|
true
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// -------- Agent state holder --------
|
||||||
|
|
||||||
|
/// Everything the agent needs to run a trial. Mirrors the fields of
|
||||||
|
/// `VectorState` the agent actually uses. Kept separate so the service
|
||||||
|
/// layer builds it explicitly — no clone of unneeded state.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct AgentDeps {
|
||||||
|
pub store: Arc<dyn ObjectStore>,
|
||||||
|
pub ai_client: AiClient,
|
||||||
|
pub catalog: CatalogRegistry,
|
||||||
|
pub index_registry: IndexRegistry,
|
||||||
|
pub hnsw_store: HnswStore,
|
||||||
|
pub embedding_cache: EmbeddingCache,
|
||||||
|
pub trial_journal: TrialJournal,
|
||||||
|
pub promotion_registry: PromotionRegistry,
|
||||||
|
}
|
||||||
|
|
||||||
|
// -------- Spawn --------
|
||||||
|
|
||||||
|
/// Start the agent loop in a background tokio task. Returns a handle
|
||||||
|
/// the caller uses to enqueue events and read status.
|
||||||
|
pub fn spawn(config: AgentConfig, deps: AgentDeps) -> AgentHandle {
|
||||||
|
let (trigger_tx, trigger_rx) = mpsc::channel::<TriggerEvent>(256);
|
||||||
|
let (stop_tx, stop_rx) = oneshot::channel::<()>();
|
||||||
|
|
||||||
|
let status = AgentStatus {
|
||||||
|
running: config.enabled,
|
||||||
|
config: config.clone(),
|
||||||
|
queue_depth: 0,
|
||||||
|
trials_run: 0,
|
||||||
|
promotions: 0,
|
||||||
|
trials_in_last_hour: 0,
|
||||||
|
last_event: None,
|
||||||
|
started_at: if config.enabled { Some(Utc::now()) } else { None },
|
||||||
|
};
|
||||||
|
|
||||||
|
let inner = Arc::new(AgentInner {
|
||||||
|
status: RwLock::new(status),
|
||||||
|
stop_tx: Mutex::new(Some(stop_tx)),
|
||||||
|
queue_len: Mutex::new(0),
|
||||||
|
recent_trials: Mutex::new(VecDeque::with_capacity(64)),
|
||||||
|
});
|
||||||
|
|
||||||
|
if config.enabled {
|
||||||
|
tracing::info!(
|
||||||
|
"autotune agent started (cycle={}s, cooldown={}s, cap={}/hr, min_recall={})",
|
||||||
|
config.cycle_interval_secs, config.cooldown_between_trials_secs,
|
||||||
|
config.max_trials_per_hour, config.min_recall,
|
||||||
|
);
|
||||||
|
let loop_inner = inner.clone();
|
||||||
|
let loop_deps = deps.clone();
|
||||||
|
let loop_config = config.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
run_loop(loop_config, loop_deps, trigger_rx, stop_rx, loop_inner).await;
|
||||||
|
});
|
||||||
|
} else {
|
||||||
|
// Agent disabled — still drain the channel so sends don't back up.
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let mut rx = trigger_rx;
|
||||||
|
while rx.recv().await.is_some() {}
|
||||||
|
});
|
||||||
|
tracing::info!("autotune agent configured but disabled (set [agent].enabled=true)");
|
||||||
|
}
|
||||||
|
|
||||||
|
AgentHandle { trigger_tx, inner }
|
||||||
|
}
|
||||||
|
|
||||||
|
// -------- Main loop --------
|
||||||
|
|
||||||
|
async fn run_loop(
|
||||||
|
config: AgentConfig,
|
||||||
|
deps: AgentDeps,
|
||||||
|
mut trigger_rx: mpsc::Receiver<TriggerEvent>,
|
||||||
|
mut stop_rx: oneshot::Receiver<()>,
|
||||||
|
inner: Arc<AgentInner>,
|
||||||
|
) {
|
||||||
|
let mut periodic = tokio::time::interval(std::time::Duration::from_secs(config.cycle_interval_secs));
|
||||||
|
// First tick fires immediately — skip it so we don't double-fire on startup.
|
||||||
|
periodic.tick().await;
|
||||||
|
|
||||||
|
loop {
|
||||||
|
let event = tokio::select! {
|
||||||
|
_ = &mut stop_rx => {
|
||||||
|
tracing::info!("autotune agent: stop signal received");
|
||||||
|
let mut s = inner.status.write().await;
|
||||||
|
s.running = false;
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
maybe = trigger_rx.recv() => match maybe {
|
||||||
|
Some(ev) => {
|
||||||
|
let mut guard = inner.queue_len.lock().await;
|
||||||
|
*guard = guard.saturating_sub(1);
|
||||||
|
let mut s = inner.status.write().await;
|
||||||
|
s.queue_depth = *guard;
|
||||||
|
ev
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
tracing::info!("autotune agent: trigger channel closed");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
},
|
||||||
|
_ = periodic.tick() => {
|
||||||
|
// Periodic wake — pick an index with existing trials.
|
||||||
|
// If nothing's been tuned yet there's nothing to propose.
|
||||||
|
match pick_periodic_target(&deps).await {
|
||||||
|
Some(idx) => TriggerEvent::periodic(idx),
|
||||||
|
None => continue,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Rate limit: count recent trials, skip if over budget.
|
||||||
|
if over_rate_limit(&inner, config.max_trials_per_hour).await {
|
||||||
|
record_event(&inner, "skipped_rate_limit", Some(&event.index_name),
|
||||||
|
format!("hit cap of {}/hour", config.max_trials_per_hour)).await;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Per-index cooldown.
|
||||||
|
if cooling_down(&inner, &event.index_name, config.cooldown_between_trials_secs).await {
|
||||||
|
record_event(&inner, "skipped_cooldown", Some(&event.index_name),
|
||||||
|
format!("last trial too recent (<{}s)", config.cooldown_between_trials_secs)).await;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Run one trial.
|
||||||
|
match run_one_cycle(&event, &deps, config.min_recall).await {
|
||||||
|
Ok(outcome) => {
|
||||||
|
mark_recent_trial(&inner).await;
|
||||||
|
{
|
||||||
|
let mut s = inner.status.write().await;
|
||||||
|
s.trials_run += 1;
|
||||||
|
if outcome.promoted { s.promotions += 1; }
|
||||||
|
}
|
||||||
|
record_event(&inner, if outcome.promoted { "promoted" } else { "trial_completed" },
|
||||||
|
Some(&event.index_name),
|
||||||
|
format!("config=ec{}/es{} recall={:.3} p50={:.0}us {}",
|
||||||
|
outcome.trial.config.ef_construction,
|
||||||
|
outcome.trial.config.ef_search,
|
||||||
|
outcome.trial.metrics.recall_at_k,
|
||||||
|
outcome.trial.metrics.search_latency_p50_us,
|
||||||
|
if outcome.promoted { "★ PROMOTED" } else { "" })).await;
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
record_event(&inner, "trial_error", Some(&event.index_name), e).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Result of one cycle — ran a trial, maybe promoted it.
|
||||||
|
struct CycleOutcome {
|
||||||
|
trial: Trial,
|
||||||
|
promoted: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Core cycle: propose → build → bench → record → maybe promote.
|
||||||
|
async fn run_one_cycle(
|
||||||
|
event: &TriggerEvent,
|
||||||
|
deps: &AgentDeps,
|
||||||
|
min_recall: f32,
|
||||||
|
) -> Result<CycleOutcome, String> {
|
||||||
|
// Read history.
|
||||||
|
let history = deps.trial_journal.list(&event.index_name).await
|
||||||
|
.map_err(|e| format!("read journal: {e}"))?;
|
||||||
|
if history.is_empty() {
|
||||||
|
return Err(format!(
|
||||||
|
"no trials yet for '{}' — seed with at least one POST /hnsw/trial first",
|
||||||
|
event.index_name,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Current champion (if any) is the promoted config.
|
||||||
|
let champion = deps.promotion_registry.get_current(&event.index_name).await;
|
||||||
|
let champion_trial = champion.as_ref().and_then(|p| {
|
||||||
|
history.iter().find(|t| t.id == p.trial_id).cloned()
|
||||||
|
});
|
||||||
|
|
||||||
|
// Propose the next config.
|
||||||
|
let Some(next_config) = propose_next_config(&history, champion_trial.as_ref()) else {
|
||||||
|
return Err("proposer returned None — search space exhausted".into());
|
||||||
|
};
|
||||||
|
|
||||||
|
// Validate bounds defensively.
|
||||||
|
if !(10..=400).contains(&next_config.ef_construction) {
|
||||||
|
return Err(format!("proposed ef_construction={} out of bounds", next_config.ef_construction));
|
||||||
|
}
|
||||||
|
if !(10..=200).contains(&next_config.ef_search) {
|
||||||
|
return Err(format!("proposed ef_search={} out of bounds", next_config.ef_search));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Need a harness to measure. Use the most recent one in history.
|
||||||
|
// (A future refinement: remember per-index "canonical harness" on
|
||||||
|
// the index metadata. For now: latest wins.)
|
||||||
|
let harness_name = history.last().unwrap().eval_set.clone();
|
||||||
|
let mut harness_set = harness::EvalSet::load(&deps.store, &harness_name).await
|
||||||
|
.map_err(|e| format!("load harness '{harness_name}': {e}"))?;
|
||||||
|
|
||||||
|
let embeddings = deps.embedding_cache.get_or_load(&event.index_name).await
|
||||||
|
.map_err(|e| format!("embeddings: {e}"))?;
|
||||||
|
|
||||||
|
if !harness_set.ground_truth_built {
|
||||||
|
harness::compute_ground_truth(&mut harness_set, &embeddings, &deps.ai_client).await
|
||||||
|
.map_err(|e| format!("ground truth: {e}"))?;
|
||||||
|
harness_set.save(&deps.store).await.ok();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build + bench.
|
||||||
|
let trial_id = Trial::new_id();
|
||||||
|
let slot = format!("{}__{}", event.index_name, trial_id);
|
||||||
|
let build = deps.hnsw_store
|
||||||
|
.build_index_with_config(&slot, (*embeddings).clone(), &next_config)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let query_vectors: Vec<Vec<f32>> = harness_set.queries
|
||||||
|
.iter().filter_map(|q| q.query_embedding.clone()).collect();
|
||||||
|
let bench = deps.hnsw_store.bench_search(&slot, &query_vectors, harness_set.k).await?;
|
||||||
|
|
||||||
|
let mut recalls = Vec::with_capacity(harness_set.queries.len());
|
||||||
|
for (q, hits) in harness_set.queries.iter().zip(bench.retrieved.iter()) {
|
||||||
|
if let Some(gt) = &q.ground_truth {
|
||||||
|
recalls.push(harness::recall_at_k(hits, gt, harness_set.k));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mean_recall = if recalls.is_empty() { 0.0 } else {
|
||||||
|
recalls.iter().sum::<f32>() / recalls.len() as f32
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut lats = bench.latencies_us.clone();
|
||||||
|
lats.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
let p = |pct: f32| -> f32 {
|
||||||
|
if lats.is_empty() { return 0.0; }
|
||||||
|
let idx = ((lats.len() as f32 - 1.0) * pct).round() as usize;
|
||||||
|
lats[idx.min(lats.len() - 1)]
|
||||||
|
};
|
||||||
|
|
||||||
|
let dims = embeddings.first().map(|e| e.vector.len()).unwrap_or(0);
|
||||||
|
let memory_bytes = (embeddings.len() * dims * std::mem::size_of::<f32>() + embeddings.len() * 128) as u64;
|
||||||
|
|
||||||
|
let note = match &event.reason {
|
||||||
|
TriggerReason::DatasetAppended { dataset } => format!("agent: dataset_appended({dataset})"),
|
||||||
|
TriggerReason::Manual => "agent: manual".to_string(),
|
||||||
|
TriggerReason::Periodic => "agent: periodic".to_string(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let trial = Trial {
|
||||||
|
id: trial_id,
|
||||||
|
index_name: event.index_name.clone(),
|
||||||
|
eval_set: harness_set.name.clone(),
|
||||||
|
config: next_config.clone(),
|
||||||
|
metrics: TrialMetrics {
|
||||||
|
build_time_secs: build.build_time_secs,
|
||||||
|
search_latency_p50_us: p(0.50),
|
||||||
|
search_latency_p95_us: p(0.95),
|
||||||
|
search_latency_p99_us: p(0.99),
|
||||||
|
recall_at_k: mean_recall,
|
||||||
|
memory_bytes,
|
||||||
|
vectors: build.vectors,
|
||||||
|
eval_queries: harness_set.queries.len(),
|
||||||
|
brute_force_latency_us: 0.0,
|
||||||
|
},
|
||||||
|
created_at: Utc::now(),
|
||||||
|
note: Some(note),
|
||||||
|
};
|
||||||
|
deps.trial_journal.append(&trial).await.ok();
|
||||||
|
deps.hnsw_store.drop(&slot).await;
|
||||||
|
|
||||||
|
// Promotion decision: the new trial must meet recall gate AND beat
|
||||||
|
// the current champion (higher recall OR same recall + lower p50).
|
||||||
|
let promoted = if trial.metrics.recall_at_k < min_recall {
|
||||||
|
false
|
||||||
|
} else {
|
||||||
|
let beats = match &champion_trial {
|
||||||
|
None => true, // no champion yet — anything passing the gate wins
|
||||||
|
Some(c) => beats_champion(&trial, c),
|
||||||
|
};
|
||||||
|
if beats {
|
||||||
|
let entry = PromotionEntry {
|
||||||
|
config: trial.config.clone(),
|
||||||
|
trial_id: trial.id.clone(),
|
||||||
|
promoted_at: Utc::now(),
|
||||||
|
promoted_by: "agent".to_string(),
|
||||||
|
note: Some(format!(
|
||||||
|
"auto-promote: recall={:.3} p50={:.0}us (was {:.3}/{:.0}us)",
|
||||||
|
trial.metrics.recall_at_k, trial.metrics.search_latency_p50_us,
|
||||||
|
champion_trial.as_ref().map(|t| t.metrics.recall_at_k).unwrap_or(0.0),
|
||||||
|
champion_trial.as_ref().map(|t| t.metrics.search_latency_p50_us).unwrap_or(0.0),
|
||||||
|
)),
|
||||||
|
};
|
||||||
|
deps.promotion_registry.promote(&event.index_name, entry).await.is_ok()
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(CycleOutcome { trial, promoted })
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Champion-beat test: strictly higher recall, OR equal recall with
|
||||||
|
/// lower p50. Same rule as autotune::pick_winner — kept consistent so the
|
||||||
|
/// agent and the synchronous autotune agree on what "better" means.
|
||||||
|
fn beats_champion(candidate: &Trial, champion: &Trial) -> bool {
|
||||||
|
if candidate.metrics.recall_at_k > champion.metrics.recall_at_k {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
if (candidate.metrics.recall_at_k - champion.metrics.recall_at_k).abs() < 1e-4
|
||||||
|
&& candidate.metrics.search_latency_p50_us < champion.metrics.search_latency_p50_us {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
false
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Propose the next HnswConfig given trial history and the current
|
||||||
|
/// champion.
|
||||||
|
///
|
||||||
|
/// ============================================================
|
||||||
|
/// J: THIS IS YOURS TO IMPLEMENT
|
||||||
|
/// ============================================================
|
||||||
|
///
|
||||||
|
/// Inputs:
|
||||||
|
/// - `history`: every trial ever run on this index, oldest first
|
||||||
|
/// - `champion`: the currently-promoted trial, if any
|
||||||
|
///
|
||||||
|
/// Output:
|
||||||
|
/// - `Some(HnswConfig)` with the config to try next
|
||||||
|
/// - `None` if you decide there's nothing worth trying (agent logs
|
||||||
|
/// "search space exhausted" and moves on)
|
||||||
|
///
|
||||||
|
/// Hard bounds the caller enforces: ef_construction ∈ [10, 400],
|
||||||
|
/// ef_search ∈ [10, 200]. Stay inside those — configs outside get
|
||||||
|
/// rejected and count as a wasted cycle.
|
||||||
|
///
|
||||||
|
/// Design options (pick one, or mix):
|
||||||
|
///
|
||||||
|
/// 1. LOCAL REFINEMENT (exploit-heavy):
|
||||||
|
/// Sample near champion ± small delta. Converges fast, risks local
|
||||||
|
/// minima. Good for "we know roughly where the sweet spot is."
|
||||||
|
///
|
||||||
|
/// 2. ε-GREEDY (mixed):
|
||||||
|
/// With prob ε, random sample from full bounds (explore). Otherwise
|
||||||
|
/// refinement around champion (exploit). ε=0.2 is a reasonable start.
|
||||||
|
/// Good for long-running tune with no prior knowledge.
|
||||||
|
///
|
||||||
|
/// 3. COARSE→FINE (annealed):
|
||||||
|
/// First N trials: wide random. Then shrink the neighborhood around
|
||||||
|
/// champion as more trials accumulate. Mimics simulated annealing.
|
||||||
|
///
|
||||||
|
/// 4. DEDUP-AWARE:
|
||||||
|
/// Whatever strategy, skip configs already in history. Prevents the
|
||||||
|
/// agent from re-running the same (ec, es) pair twice.
|
||||||
|
///
|
||||||
|
/// A simple starter implementation is provided below (local refinement
|
||||||
|
/// + dedup). Replace with your preferred strategy.
|
||||||
|
pub fn propose_next_config(history: &[Trial], champion: Option<&Trial>) -> Option<HnswConfig> {
|
||||||
|
// ε-greedy around the champion, dedup-aware.
|
||||||
|
//
|
||||||
|
// - With probability ε (≈0.25), sample a random config from the full
|
||||||
|
// bounds. Keeps exploration alive so we don't get stuck hill-climbing
|
||||||
|
// one axis.
|
||||||
|
// - Otherwise: perturb the champion symmetrically on BOTH axes (not
|
||||||
|
// just +20 / +40 like the starter did). Prefers small moves first
|
||||||
|
// so recall stays near the current level.
|
||||||
|
// - Always skip configs already in history — no point re-running.
|
||||||
|
// - Deterministic per-history: RNG is seeded from history length so
|
||||||
|
// the same journal state always proposes the same next config.
|
||||||
|
// Makes tests + offline replay reproducible.
|
||||||
|
|
||||||
|
let base = champion
|
||||||
|
.map(|t| t.config.clone())
|
||||||
|
.or_else(|| history.last().map(|t| t.config.clone()))
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
let tried = |ec: usize, es: usize| -> bool {
|
||||||
|
history.iter().any(|t|
|
||||||
|
t.config.ef_construction == ec && t.config.ef_search == es
|
||||||
|
)
|
||||||
|
};
|
||||||
|
let clamp = |ec: i32, es: i32| -> (usize, usize) {
|
||||||
|
(ec.clamp(10, 400) as usize, es.clamp(10, 200) as usize)
|
||||||
|
};
|
||||||
|
|
||||||
|
// Tiny xorshift — no rand crate dep. Seeded from history length so the
|
||||||
|
// proposer is deterministic for a given journal state.
|
||||||
|
let mut rng = (history.len() as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
|
||||||
|
let mut next_u = || {
|
||||||
|
rng ^= rng << 13; rng ^= rng >> 7; rng ^= rng << 17;
|
||||||
|
rng
|
||||||
|
};
|
||||||
|
|
||||||
|
for _attempt in 0..32 {
|
||||||
|
let explore = (next_u() % 100) < 25; // ε = 0.25
|
||||||
|
let (ec, es) = if explore {
|
||||||
|
clamp(10 + (next_u() % 391) as i32, 10 + (next_u() % 191) as i32)
|
||||||
|
} else {
|
||||||
|
// Symmetric perturbation — signed steps on both axes.
|
||||||
|
let dec = [-40, -20, -10, 10, 20, 40][(next_u() % 6) as usize];
|
||||||
|
let des = [-20, -10, -5, 5, 10, 20][(next_u() % 6) as usize];
|
||||||
|
clamp(base.ef_construction as i32 + dec, base.ef_search as i32 + des)
|
||||||
|
};
|
||||||
|
if !tried(ec, es) {
|
||||||
|
return Some(HnswConfig { ef_construction: ec, ef_search: es, seed: Some(42) });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None // 32 attempts all landed on duplicates — likely saturated
|
||||||
|
}
|
||||||
|
|
||||||
|
// -------- Helpers --------
|
||||||
|
|
||||||
|
/// Find an index to poke on a periodic wake. Strategy: the one with the
|
||||||
|
/// most recent trial. If nothing's been trialed yet, return None.
|
||||||
|
async fn pick_periodic_target(deps: &AgentDeps) -> Option<String> {
|
||||||
|
// `/agent` runs against any index that has a trial journal. We don't
|
||||||
|
// have a "list all journals" helper, so we derive candidates from the
|
||||||
|
// promotion registry (indexes with a human ever promoting are live).
|
||||||
|
let promos = deps.promotion_registry.list_all().await.ok()?;
|
||||||
|
// Prefer the one most recently promoted — it's the one a human cares
|
||||||
|
// about right now.
|
||||||
|
promos.into_iter()
|
||||||
|
.filter_map(|f| f.current.map(|c| (f.index_name, c.promoted_at)))
|
||||||
|
.max_by_key(|(_, at)| *at)
|
||||||
|
.map(|(name, _)| name)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn over_rate_limit(inner: &Arc<AgentInner>, cap: u32) -> bool {
|
||||||
|
let cutoff = Utc::now() - chrono::Duration::hours(1);
|
||||||
|
let ring = inner.recent_trials.lock().await;
|
||||||
|
ring.iter().filter(|t| **t >= cutoff).count() as u32 >= cap
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn cooling_down(inner: &Arc<AgentInner>, _index: &str, cooldown_secs: u64) -> bool {
|
||||||
|
// Minimal impl: gate on global most-recent trial rather than per-index.
|
||||||
|
// Per-index cooldown would be easy to add — keep a HashMap<String, DateTime>
|
||||||
|
// — but for Phase 16.2 MVP, global is fine. Ollama is the shared resource.
|
||||||
|
let ring = inner.recent_trials.lock().await;
|
||||||
|
if let Some(last) = ring.back() {
|
||||||
|
let since = Utc::now().signed_duration_since(*last);
|
||||||
|
return since < chrono::Duration::seconds(cooldown_secs as i64);
|
||||||
|
}
|
||||||
|
false
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mark_recent_trial(inner: &Arc<AgentInner>) {
|
||||||
|
let mut ring = inner.recent_trials.lock().await;
|
||||||
|
ring.push_back(Utc::now());
|
||||||
|
// Keep bounded.
|
||||||
|
while ring.len() > 256 {
|
||||||
|
ring.pop_front();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_event(
|
||||||
|
inner: &Arc<AgentInner>,
|
||||||
|
kind: &str,
|
||||||
|
index: Option<&str>,
|
||||||
|
detail: String,
|
||||||
|
) {
|
||||||
|
tracing::info!("agent: {} {}{}", kind,
|
||||||
|
index.map(|i| format!("[{i}] ")).unwrap_or_default(), detail);
|
||||||
|
let mut s = inner.status.write().await;
|
||||||
|
s.last_event = Some(AgentEvent {
|
||||||
|
at: Utc::now(),
|
||||||
|
kind: kind.to_string(),
|
||||||
|
index_name: index.map(String::from),
|
||||||
|
detail,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn mk_trial(ec: usize, es: usize, recall: f32, p50: f32) -> Trial {
|
||||||
|
Trial {
|
||||||
|
id: format!("t-{ec}-{es}"),
|
||||||
|
index_name: "test".into(),
|
||||||
|
eval_set: "eval".into(),
|
||||||
|
config: HnswConfig { ef_construction: ec, ef_search: es, seed: Some(42) },
|
||||||
|
metrics: TrialMetrics {
|
||||||
|
build_time_secs: 1.0,
|
||||||
|
search_latency_p50_us: p50,
|
||||||
|
search_latency_p95_us: p50 * 1.5,
|
||||||
|
search_latency_p99_us: p50 * 2.0,
|
||||||
|
recall_at_k: recall,
|
||||||
|
memory_bytes: 0, vectors: 1000, eval_queries: 10,
|
||||||
|
brute_force_latency_us: 100.0,
|
||||||
|
},
|
||||||
|
created_at: Utc::now(),
|
||||||
|
note: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn propose_skips_duplicates() {
|
||||||
|
let hist = vec![
|
||||||
|
mk_trial(80, 30, 1.0, 500.0),
|
||||||
|
mk_trial(100, 30, 1.0, 520.0), // ec+20
|
||||||
|
];
|
||||||
|
let next = propose_next_config(&hist, Some(&hist[0])).unwrap();
|
||||||
|
// ec+20 is taken, so the proposer should skip it.
|
||||||
|
assert!(next.ef_construction != 100 || next.ef_search != 30);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn beats_champion_strict_recall() {
|
||||||
|
let champ = mk_trial(80, 30, 0.95, 500.0);
|
||||||
|
let better_recall = mk_trial(80, 30, 0.99, 600.0);
|
||||||
|
let worse_recall = mk_trial(80, 30, 0.90, 100.0);
|
||||||
|
assert!(beats_champion(&better_recall, &champ));
|
||||||
|
assert!(!beats_champion(&worse_recall, &champ));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn beats_champion_same_recall_lower_latency() {
|
||||||
|
let champ = mk_trial(80, 30, 1.0, 500.0);
|
||||||
|
let faster = mk_trial(60, 30, 1.0, 400.0);
|
||||||
|
let slower = mk_trial(60, 30, 1.0, 600.0);
|
||||||
|
assert!(beats_champion(&faster, &champ));
|
||||||
|
assert!(!beats_champion(&slower, &champ));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,328 @@
|
|||||||
|
//! Phase 16: Autotune — the agent loop for HNSW tuning.
|
||||||
|
//!
|
||||||
|
//! Given an index and a harness, run a grid of (ef_construction,
|
||||||
|
//! ef_search) trials, pick the Pareto winner, and promote it.
|
||||||
|
//!
|
||||||
|
//! This is deliberately NOT a sophisticated optimizer — random-grid
|
||||||
|
//! search is good enough when the config space is small (a dozen
|
||||||
|
//! combinations) and the work per trial is seconds, not hours. Think
|
||||||
|
//! of it as a "smart operator" rather than Bayesian optimization.
|
||||||
|
//!
|
||||||
|
//! Score: we want high recall and low p50 latency. The function below
|
||||||
|
//! picks the Pareto-dominating trial, breaking ties by prefer-lower-
|
||||||
|
//! build-time (cheaper to rebuild on data changes).
|
||||||
|
//!
|
||||||
|
//! Safeguards (PRD risk mitigation):
|
||||||
|
//! - Never promote a trial with recall < `min_recall` (default 0.9)
|
||||||
|
//! - Hard config bounds: ef_construction ∈ [10, 400], ef_search ∈ [10, 200]
|
||||||
|
//! - Concurrency cap: one autotune per index at a time (tracked by
|
||||||
|
//! JobTracker — caller is responsible for not firing in parallel
|
||||||
|
//! for the same index)
|
||||||
|
|
||||||
|
use chrono::Utc;
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use aibridge::client::AiClient;
|
||||||
|
use catalogd::registry::Registry as CatalogRegistry;
|
||||||
|
|
||||||
|
use crate::embedding_cache::EmbeddingCache;
|
||||||
|
use crate::harness;
|
||||||
|
use crate::hnsw::HnswStore;
|
||||||
|
use crate::index_registry::IndexRegistry;
|
||||||
|
use crate::jobs::JobTracker;
|
||||||
|
use crate::promotion::{PromotionEntry, PromotionRegistry};
|
||||||
|
use crate::trial::{HnswConfig, Trial, TrialJournal, TrialMetrics};
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Deserialize)]
|
||||||
|
pub struct AutotuneRequest {
|
||||||
|
pub index_name: String,
|
||||||
|
pub harness: String,
|
||||||
|
/// Minimum recall a trial must achieve to be eligible for promotion.
|
||||||
|
/// Default 0.9 — below that, the tuning is not producing useful
|
||||||
|
/// answers. Never promote a bad index.
|
||||||
|
#[serde(default = "default_min_recall")]
|
||||||
|
pub min_recall: f32,
|
||||||
|
/// Override the default config grid (rarely needed). When omitted,
|
||||||
|
/// a small curated grid runs.
|
||||||
|
#[serde(default)]
|
||||||
|
pub grid: Option<Vec<HnswConfig>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_min_recall() -> f32 { 0.9 }
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct AutotuneResult {
|
||||||
|
pub index_name: String,
|
||||||
|
pub trials_run: usize,
|
||||||
|
pub trials_eligible: usize,
|
||||||
|
pub winner: Option<WinnerReport>,
|
||||||
|
pub promoted: bool,
|
||||||
|
pub duration_secs: f32,
|
||||||
|
pub skipped: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct WinnerReport {
|
||||||
|
pub trial_id: String,
|
||||||
|
pub config: HnswConfig,
|
||||||
|
pub recall_at_k: f32,
|
||||||
|
pub search_latency_p50_us: f32,
|
||||||
|
pub build_time_secs: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Default exploration grid. Covers the practical range of settings
|
||||||
|
/// from "fast build, OK recall" to "thorough build, great recall".
|
||||||
|
fn default_grid() -> Vec<HnswConfig> {
|
||||||
|
vec![
|
||||||
|
HnswConfig { ef_construction: 20, ef_search: 30, seed: Some(42) },
|
||||||
|
HnswConfig { ef_construction: 40, ef_search: 30, seed: Some(42) },
|
||||||
|
HnswConfig { ef_construction: 80, ef_search: 30, seed: Some(42) },
|
||||||
|
HnswConfig { ef_construction: 80, ef_search: 60, seed: Some(42) },
|
||||||
|
HnswConfig { ef_construction: 160, ef_search: 30, seed: Some(42) },
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Config-space bounds. Rejects anything outside these before running a
|
||||||
|
/// trial — defends against an agent getting creative with absurd values.
|
||||||
|
fn sanitize_grid(grid: Vec<HnswConfig>) -> (Vec<HnswConfig>, Vec<String>) {
|
||||||
|
let mut ok = Vec::new();
|
||||||
|
let mut rejected = Vec::new();
|
||||||
|
for cfg in grid {
|
||||||
|
if !(10..=400).contains(&cfg.ef_construction) {
|
||||||
|
rejected.push(format!(
|
||||||
|
"ef_construction={} out of bounds [10,400]", cfg.ef_construction,
|
||||||
|
));
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if !(10..=200).contains(&cfg.ef_search) {
|
||||||
|
rejected.push(format!(
|
||||||
|
"ef_search={} out of bounds [10,200]", cfg.ef_search,
|
||||||
|
));
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
ok.push(cfg);
|
||||||
|
}
|
||||||
|
(ok, rejected)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pareto winner: highest recall, then lowest p50 latency among
|
||||||
|
/// eligible trials (recall >= min_recall).
|
||||||
|
fn pick_winner(trials: &[Trial], min_recall: f32) -> Option<&Trial> {
|
||||||
|
trials
|
||||||
|
.iter()
|
||||||
|
.filter(|t| t.metrics.recall_at_k >= min_recall)
|
||||||
|
.min_by(|a, b| {
|
||||||
|
// Compare by (−recall, p50_latency) so minimum is the
|
||||||
|
// best trial: higher recall wins, then lower latency.
|
||||||
|
let a_key = (-a.metrics.recall_at_k, a.metrics.search_latency_p50_us);
|
||||||
|
let b_key = (-b.metrics.recall_at_k, b.metrics.search_latency_p50_us);
|
||||||
|
a_key.partial_cmp(&b_key).unwrap_or(std::cmp::Ordering::Equal)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run the autotune loop. Synchronous — caller can wrap in tokio::spawn
|
||||||
|
/// to make it fire-and-forget. Uses the existing trial flow so every
|
||||||
|
/// autotune step lands in the trial journal.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
pub async fn run_autotune(
|
||||||
|
req: AutotuneRequest,
|
||||||
|
store: &Arc<dyn ObjectStore>,
|
||||||
|
catalog: &CatalogRegistry,
|
||||||
|
ai_client: &AiClient,
|
||||||
|
embedding_cache: &EmbeddingCache,
|
||||||
|
hnsw_store: &HnswStore,
|
||||||
|
index_registry: &IndexRegistry,
|
||||||
|
trial_journal: &TrialJournal,
|
||||||
|
promotion_registry: &PromotionRegistry,
|
||||||
|
_job_tracker: &JobTracker,
|
||||||
|
) -> Result<AutotuneResult, String> {
|
||||||
|
let t0 = std::time::Instant::now();
|
||||||
|
|
||||||
|
// Sanity: the index has to exist.
|
||||||
|
if index_registry.get(&req.index_name).await.is_none() {
|
||||||
|
return Err(format!("index not found: {}", req.index_name));
|
||||||
|
}
|
||||||
|
let _ = catalog; // reserved for future audit emission
|
||||||
|
|
||||||
|
// Load the harness once, compute ground truth once.
|
||||||
|
let mut harness_set = harness::EvalSet::load(store, &req.harness)
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("load harness: {e}"))?;
|
||||||
|
|
||||||
|
let embeddings = embedding_cache
|
||||||
|
.get_or_load(&req.index_name)
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("embeddings: {e}"))?;
|
||||||
|
|
||||||
|
if !harness_set.ground_truth_built {
|
||||||
|
harness::compute_ground_truth(&mut harness_set, &embeddings, ai_client)
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("ground truth: {e}"))?;
|
||||||
|
harness_set.save(store).await.ok();
|
||||||
|
}
|
||||||
|
|
||||||
|
let (grid, rejected) = sanitize_grid(req.grid.clone().unwrap_or_else(default_grid));
|
||||||
|
if grid.is_empty() {
|
||||||
|
return Err("all configs rejected by bounds check".into());
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
"autotune '{}': {} trial configs, min_recall={}",
|
||||||
|
req.index_name, grid.len(), req.min_recall,
|
||||||
|
);
|
||||||
|
|
||||||
|
// Run each trial through the existing pipeline so every run lands
|
||||||
|
// in the trial journal and contributes to the long-term tuning
|
||||||
|
// memory. This is the llms3.com-style "trials are data" pattern.
|
||||||
|
let mut run_trials: Vec<Trial> = Vec::with_capacity(grid.len());
|
||||||
|
for config in &grid {
|
||||||
|
let trial = run_single_trial(
|
||||||
|
&req.index_name,
|
||||||
|
&harness_set,
|
||||||
|
&embeddings,
|
||||||
|
config,
|
||||||
|
hnsw_store,
|
||||||
|
trial_journal,
|
||||||
|
ai_client,
|
||||||
|
).await;
|
||||||
|
match trial {
|
||||||
|
Ok(t) => run_trials.push(t),
|
||||||
|
Err(e) => tracing::warn!("autotune: trial ec={} es={} failed: {e}",
|
||||||
|
config.ef_construction, config.ef_search),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let trials_eligible = run_trials
|
||||||
|
.iter()
|
||||||
|
.filter(|t| t.metrics.recall_at_k >= req.min_recall)
|
||||||
|
.count();
|
||||||
|
|
||||||
|
let winner = pick_winner(&run_trials, req.min_recall);
|
||||||
|
let mut promoted = false;
|
||||||
|
|
||||||
|
let winner_report = if let Some(w) = winner {
|
||||||
|
let entry = PromotionEntry {
|
||||||
|
config: w.config.clone(),
|
||||||
|
trial_id: w.id.clone(),
|
||||||
|
promoted_at: Utc::now(),
|
||||||
|
promoted_by: "autotune".to_string(),
|
||||||
|
note: Some(format!(
|
||||||
|
"autotune winner: recall={:.3} p50={:.0}us",
|
||||||
|
w.metrics.recall_at_k, w.metrics.search_latency_p50_us,
|
||||||
|
)),
|
||||||
|
};
|
||||||
|
if let Err(e) = promotion_registry.promote(&req.index_name, entry).await {
|
||||||
|
tracing::warn!("autotune: promote failed: {e}");
|
||||||
|
} else {
|
||||||
|
promoted = true;
|
||||||
|
}
|
||||||
|
Some(WinnerReport {
|
||||||
|
trial_id: w.id.clone(),
|
||||||
|
config: w.config.clone(),
|
||||||
|
recall_at_k: w.metrics.recall_at_k,
|
||||||
|
search_latency_p50_us: w.metrics.search_latency_p50_us,
|
||||||
|
build_time_secs: w.metrics.build_time_secs,
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
tracing::warn!(
|
||||||
|
"autotune '{}': no trial met min_recall={}; not promoting",
|
||||||
|
req.index_name, req.min_recall,
|
||||||
|
);
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(AutotuneResult {
|
||||||
|
index_name: req.index_name,
|
||||||
|
trials_run: run_trials.len(),
|
||||||
|
trials_eligible,
|
||||||
|
winner: winner_report,
|
||||||
|
promoted,
|
||||||
|
duration_secs: t0.elapsed().as_secs_f32(),
|
||||||
|
skipped: rejected,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run one trial and journal it. Returns the recorded Trial.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
async fn run_single_trial(
|
||||||
|
index_name: &str,
|
||||||
|
harness_set: &harness::EvalSet,
|
||||||
|
embeddings: &Arc<Vec<crate::store::StoredEmbedding>>,
|
||||||
|
config: &HnswConfig,
|
||||||
|
hnsw_store: &HnswStore,
|
||||||
|
trial_journal: &TrialJournal,
|
||||||
|
_ai_client: &AiClient,
|
||||||
|
) -> Result<Trial, String> {
|
||||||
|
let trial_id = Trial::new_id();
|
||||||
|
let slot = format!("{}__{}", index_name, trial_id);
|
||||||
|
|
||||||
|
let build = hnsw_store
|
||||||
|
.build_index_with_config(&slot, (**embeddings).clone(), config)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let query_vectors: Vec<Vec<f32>> = harness_set
|
||||||
|
.queries
|
||||||
|
.iter()
|
||||||
|
.filter_map(|q| q.query_embedding.clone())
|
||||||
|
.collect();
|
||||||
|
let bench = hnsw_store
|
||||||
|
.bench_search(&slot, &query_vectors, harness_set.k)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let mut recalls = Vec::with_capacity(harness_set.queries.len());
|
||||||
|
for (q, hits) in harness_set.queries.iter().zip(bench.retrieved.iter()) {
|
||||||
|
if let Some(gt) = &q.ground_truth {
|
||||||
|
recalls.push(harness::recall_at_k(hits, gt, harness_set.k));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mean_recall = if recalls.is_empty() {
|
||||||
|
0.0
|
||||||
|
} else {
|
||||||
|
recalls.iter().sum::<f32>() / recalls.len() as f32
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut lats = bench.latencies_us.clone();
|
||||||
|
lats.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
let p = |pct: f32| -> f32 {
|
||||||
|
if lats.is_empty() { return 0.0; }
|
||||||
|
let idx = ((lats.len() as f32 - 1.0) * pct).round() as usize;
|
||||||
|
lats[idx.min(lats.len() - 1)]
|
||||||
|
};
|
||||||
|
|
||||||
|
let brute_us = if let Some(qv) = query_vectors.first() {
|
||||||
|
let t = std::time::Instant::now();
|
||||||
|
let _ = harness::brute_force_top_k(qv, embeddings, harness_set.k);
|
||||||
|
t.elapsed().as_micros() as f32
|
||||||
|
} else {
|
||||||
|
0.0
|
||||||
|
};
|
||||||
|
|
||||||
|
let dims = embeddings.first().map(|e| e.vector.len()).unwrap_or(0);
|
||||||
|
let memory_bytes =
|
||||||
|
(embeddings.len() * dims * std::mem::size_of::<f32>() + embeddings.len() * 128) as u64;
|
||||||
|
|
||||||
|
let trial = Trial {
|
||||||
|
id: trial_id.clone(),
|
||||||
|
index_name: index_name.to_string(),
|
||||||
|
eval_set: harness_set.name.clone(),
|
||||||
|
config: config.clone(),
|
||||||
|
metrics: TrialMetrics {
|
||||||
|
build_time_secs: build.build_time_secs,
|
||||||
|
search_latency_p50_us: p(0.50),
|
||||||
|
search_latency_p95_us: p(0.95),
|
||||||
|
search_latency_p99_us: p(0.99),
|
||||||
|
recall_at_k: mean_recall,
|
||||||
|
memory_bytes,
|
||||||
|
vectors: build.vectors,
|
||||||
|
eval_queries: harness_set.queries.len(),
|
||||||
|
brute_force_latency_us: brute_us,
|
||||||
|
},
|
||||||
|
created_at: Utc::now(),
|
||||||
|
note: Some("autotune".to_string()),
|
||||||
|
};
|
||||||
|
let _ = trial_journal.append(&trial).await;
|
||||||
|
let _ = hnsw_store.drop(&slot).await;
|
||||||
|
Ok(trial)
|
||||||
|
}
|
||||||
@@ -0,0 +1,224 @@
|
|||||||
|
/// Text chunking strategies for embedding.
|
||||||
|
/// Chunks need to be small enough for the embedding model (typically <512 tokens)
|
||||||
|
/// but large enough to carry meaning.
|
||||||
|
|
||||||
|
/// A chunk of text with metadata pointing back to its source.
|
||||||
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct TextChunk {
|
||||||
|
/// Source identifier (dataset name, filename, etc.)
|
||||||
|
pub source: String,
|
||||||
|
/// Row or document ID within the source
|
||||||
|
pub doc_id: String,
|
||||||
|
/// Chunk index within the document (0, 1, 2, ...)
|
||||||
|
pub chunk_idx: u32,
|
||||||
|
/// The actual text content
|
||||||
|
pub text: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Round `idx` UP to the nearest UTF-8 char boundary in `s`.
|
||||||
|
/// Needed because `&str[a..b]` panics if `a` or `b` lands mid-character —
|
||||||
|
/// and any text imported from Postgres/user data will contain multi-byte
|
||||||
|
/// chars like `\u{202f}` (narrow no-break space), em-dashes, etc.
|
||||||
|
fn ceil_char_boundary(s: &str, mut idx: usize) -> usize {
|
||||||
|
if idx >= s.len() { return s.len(); }
|
||||||
|
while idx <= s.len() && !s.is_char_boundary(idx) {
|
||||||
|
idx += 1;
|
||||||
|
}
|
||||||
|
idx
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Split text into overlapping chunks.
|
||||||
|
/// - `chunk_size`: target bytes per chunk (UTF-8 boundary-respecting)
|
||||||
|
/// - `overlap`: bytes of overlap between consecutive chunks
|
||||||
|
///
|
||||||
|
/// Safety contract:
|
||||||
|
/// - Never slices at a byte that isn't a UTF-8 char boundary
|
||||||
|
/// - Always advances forward (monotonic `start`) — no infinite loops on
|
||||||
|
/// pathological inputs
|
||||||
|
/// - Bounded iteration — caps total chunks at `text.len()` to prevent
|
||||||
|
/// runaway allocation on degenerate text/parameter combinations
|
||||||
|
pub fn chunk_text(
|
||||||
|
text: &str,
|
||||||
|
source: &str,
|
||||||
|
doc_id: &str,
|
||||||
|
chunk_size: usize,
|
||||||
|
overlap: usize,
|
||||||
|
) -> Vec<TextChunk> {
|
||||||
|
let text = text.trim();
|
||||||
|
if text.is_empty() {
|
||||||
|
return vec![];
|
||||||
|
}
|
||||||
|
|
||||||
|
// Short text — single chunk
|
||||||
|
if text.len() <= chunk_size {
|
||||||
|
return vec![TextChunk {
|
||||||
|
source: source.to_string(),
|
||||||
|
doc_id: doc_id.to_string(),
|
||||||
|
chunk_idx: 0,
|
||||||
|
text: text.to_string(),
|
||||||
|
}];
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut chunks = Vec::new();
|
||||||
|
let mut start = 0;
|
||||||
|
let mut idx = 0u32;
|
||||||
|
|
||||||
|
// Safety cap: a well-formed chunk_size advances `start` by at least
|
||||||
|
// `chunk_size - overlap` each iteration. Even in degenerate cases we
|
||||||
|
// should never emit more chunks than there are bytes of text.
|
||||||
|
let max_chunks = text.len();
|
||||||
|
let mut iterations = 0;
|
||||||
|
|
||||||
|
while start < text.len() {
|
||||||
|
iterations += 1;
|
||||||
|
if iterations > max_chunks {
|
||||||
|
tracing::warn!(
|
||||||
|
"chunker: exceeded safety cap ({}) on doc '{}' of length {} — stopping",
|
||||||
|
max_chunks, doc_id, text.len(),
|
||||||
|
);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
// UTF-8 safe upper bound: move forward to a char boundary.
|
||||||
|
let end = ceil_char_boundary(text, (start + chunk_size).min(text.len()));
|
||||||
|
|
||||||
|
let chunk_text_slice = &text[start..end];
|
||||||
|
let actual_end = if end < text.len() {
|
||||||
|
// Look for last sentence boundary inside the slice.
|
||||||
|
if let Some(pos) = chunk_text_slice.rfind(". ") {
|
||||||
|
start + pos + 2
|
||||||
|
} else if let Some(pos) = chunk_text_slice.rfind('\n') {
|
||||||
|
start + pos + 1
|
||||||
|
} else if let Some(pos) = chunk_text_slice.rfind(' ') {
|
||||||
|
start + pos + 1
|
||||||
|
} else {
|
||||||
|
end
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
end
|
||||||
|
};
|
||||||
|
// All the rfind'd positions are at char boundaries (they match
|
||||||
|
// ASCII byte patterns), so actual_end should already be safe.
|
||||||
|
// Defensive belt-and-suspenders: round up anyway.
|
||||||
|
let actual_end = ceil_char_boundary(text, actual_end);
|
||||||
|
|
||||||
|
let chunk = text[start..actual_end].trim();
|
||||||
|
if !chunk.is_empty() {
|
||||||
|
chunks.push(TextChunk {
|
||||||
|
source: source.to_string(),
|
||||||
|
doc_id: doc_id.to_string(),
|
||||||
|
chunk_idx: idx,
|
||||||
|
text: chunk.to_string(),
|
||||||
|
});
|
||||||
|
idx += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Advance with overlap — must strictly progress to avoid infinite
|
||||||
|
// loops on degenerate inputs (e.g. chunk_size ≤ overlap, or
|
||||||
|
// boundary-finding returning a position at `start`).
|
||||||
|
if actual_end >= text.len() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
let tentative = if actual_end > overlap { actual_end - overlap } else { actual_end };
|
||||||
|
let next_start = ceil_char_boundary(text, tentative);
|
||||||
|
if next_start <= start {
|
||||||
|
// No progress — force forward one char to guarantee termination.
|
||||||
|
let mut forced = start + 1;
|
||||||
|
while forced < text.len() && !text.is_char_boundary(forced) {
|
||||||
|
forced += 1;
|
||||||
|
}
|
||||||
|
start = forced;
|
||||||
|
} else {
|
||||||
|
start = next_start;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
chunks
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Chunk a dataset's text column. Returns all chunks from all rows.
|
||||||
|
pub fn chunk_column(
|
||||||
|
source: &str,
|
||||||
|
doc_ids: &[String],
|
||||||
|
texts: &[String],
|
||||||
|
chunk_size: usize,
|
||||||
|
overlap: usize,
|
||||||
|
) -> Vec<TextChunk> {
|
||||||
|
let mut all_chunks = Vec::new();
|
||||||
|
for (doc_id, text) in doc_ids.iter().zip(texts.iter()) {
|
||||||
|
let chunks = chunk_text(text, source, doc_id, chunk_size, overlap);
|
||||||
|
all_chunks.extend(chunks);
|
||||||
|
}
|
||||||
|
all_chunks
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn short_text_single_chunk() {
|
||||||
|
let chunks = chunk_text("Hello world", "test", "1", 500, 50);
|
||||||
|
assert_eq!(chunks.len(), 1);
|
||||||
|
assert_eq!(chunks[0].text, "Hello world");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn long_text_multiple_chunks() {
|
||||||
|
let text = "First sentence. Second sentence. Third sentence. Fourth sentence. Fifth sentence. Sixth sentence. Seventh sentence. Eighth sentence.";
|
||||||
|
let chunks = chunk_text(text, "test", "1", 50, 10);
|
||||||
|
assert!(chunks.len() > 1);
|
||||||
|
// All text should be covered
|
||||||
|
for chunk in &chunks {
|
||||||
|
assert!(!chunk.text.is_empty());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn chunk_preserves_source() {
|
||||||
|
let chunks = chunk_text("Some text here", "candidates", "CAND-001", 500, 50);
|
||||||
|
assert_eq!(chunks[0].source, "candidates");
|
||||||
|
assert_eq!(chunks[0].doc_id, "CAND-001");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn handles_multibyte_utf8_at_chunk_boundary() {
|
||||||
|
// Regression: multi-byte unicode (narrow no-break space \u{202f},
|
||||||
|
// 3 bytes) landing inside a chunk-size window used to panic with
|
||||||
|
// "byte index N is not a char boundary". Found on pg-imported
|
||||||
|
// data that contained editorial unicode punctuation.
|
||||||
|
let text = "**Claim 1**\u{202f}\u{2013} *Each large language model has its own personality \
|
||||||
|
and way of interacting.*\u{202f}\u{2013}\u{202f}Verdict:\u{202f}VERIFIED. \
|
||||||
|
Reasoning: widely accepted observation that different models exhibit distinct response \
|
||||||
|
styles stemming from variations in training data.".repeat(5);
|
||||||
|
let chunks = chunk_text(&text, "test", "unicode-1", 500, 50);
|
||||||
|
// Must not panic, must produce at least one chunk.
|
||||||
|
assert!(!chunks.is_empty());
|
||||||
|
// All chunks must be valid UTF-8 substrings of the original.
|
||||||
|
for chunk in &chunks {
|
||||||
|
assert!(!chunk.text.is_empty());
|
||||||
|
assert!(text.contains(&chunk.text));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn no_infinite_loop_on_no_spaces() {
|
||||||
|
// Long text with no spaces, sentences, or newlines — chunker must
|
||||||
|
// still advance via the hard `end` fallback.
|
||||||
|
let text = "a".repeat(5000);
|
||||||
|
let chunks = chunk_text(&text, "test", "noSpace", 500, 50);
|
||||||
|
assert!(!chunks.is_empty());
|
||||||
|
// Expected ~5000/450 chunks (with overlap=50 → step=450)
|
||||||
|
assert!(chunks.len() < 50, "chunk count {} seems excessive", chunks.len());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn no_infinite_loop_on_degenerate_params() {
|
||||||
|
// chunk_size equals overlap — previously would have made zero
|
||||||
|
// forward progress. Our safety forced-advance must catch this.
|
||||||
|
let text = "a b c d e f g h i j k l m n o p q r s t u v w x y z".repeat(20);
|
||||||
|
let chunks = chunk_text(&text, "test", "degen", 100, 100);
|
||||||
|
assert!(!chunks.is_empty());
|
||||||
|
assert!(chunks.len() <= text.len(), "runaway chunk count {}", chunks.len());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,98 @@
|
|||||||
|
/// In-memory cache for StoredEmbedding vectors.
|
||||||
|
///
|
||||||
|
/// Rationale: loading 100K embeddings from Parquet takes ~2-5s. When an AI agent
|
||||||
|
/// iterates on HNSW parameters, each trial would repeat that cost. The cache
|
||||||
|
/// pins embeddings in memory so trials reuse them.
|
||||||
|
///
|
||||||
|
/// This is a pure performance layer — the Parquet file is still the source of
|
||||||
|
/// truth (ADR-008). Eviction is safe; worst case is one slow reload.
|
||||||
|
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use serde::Serialize;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
use crate::store::{self, StoredEmbedding};
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct EmbeddingCache {
|
||||||
|
store: Arc<dyn ObjectStore>,
|
||||||
|
cache: Arc<RwLock<HashMap<String, Arc<Vec<StoredEmbedding>>>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct CacheEntry {
|
||||||
|
pub index_name: String,
|
||||||
|
pub vectors: usize,
|
||||||
|
pub dimensions: usize,
|
||||||
|
pub memory_bytes: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct CacheStats {
|
||||||
|
pub entries: Vec<CacheEntry>,
|
||||||
|
pub total_memory_bytes: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl EmbeddingCache {
|
||||||
|
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||||
|
Self {
|
||||||
|
store,
|
||||||
|
cache: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Return cached embeddings, loading from object storage on first request.
|
||||||
|
pub async fn get_or_load(
|
||||||
|
&self,
|
||||||
|
index_name: &str,
|
||||||
|
) -> Result<Arc<Vec<StoredEmbedding>>, String> {
|
||||||
|
if let Some(cached) = self.cache.read().await.get(index_name) {
|
||||||
|
return Ok(cached.clone());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Load under a write lock so concurrent callers only hit disk once.
|
||||||
|
let mut guard = self.cache.write().await;
|
||||||
|
if let Some(cached) = guard.get(index_name) {
|
||||||
|
return Ok(cached.clone());
|
||||||
|
}
|
||||||
|
tracing::info!("embedding_cache: loading '{index_name}' from object storage");
|
||||||
|
let t0 = std::time::Instant::now();
|
||||||
|
let loaded = store::load_embeddings(&self.store, index_name).await?;
|
||||||
|
let n = loaded.len();
|
||||||
|
let arc = Arc::new(loaded);
|
||||||
|
guard.insert(index_name.to_string(), arc.clone());
|
||||||
|
tracing::info!(
|
||||||
|
"embedding_cache: loaded '{index_name}' — {n} vectors in {:.2}s",
|
||||||
|
t0.elapsed().as_secs_f32()
|
||||||
|
);
|
||||||
|
Ok(arc)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn evict(&self, index_name: &str) -> bool {
|
||||||
|
self.cache.write().await.remove(index_name).is_some()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn stats(&self) -> CacheStats {
|
||||||
|
let cache = self.cache.read().await;
|
||||||
|
let mut entries = Vec::with_capacity(cache.len());
|
||||||
|
let mut total: u64 = 0;
|
||||||
|
for (name, embs) in cache.iter() {
|
||||||
|
let dims = embs.first().map(|e| e.vector.len()).unwrap_or(0);
|
||||||
|
// Rough estimate: vector data + chunk_text + metadata overhead.
|
||||||
|
let vec_bytes = (embs.len() * dims * std::mem::size_of::<f32>()) as u64;
|
||||||
|
let text_bytes: u64 = embs.iter().map(|e| e.chunk_text.len() as u64).sum();
|
||||||
|
let overhead = (embs.len() * 128) as u64; // strings + struct overhead
|
||||||
|
let mem = vec_bytes + text_bytes + overhead;
|
||||||
|
total += mem;
|
||||||
|
entries.push(CacheEntry {
|
||||||
|
index_name: name.clone(),
|
||||||
|
vectors: embs.len(),
|
||||||
|
dimensions: dims,
|
||||||
|
memory_bytes: mem,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
CacheStats { entries, total_memory_bytes: total }
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,216 @@
|
|||||||
|
/// Eval harness for HNSW tuning.
|
||||||
|
///
|
||||||
|
/// An EvalSet is a named collection of queries against a specific vector index.
|
||||||
|
/// Each query has a query_text and (optionally pre-computed) ground_truth
|
||||||
|
/// doc_ids — the "correct" top-k results that brute-force cosine returns.
|
||||||
|
/// HNSW trials are scored by how well they recreate that top-k (recall@k).
|
||||||
|
///
|
||||||
|
/// Storage: `_hnsw_evals/{name}.json` as a single JSON document. Small.
|
||||||
|
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use aibridge::client::{AiClient, EmbedRequest};
|
||||||
|
use storaged::ops;
|
||||||
|
|
||||||
|
use crate::store::StoredEmbedding;
|
||||||
|
|
||||||
|
/// A single eval query with optional pre-computed ground truth.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct EvalQuery {
|
||||||
|
pub id: String,
|
||||||
|
pub query_text: String,
|
||||||
|
/// Ordered list of doc_ids that brute-force returns as top-k.
|
||||||
|
/// `None` means the ground truth hasn't been computed yet; `compute_ground_truth`
|
||||||
|
/// fills it in against the current embedding set.
|
||||||
|
#[serde(default)]
|
||||||
|
pub ground_truth: Option<Vec<String>>,
|
||||||
|
/// Optional pre-computed query embedding. Filling this avoids re-embedding
|
||||||
|
/// the same query on every trial.
|
||||||
|
#[serde(default)]
|
||||||
|
pub query_embedding: Option<Vec<f32>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A named eval set tied to a specific vector index.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct EvalSet {
|
||||||
|
pub name: String,
|
||||||
|
pub index_name: String,
|
||||||
|
pub k: usize, // top-k used for recall calculation
|
||||||
|
pub queries: Vec<EvalQuery>,
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
pub ground_truth_built: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl EvalSet {
|
||||||
|
pub fn new(name: &str, index_name: &str, k: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
name: name.to_string(),
|
||||||
|
index_name: index_name.to_string(),
|
||||||
|
k,
|
||||||
|
queries: Vec::new(),
|
||||||
|
created_at: Utc::now(),
|
||||||
|
ground_truth_built: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn storage_key(&self) -> String {
|
||||||
|
format!("_hnsw_evals/{}.json", self.name)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn save(&self, store: &Arc<dyn ObjectStore>) -> Result<(), String> {
|
||||||
|
let json = serde_json::to_vec_pretty(self).map_err(|e| e.to_string())?;
|
||||||
|
ops::put(store, &self.storage_key(), json.into()).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn load(store: &Arc<dyn ObjectStore>, name: &str) -> Result<Self, String> {
|
||||||
|
let key = format!("_hnsw_evals/{}.json", name);
|
||||||
|
let data = ops::get(store, &key).await?;
|
||||||
|
serde_json::from_slice(&data).map_err(|e| format!("parse eval set: {e}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list(store: &Arc<dyn ObjectStore>) -> Result<Vec<String>, String> {
|
||||||
|
let keys = ops::list(store, Some("_hnsw_evals/")).await?;
|
||||||
|
Ok(keys
|
||||||
|
.into_iter()
|
||||||
|
.filter(|k| k.ends_with(".json"))
|
||||||
|
.map(|k| {
|
||||||
|
k.trim_start_matches("_hnsw_evals/")
|
||||||
|
.trim_end_matches(".json")
|
||||||
|
.to_string()
|
||||||
|
})
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Cosine similarity for two same-length f32 slices.
|
||||||
|
fn cosine(a: &[f32], b: &[f32]) -> f32 {
|
||||||
|
let mut dot = 0.0f32;
|
||||||
|
let mut na = 0.0f32;
|
||||||
|
let mut nb = 0.0f32;
|
||||||
|
for i in 0..a.len() {
|
||||||
|
dot += a[i] * b[i];
|
||||||
|
na += a[i] * a[i];
|
||||||
|
nb += b[i] * b[i];
|
||||||
|
}
|
||||||
|
if na == 0.0 || nb == 0.0 {
|
||||||
|
return 0.0;
|
||||||
|
}
|
||||||
|
dot / (na.sqrt() * nb.sqrt())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Brute-force top-k search for a single query against all embeddings.
|
||||||
|
/// This is the ground-truth oracle that HNSW trials must approximate.
|
||||||
|
pub fn brute_force_top_k(
|
||||||
|
query: &[f32],
|
||||||
|
embeddings: &[StoredEmbedding],
|
||||||
|
k: usize,
|
||||||
|
) -> Vec<String> {
|
||||||
|
let mut scored: Vec<(f32, usize)> = embeddings
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, e)| (cosine(query, &e.vector), i))
|
||||||
|
.collect();
|
||||||
|
// Partial sort — we only need top-k.
|
||||||
|
scored.sort_unstable_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
scored
|
||||||
|
.into_iter()
|
||||||
|
.take(k)
|
||||||
|
.map(|(_, i)| embeddings[i].doc_id.clone())
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Recall@k — fraction of ground-truth doc_ids that appear in predicted.
|
||||||
|
pub fn recall_at_k(predicted: &[String], ground_truth: &[String], k: usize) -> f32 {
|
||||||
|
if ground_truth.is_empty() || k == 0 {
|
||||||
|
return 0.0;
|
||||||
|
}
|
||||||
|
let gt_set: std::collections::HashSet<&String> =
|
||||||
|
ground_truth.iter().take(k).collect();
|
||||||
|
let hits = predicted
|
||||||
|
.iter()
|
||||||
|
.take(k)
|
||||||
|
.filter(|d| gt_set.contains(d))
|
||||||
|
.count();
|
||||||
|
hits as f32 / gt_set.len() as f32
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Populate query_embedding and ground_truth for every query that lacks them.
|
||||||
|
pub async fn compute_ground_truth(
|
||||||
|
eval: &mut EvalSet,
|
||||||
|
embeddings: &[StoredEmbedding],
|
||||||
|
ai_client: &AiClient,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let need_embed: Vec<(usize, String)> = eval
|
||||||
|
.queries
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.filter(|(_, q)| q.query_embedding.is_none())
|
||||||
|
.map(|(i, q)| (i, q.query_text.clone()))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
if !need_embed.is_empty() {
|
||||||
|
// Embed in one batch to keep things simple; for very large eval sets
|
||||||
|
// we'd batch this in chunks of 32.
|
||||||
|
let texts: Vec<String> = need_embed.iter().map(|(_, t)| t.clone()).collect();
|
||||||
|
let resp = ai_client
|
||||||
|
.embed(EmbedRequest { texts, model: None })
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("embed queries: {e}"))?;
|
||||||
|
for ((idx, _), vec) in need_embed.iter().zip(resp.embeddings.iter()) {
|
||||||
|
let v: Vec<f32> = vec.iter().map(|&x| x as f32).collect();
|
||||||
|
eval.queries[*idx].query_embedding = Some(v);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for q in eval.queries.iter_mut() {
|
||||||
|
if q.ground_truth.is_some() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let emb = q.query_embedding.as_ref().ok_or("missing embedding")?;
|
||||||
|
q.ground_truth = Some(brute_force_top_k(emb, embeddings, eval.k));
|
||||||
|
}
|
||||||
|
eval.ground_truth_built = true;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Auto-generate a synthetic eval set by sampling every Nth chunk's text as
|
||||||
|
/// its own query. Useful for a quick-start eval when the user doesn't have
|
||||||
|
/// real natural-language queries yet.
|
||||||
|
pub fn synthetic_from_chunks(
|
||||||
|
eval_name: &str,
|
||||||
|
index_name: &str,
|
||||||
|
embeddings: &[StoredEmbedding],
|
||||||
|
sample_count: usize,
|
||||||
|
k: usize,
|
||||||
|
) -> EvalSet {
|
||||||
|
let n = embeddings.len();
|
||||||
|
let sample_count = sample_count.min(n);
|
||||||
|
let stride = (n / sample_count.max(1)).max(1);
|
||||||
|
|
||||||
|
let mut queries = Vec::with_capacity(sample_count);
|
||||||
|
for i in 0..sample_count {
|
||||||
|
let idx = (i * stride).min(n - 1);
|
||||||
|
let chunk = &embeddings[idx];
|
||||||
|
// Use the first ~200 chars of the chunk as the "query" — it should find
|
||||||
|
// itself and nearby chunks as top results.
|
||||||
|
let query_text: String = chunk.chunk_text.chars().take(200).collect();
|
||||||
|
queries.push(EvalQuery {
|
||||||
|
id: format!("syn-{}", i),
|
||||||
|
query_text,
|
||||||
|
ground_truth: None,
|
||||||
|
query_embedding: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
EvalSet {
|
||||||
|
name: eval_name.to_string(),
|
||||||
|
index_name: index_name.to_string(),
|
||||||
|
k,
|
||||||
|
queries,
|
||||||
|
created_at: Utc::now(),
|
||||||
|
ground_truth_built: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,229 @@
|
|||||||
|
/// HNSW (Hierarchical Navigable Small World) index for fast approximate nearest neighbor search.
|
||||||
|
/// Wraps instant-distance to provide <50ms search over 100K+ vectors.
|
||||||
|
/// Falls back to brute-force for small datasets.
|
||||||
|
|
||||||
|
use instant_distance::{Builder, HnswMap, Search, Point};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
use crate::store::StoredEmbedding;
|
||||||
|
use crate::trial::HnswConfig;
|
||||||
|
|
||||||
|
/// A vector point for HNSW — wraps f32 slice with cosine distance.
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct VecPoint(Vec<f32>);
|
||||||
|
|
||||||
|
impl Point for VecPoint {
|
||||||
|
fn distance(&self, other: &Self) -> f32 {
|
||||||
|
// Cosine distance = 1 - cosine_similarity
|
||||||
|
let dot: f32 = self.0.iter().zip(other.0.iter()).map(|(a, b)| a * b).sum();
|
||||||
|
let norm_a: f32 = self.0.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||||
|
let norm_b: f32 = other.0.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||||
|
if norm_a == 0.0 || norm_b == 0.0 {
|
||||||
|
return 1.0;
|
||||||
|
}
|
||||||
|
1.0 - (dot / (norm_a * norm_b))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// HNSW search result.
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct HnswResult {
|
||||||
|
pub source: String,
|
||||||
|
pub doc_id: String,
|
||||||
|
pub chunk_idx: u32,
|
||||||
|
pub chunk_text: String,
|
||||||
|
pub score: f32, // cosine similarity (1.0 = identical)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// An HNSW index built from stored embeddings.
|
||||||
|
pub struct HnswIndex {
|
||||||
|
map: HnswMap<VecPoint, usize>, // value is index into metadata vec
|
||||||
|
metadata: Vec<EmbeddingMeta>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct EmbeddingMeta {
|
||||||
|
source: String,
|
||||||
|
doc_id: String,
|
||||||
|
chunk_idx: u32,
|
||||||
|
chunk_text: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Shared, thread-safe index store.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct HnswStore {
|
||||||
|
indexes: Arc<RwLock<std::collections::HashMap<String, Arc<HnswIndex>>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl HnswStore {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
indexes: Arc::new(RwLock::new(std::collections::HashMap::new())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build an HNSW index from stored embeddings with default config.
|
||||||
|
pub async fn build_index(
|
||||||
|
&self,
|
||||||
|
index_name: &str,
|
||||||
|
embeddings: Vec<StoredEmbedding>,
|
||||||
|
) -> Result<BuildStats, String> {
|
||||||
|
self.build_index_with_config(index_name, embeddings, &HnswConfig::default()).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build an HNSW index from stored embeddings with explicit config.
|
||||||
|
/// Used by the trial system — each trial calls this with different params.
|
||||||
|
pub async fn build_index_with_config(
|
||||||
|
&self,
|
||||||
|
index_name: &str,
|
||||||
|
embeddings: Vec<StoredEmbedding>,
|
||||||
|
config: &HnswConfig,
|
||||||
|
) -> Result<BuildStats, String> {
|
||||||
|
let n = embeddings.len();
|
||||||
|
if n == 0 {
|
||||||
|
return Err("no embeddings to index".into());
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
"building HNSW '{}': {} vectors, ef_construction={} ef_search={} seed={:?}",
|
||||||
|
index_name, n, config.ef_construction, config.ef_search, config.seed,
|
||||||
|
);
|
||||||
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
|
// Separate points and metadata
|
||||||
|
let mut points = Vec::with_capacity(n);
|
||||||
|
let mut metadata = Vec::with_capacity(n);
|
||||||
|
let mut values = Vec::with_capacity(n);
|
||||||
|
|
||||||
|
for (i, emb) in embeddings.into_iter().enumerate() {
|
||||||
|
points.push(VecPoint(emb.vector));
|
||||||
|
metadata.push(EmbeddingMeta {
|
||||||
|
source: emb.source,
|
||||||
|
doc_id: emb.doc_id,
|
||||||
|
chunk_idx: emb.chunk_idx,
|
||||||
|
chunk_text: emb.chunk_text,
|
||||||
|
});
|
||||||
|
values.push(i);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build HNSW — the expensive part
|
||||||
|
let mut builder = Builder::default()
|
||||||
|
.ef_construction(config.ef_construction)
|
||||||
|
.ef_search(config.ef_search);
|
||||||
|
if let Some(seed) = config.seed {
|
||||||
|
builder = builder.seed(seed);
|
||||||
|
}
|
||||||
|
let map = builder.build(points, values);
|
||||||
|
|
||||||
|
let build_time = start.elapsed().as_secs_f32();
|
||||||
|
tracing::info!("HNSW '{}' built: {} vectors in {:.1}s", index_name, n, build_time);
|
||||||
|
|
||||||
|
let index = Arc::new(HnswIndex { map, metadata });
|
||||||
|
self.indexes.write().await.insert(index_name.to_string(), index);
|
||||||
|
|
||||||
|
Ok(BuildStats {
|
||||||
|
index_name: index_name.to_string(),
|
||||||
|
vectors: n,
|
||||||
|
build_time_secs: build_time,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run a batch of search queries and return raw per-query latencies in microseconds.
|
||||||
|
/// Also returns the retrieved doc_ids per query (for recall calculation).
|
||||||
|
pub async fn bench_search(
|
||||||
|
&self,
|
||||||
|
index_name: &str,
|
||||||
|
query_vectors: &[Vec<f32>],
|
||||||
|
top_k: usize,
|
||||||
|
) -> Result<BenchResult, String> {
|
||||||
|
let indexes = self.indexes.read().await;
|
||||||
|
let index = indexes
|
||||||
|
.get(index_name)
|
||||||
|
.ok_or_else(|| format!("HNSW index not found: {index_name}"))?
|
||||||
|
.clone();
|
||||||
|
drop(indexes);
|
||||||
|
|
||||||
|
let mut latencies_us = Vec::with_capacity(query_vectors.len());
|
||||||
|
let mut retrieved: Vec<Vec<String>> = Vec::with_capacity(query_vectors.len());
|
||||||
|
|
||||||
|
for qv in query_vectors {
|
||||||
|
let query_point = VecPoint(qv.clone());
|
||||||
|
let t0 = std::time::Instant::now();
|
||||||
|
let mut search = Search::default();
|
||||||
|
let results = index.map.search(&query_point, &mut search);
|
||||||
|
let ids: Vec<String> = results
|
||||||
|
.take(top_k)
|
||||||
|
.map(|item| {
|
||||||
|
let meta_idx = *item.value;
|
||||||
|
index.metadata[meta_idx].doc_id.clone()
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
latencies_us.push(t0.elapsed().as_micros() as f32);
|
||||||
|
retrieved.push(ids);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(BenchResult { latencies_us, retrieved })
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Search an HNSW index. Returns approximate nearest neighbors.
|
||||||
|
pub async fn search(
|
||||||
|
&self,
|
||||||
|
index_name: &str,
|
||||||
|
query: &[f32],
|
||||||
|
top_k: usize,
|
||||||
|
) -> Result<Vec<HnswResult>, String> {
|
||||||
|
let indexes = self.indexes.read().await;
|
||||||
|
let index = indexes.get(index_name)
|
||||||
|
.ok_or_else(|| format!("HNSW index not found: {index_name}"))?;
|
||||||
|
|
||||||
|
let query_point = VecPoint(query.to_vec());
|
||||||
|
let mut search = Search::default();
|
||||||
|
let results = index.map.search(&query_point, &mut search);
|
||||||
|
|
||||||
|
let mut out = Vec::with_capacity(top_k);
|
||||||
|
for item in results.take(top_k) {
|
||||||
|
let meta_idx = *item.value;
|
||||||
|
let meta = &index.metadata[meta_idx];
|
||||||
|
// Convert distance back to similarity
|
||||||
|
let similarity = 1.0 - item.distance;
|
||||||
|
out.push(HnswResult {
|
||||||
|
source: meta.source.clone(),
|
||||||
|
doc_id: meta.doc_id.clone(),
|
||||||
|
chunk_idx: meta.chunk_idx,
|
||||||
|
chunk_text: meta.chunk_text.clone(),
|
||||||
|
score: similarity,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(out)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if an index exists.
|
||||||
|
pub async fn has_index(&self, name: &str) -> bool {
|
||||||
|
self.indexes.read().await.contains_key(name)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all loaded indexes.
|
||||||
|
pub async fn list(&self) -> Vec<String> {
|
||||||
|
self.indexes.read().await.keys().cloned().collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Drop an index from memory.
|
||||||
|
pub async fn drop(&self, name: &str) -> bool {
|
||||||
|
self.indexes.write().await.remove(name).is_some()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct BuildStats {
|
||||||
|
pub index_name: String,
|
||||||
|
pub vectors: usize,
|
||||||
|
pub build_time_secs: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct BenchResult {
|
||||||
|
pub latencies_us: Vec<f32>,
|
||||||
|
pub retrieved: Vec<Vec<String>>,
|
||||||
|
}
|
||||||
@@ -0,0 +1,125 @@
|
|||||||
|
/// Vector index registry — tracks all indexes with model versioning.
|
||||||
|
/// Each index knows which model created it, enabling:
|
||||||
|
/// - Multi-version indexes (same data, different models, coexist)
|
||||||
|
/// - Incremental re-embed (only new/changed docs on model upgrade)
|
||||||
|
/// - A/B search comparison between model versions
|
||||||
|
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
use storaged::ops;
|
||||||
|
|
||||||
|
/// Metadata for a vector index.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct IndexMeta {
|
||||||
|
pub index_name: String,
|
||||||
|
pub source: String, // dataset this was built from
|
||||||
|
pub model_name: String, // "nomic-embed-text"
|
||||||
|
pub model_version: String, // "latest" or specific version
|
||||||
|
pub dimensions: u32, // 768
|
||||||
|
pub chunk_count: usize,
|
||||||
|
pub doc_count: usize,
|
||||||
|
pub chunk_size: usize,
|
||||||
|
pub overlap: usize,
|
||||||
|
pub storage_key: String, // "vectors/resumes_v1_nomic.parquet"
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
pub build_time_secs: f32,
|
||||||
|
pub chunks_per_sec: f32,
|
||||||
|
/// Federation layer 2: which bucket holds this index's artifacts
|
||||||
|
/// (trial journal + promotion file). Defaults to "primary" for
|
||||||
|
/// pre-federation indexes — the serde default keeps old metadata
|
||||||
|
/// files readable without migration.
|
||||||
|
#[serde(default = "default_bucket")]
|
||||||
|
pub bucket: String,
|
||||||
|
/// ADR-019: which physical backend stores this index. `Parquet`
|
||||||
|
/// means storage_key points at our binary-blob Parquet file;
|
||||||
|
/// `Lance` means it points at a Lance dataset directory.
|
||||||
|
#[serde(default)]
|
||||||
|
pub vector_backend: shared::types::VectorBackend,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_bucket() -> String { "primary".to_string() }
|
||||||
|
|
||||||
|
/// Registry of all vector indexes.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct IndexRegistry {
|
||||||
|
indexes: Arc<RwLock<HashMap<String, IndexMeta>>>,
|
||||||
|
store: Arc<dyn ObjectStore>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl IndexRegistry {
|
||||||
|
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||||
|
Self {
|
||||||
|
indexes: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
store,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Rebuild from persisted index metadata on startup.
|
||||||
|
pub async fn rebuild(&self) -> Result<usize, String> {
|
||||||
|
let keys = ops::list(&self.store, Some("vectors/meta/")).await?;
|
||||||
|
let mut reg = self.indexes.write().await;
|
||||||
|
reg.clear();
|
||||||
|
|
||||||
|
for key in &keys {
|
||||||
|
if !key.ends_with(".json") { continue; }
|
||||||
|
let data = ops::get(&self.store, key).await?;
|
||||||
|
match serde_json::from_slice::<IndexMeta>(&data) {
|
||||||
|
Ok(meta) => { reg.insert(meta.index_name.clone(), meta); }
|
||||||
|
Err(e) => tracing::warn!("failed to load index meta {key}: {e}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let count = reg.len();
|
||||||
|
if count > 0 {
|
||||||
|
tracing::info!("loaded {count} vector index metadata entries");
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Register a new index.
|
||||||
|
pub async fn register(&self, meta: IndexMeta) -> Result<(), String> {
|
||||||
|
let key = format!("vectors/meta/{}.json", meta.index_name);
|
||||||
|
let json = serde_json::to_vec_pretty(&meta).map_err(|e| e.to_string())?;
|
||||||
|
ops::put(&self.store, &key, json.into()).await?;
|
||||||
|
self.indexes.write().await.insert(meta.index_name.clone(), meta);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get metadata for an index.
|
||||||
|
pub async fn get(&self, index_name: &str) -> Option<IndexMeta> {
|
||||||
|
self.indexes.read().await.get(index_name).cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all indexes, optionally filtered by source or model.
|
||||||
|
pub async fn list(&self, source: Option<&str>, model: Option<&str>) -> Vec<IndexMeta> {
|
||||||
|
self.indexes.read().await.values()
|
||||||
|
.filter(|m| source.map_or(true, |s| m.source == s))
|
||||||
|
.filter(|m| model.map_or(true, |mo| m.model_name == mo))
|
||||||
|
.cloned()
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Find all versions of an index for a given source dataset.
|
||||||
|
/// Returns indexes sorted by creation time (newest first).
|
||||||
|
pub async fn versions_for_source(&self, source: &str) -> Vec<IndexMeta> {
|
||||||
|
let mut versions: Vec<IndexMeta> = self.indexes.read().await.values()
|
||||||
|
.filter(|m| m.source == source)
|
||||||
|
.cloned()
|
||||||
|
.collect();
|
||||||
|
versions.sort_by(|a, b| b.created_at.cmp(&a.created_at));
|
||||||
|
versions
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Delete an index (metadata only — vector Parquet stays for safety).
|
||||||
|
pub async fn delete(&self, index_name: &str) -> Result<(), String> {
|
||||||
|
let key = format!("vectors/meta/{index_name}.json");
|
||||||
|
ops::delete(&self.store, &key).await?;
|
||||||
|
self.indexes.write().await.remove(index_name);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,112 @@
|
|||||||
|
/// Background job system for long-running embedding tasks.
|
||||||
|
/// POST /vectors/index returns a job_id immediately.
|
||||||
|
/// GET /vectors/jobs/{id} returns progress.
|
||||||
|
/// Embedding runs in background via tokio::spawn.
|
||||||
|
|
||||||
|
use serde::Serialize;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
#[serde(rename_all = "lowercase")]
|
||||||
|
pub enum JobStatus {
|
||||||
|
Running,
|
||||||
|
Completed,
|
||||||
|
Failed,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct Job {
|
||||||
|
pub id: String,
|
||||||
|
pub status: JobStatus,
|
||||||
|
pub index_name: String,
|
||||||
|
pub total_chunks: usize,
|
||||||
|
pub embedded_chunks: usize,
|
||||||
|
pub progress_pct: f32,
|
||||||
|
pub storage_key: Option<String>,
|
||||||
|
pub error: Option<String>,
|
||||||
|
pub started_at: String,
|
||||||
|
pub completed_at: Option<String>,
|
||||||
|
pub chunks_per_sec: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Shared progress tracker that background tasks update.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct JobTracker {
|
||||||
|
jobs: Arc<RwLock<HashMap<String, Job>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl JobTracker {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
jobs: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Register a new job. Returns the job ID.
|
||||||
|
pub async fn create(&self, index_name: &str, total_chunks: usize) -> String {
|
||||||
|
let id = format!("job-{}", chrono::Utc::now().timestamp_millis());
|
||||||
|
let job = Job {
|
||||||
|
id: id.clone(),
|
||||||
|
status: JobStatus::Running,
|
||||||
|
index_name: index_name.to_string(),
|
||||||
|
total_chunks,
|
||||||
|
embedded_chunks: 0,
|
||||||
|
progress_pct: 0.0,
|
||||||
|
storage_key: None,
|
||||||
|
error: None,
|
||||||
|
started_at: chrono::Utc::now().to_rfc3339(),
|
||||||
|
completed_at: None,
|
||||||
|
chunks_per_sec: 0.0,
|
||||||
|
};
|
||||||
|
self.jobs.write().await.insert(id.clone(), job);
|
||||||
|
id
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Update progress.
|
||||||
|
pub async fn update_progress(&self, id: &str, embedded: usize, rate: f32) {
|
||||||
|
let mut jobs = self.jobs.write().await;
|
||||||
|
if let Some(job) = jobs.get_mut(id) {
|
||||||
|
job.embedded_chunks = embedded;
|
||||||
|
job.progress_pct = if job.total_chunks > 0 {
|
||||||
|
(embedded as f32 / job.total_chunks as f32) * 100.0
|
||||||
|
} else {
|
||||||
|
0.0
|
||||||
|
};
|
||||||
|
job.chunks_per_sec = rate;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mark job as completed.
|
||||||
|
pub async fn complete(&self, id: &str, storage_key: String) {
|
||||||
|
let mut jobs = self.jobs.write().await;
|
||||||
|
if let Some(job) = jobs.get_mut(id) {
|
||||||
|
job.status = JobStatus::Completed;
|
||||||
|
job.embedded_chunks = job.total_chunks;
|
||||||
|
job.progress_pct = 100.0;
|
||||||
|
job.storage_key = Some(storage_key);
|
||||||
|
job.completed_at = Some(chrono::Utc::now().to_rfc3339());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mark job as failed.
|
||||||
|
pub async fn fail(&self, id: &str, error: String) {
|
||||||
|
let mut jobs = self.jobs.write().await;
|
||||||
|
if let Some(job) = jobs.get_mut(id) {
|
||||||
|
job.status = JobStatus::Failed;
|
||||||
|
job.error = Some(error);
|
||||||
|
job.completed_at = Some(chrono::Utc::now().to_rfc3339());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get job status.
|
||||||
|
pub async fn get(&self, id: &str) -> Option<Job> {
|
||||||
|
self.jobs.read().await.get(id).cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all jobs.
|
||||||
|
pub async fn list(&self) -> Vec<Job> {
|
||||||
|
self.jobs.read().await.values().cloned().collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,131 @@
|
|||||||
|
//! ADR-019 hybrid: routing layer for the Lance vector backend.
|
||||||
|
//!
|
||||||
|
//! Holds a `LanceRegistry` which maps an index name to its
|
||||||
|
//! `LanceVectorStore` instance, lazy-creating on first touch. Path
|
||||||
|
//! resolution: the index's bucket gives us a bucket root; we append
|
||||||
|
//! `lance/{index_name}` and use that as the dataset URI. For local
|
||||||
|
//! buckets that means a directory under the bucket's root.
|
||||||
|
//!
|
||||||
|
//! S3 buckets: in principle Lance accepts s3://... URIs, but the
|
||||||
|
//! firewall crate (`vectord-lance`) doesn't pull the S3 feature in by
|
||||||
|
//! default. When we promote profile buckets onto S3 in a future phase,
|
||||||
|
//! enable that feature and update `lance_uri_for` accordingly.
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::path::PathBuf;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
use storaged::registry::BucketRegistry;
|
||||||
|
use vectord_lance::LanceVectorStore;
|
||||||
|
|
||||||
|
use crate::index_registry::IndexRegistry;
|
||||||
|
|
||||||
|
/// Convert a bucket+index pair into the URI Lance should use as the
|
||||||
|
/// dataset path. Supports both local (filesystem) and S3 buckets.
|
||||||
|
///
|
||||||
|
/// **Local buckets:** path resolution mirrors lakehouse.toml's convention.
|
||||||
|
/// Returns an absolute filesystem path.
|
||||||
|
///
|
||||||
|
/// **S3 buckets:** returns `s3://{s3_bucket}/lance/{index_name}`. Lance's
|
||||||
|
/// internal object_store crate reads `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY`
|
||||||
|
/// / `AWS_ENDPOINT` from environment (or the S3 feature's default chain).
|
||||||
|
/// For MinIO: set `AWS_ENDPOINT=http://localhost:9000` and
|
||||||
|
/// `AWS_ALLOW_HTTP=true` before starting the gateway.
|
||||||
|
///
|
||||||
|
/// Refuses unknown buckets so a typo doesn't silently land Lance data
|
||||||
|
/// in a directory / prefix the rest of the system can't see.
|
||||||
|
pub fn lance_uri_for(
|
||||||
|
buckets: &BucketRegistry,
|
||||||
|
bucket: &str,
|
||||||
|
index_name: &str,
|
||||||
|
) -> Result<String, String> {
|
||||||
|
if !buckets.contains(bucket) {
|
||||||
|
return Err(format!("bucket '{bucket}' not registered"));
|
||||||
|
}
|
||||||
|
// Check if this bucket is S3-backed by looking for a bucket config
|
||||||
|
// with backend="s3". BucketRegistry exposes backend type through
|
||||||
|
// the list() info, but that's async. The simpler signal: if the
|
||||||
|
// bucket name matches one we know is S3 (configured via lakehouse.toml
|
||||||
|
// with backend="s3"), use the s3:// URI scheme.
|
||||||
|
//
|
||||||
|
// For the synchronous path, we check a naming convention: buckets
|
||||||
|
// whose name starts with "s3:" are treated as S3 targets. The rest
|
||||||
|
// of the name is the S3 bucket name. Convention-based, explicit,
|
||||||
|
// no async needed.
|
||||||
|
//
|
||||||
|
// Additionally, any bucket registered with backend="s3" in the
|
||||||
|
// config will have its BucketConfig.bucket field set — that's the
|
||||||
|
// actual S3 bucket name. We can't access BucketConfig synchronously
|
||||||
|
// from the current registry API, so for now the naming convention
|
||||||
|
// is the primary signal.
|
||||||
|
if bucket.starts_with("s3:") {
|
||||||
|
let s3_bucket = &bucket["s3:".len()..];
|
||||||
|
return Ok(format!("s3://{s3_bucket}/lance/{index_name}"));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Local path resolution.
|
||||||
|
let root: PathBuf = match bucket {
|
||||||
|
"primary" => PathBuf::from("./data"),
|
||||||
|
"rescue" => PathBuf::from("./data/_rescue"),
|
||||||
|
"testing" => PathBuf::from("./data/_testing"),
|
||||||
|
b if b.starts_with("profile:") => {
|
||||||
|
let safe = b.replace(':', "_");
|
||||||
|
PathBuf::from(format!("./data/_profiles/{safe}"))
|
||||||
|
}
|
||||||
|
b => PathBuf::from(format!("./data/_buckets/{}", b.replace(':', "_"))),
|
||||||
|
};
|
||||||
|
let _ = std::fs::create_dir_all(root.join("lance"));
|
||||||
|
let abs = match std::fs::canonicalize(&root) {
|
||||||
|
Ok(p) => p.join("lance").join(index_name),
|
||||||
|
Err(_) => root.join("lance").join(index_name),
|
||||||
|
};
|
||||||
|
Ok(abs.to_string_lossy().to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Lookup-by-index registry of `LanceVectorStore` handles. One handle
|
||||||
|
/// per index, lazy-created on first call. Cheap to keep alive — it's
|
||||||
|
/// just a path string + handle to Lance's metadata cache.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct LanceRegistry {
|
||||||
|
buckets: Arc<BucketRegistry>,
|
||||||
|
indexes: IndexRegistry,
|
||||||
|
stores: Arc<RwLock<HashMap<String, LanceVectorStore>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl LanceRegistry {
|
||||||
|
pub fn new(buckets: Arc<BucketRegistry>, indexes: IndexRegistry) -> Self {
|
||||||
|
Self {
|
||||||
|
buckets,
|
||||||
|
indexes,
|
||||||
|
stores: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get (or lazy-create) the LanceVectorStore for an index.
|
||||||
|
/// Resolves bucket → URI via the registry; backend stays whatever
|
||||||
|
/// IndexMeta says (no enforcement here — caller decides whether to
|
||||||
|
/// route to Lance based on profile/index backend choice).
|
||||||
|
pub async fn store_for(&self, index_name: &str) -> Result<LanceVectorStore, String> {
|
||||||
|
if let Some(s) = self.stores.read().await.get(index_name) {
|
||||||
|
return Ok(s.clone());
|
||||||
|
}
|
||||||
|
let bucket = self.indexes
|
||||||
|
.get(index_name).await
|
||||||
|
.map(|m| m.bucket)
|
||||||
|
.unwrap_or_else(|| "primary".to_string());
|
||||||
|
let uri = lance_uri_for(&self.buckets, &bucket, index_name)?;
|
||||||
|
let store = LanceVectorStore::new(uri);
|
||||||
|
self.stores.write().await.insert(index_name.to_string(), store.clone());
|
||||||
|
Ok(store)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// For freshly-created indexes that don't have IndexMeta yet — caller
|
||||||
|
/// supplies the bucket explicitly.
|
||||||
|
pub async fn store_for_new(&self, index_name: &str, bucket: &str) -> Result<LanceVectorStore, String> {
|
||||||
|
let uri = lance_uri_for(&self.buckets, bucket, index_name)?;
|
||||||
|
let store = LanceVectorStore::new(uri);
|
||||||
|
self.stores.write().await.insert(index_name.to_string(), store.clone());
|
||||||
|
Ok(store)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,17 @@
|
|||||||
|
pub mod agent;
|
||||||
|
pub mod autotune;
|
||||||
|
pub mod lance_backend;
|
||||||
|
pub mod chunker;
|
||||||
|
pub mod embedding_cache;
|
||||||
|
pub mod harness;
|
||||||
|
pub mod hnsw;
|
||||||
|
pub mod index_registry;
|
||||||
|
pub mod jobs;
|
||||||
|
pub mod promotion;
|
||||||
|
pub mod refresh;
|
||||||
|
pub mod store;
|
||||||
|
pub mod search;
|
||||||
|
pub mod rag;
|
||||||
|
pub mod supervisor;
|
||||||
|
pub mod service;
|
||||||
|
pub mod trial;
|
||||||
@@ -0,0 +1,220 @@
|
|||||||
|
//! Phase 16: Promoted HNSW configs — the "active generation" pointer.
|
||||||
|
//!
|
||||||
|
//! An index's HNSW config used at build time normally defaults to the
|
||||||
|
//! system-wide default (`HnswConfig::default()`). An operator or the
|
||||||
|
//! autotune agent can *promote* a specific trial's config — subsequent
|
||||||
|
//! HNSW builds against that index use the promoted config instead.
|
||||||
|
//!
|
||||||
|
//! Every promotion is history-tracked so `rollback` can revert. The
|
||||||
|
//! history file lives at `primary://_hnsw_promotions/{index_name}.json`
|
||||||
|
//! and is small (< few KB) so we rewrite it on every promotion rather
|
||||||
|
//! than append-log.
|
||||||
|
//!
|
||||||
|
//! Not included here:
|
||||||
|
//! - Atomic graph rebuild on promote — promotion only updates the sticky
|
||||||
|
//! default. Next activation (or search that triggers lazy build) picks
|
||||||
|
//! up the new config. That's "zero-downtime swap after build" which is
|
||||||
|
//! what ADR-019 actually claimed; an instant-swap requires a
|
||||||
|
//! pre-built graph pool which we don't have yet.
|
||||||
|
//! - Agent loop — lives in `vectord::autotune`.
|
||||||
|
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use storaged::ops;
|
||||||
|
use storaged::registry::BucketRegistry;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
use crate::index_registry::IndexRegistry;
|
||||||
|
use crate::trial::HnswConfig;
|
||||||
|
|
||||||
|
const PROMOTION_PREFIX: &str = "_hnsw_promotions";
|
||||||
|
|
||||||
|
/// One promotion record. The `trial_id` is the origin of the config —
|
||||||
|
/// lets operators trace back "why was this config picked?" to the exact
|
||||||
|
/// trial in the trial journal.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct PromotionEntry {
|
||||||
|
pub config: HnswConfig,
|
||||||
|
pub trial_id: String,
|
||||||
|
pub promoted_at: DateTime<Utc>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub promoted_by: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub note: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Serialized form of an index's promotion history.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||||
|
pub struct PromotionFile {
|
||||||
|
pub index_name: String,
|
||||||
|
pub current: Option<PromotionEntry>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub history: Vec<PromotionEntry>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct PromotionRegistry {
|
||||||
|
buckets: Arc<BucketRegistry>,
|
||||||
|
index_registry: IndexRegistry,
|
||||||
|
cache: Arc<RwLock<HashMap<String, PromotionFile>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PromotionRegistry {
|
||||||
|
pub fn new(buckets: Arc<BucketRegistry>, index_registry: IndexRegistry) -> Self {
|
||||||
|
Self {
|
||||||
|
buckets,
|
||||||
|
index_registry,
|
||||||
|
cache: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn key(index_name: &str) -> String {
|
||||||
|
// Sanitize for object-store safety.
|
||||||
|
let safe: String = index_name
|
||||||
|
.chars()
|
||||||
|
.map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' })
|
||||||
|
.collect();
|
||||||
|
format!("{PROMOTION_PREFIX}/{safe}.json")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolve which bucket's store holds this index's promotion file.
|
||||||
|
/// Same rules as TrialJournal::bucket_for — follows IndexMeta.bucket,
|
||||||
|
/// defaults to primary when metadata is missing.
|
||||||
|
async fn store_for(&self, index_name: &str) -> Result<Arc<dyn object_store::ObjectStore>, String> {
|
||||||
|
let bucket = self.index_registry
|
||||||
|
.get(index_name)
|
||||||
|
.await
|
||||||
|
.map(|m| m.bucket)
|
||||||
|
.unwrap_or_else(|| "primary".to_string());
|
||||||
|
self.buckets.get(&bucket)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Load (and cache) the promotion file for an index.
|
||||||
|
pub async fn load(&self, index_name: &str) -> Result<PromotionFile, String> {
|
||||||
|
if let Some(cached) = self.cache.read().await.get(index_name) {
|
||||||
|
return Ok(cached.clone());
|
||||||
|
}
|
||||||
|
let store = self.store_for(index_name).await?;
|
||||||
|
let key = Self::key(index_name);
|
||||||
|
let file = match ops::get(&store, &key).await {
|
||||||
|
Ok(bytes) => serde_json::from_slice::<PromotionFile>(&bytes)
|
||||||
|
.map_err(|e| format!("parse promotion file: {e}"))?,
|
||||||
|
Err(_) => PromotionFile {
|
||||||
|
index_name: index_name.to_string(),
|
||||||
|
current: None,
|
||||||
|
history: Vec::new(),
|
||||||
|
},
|
||||||
|
};
|
||||||
|
self.cache.write().await.insert(index_name.to_string(), file.clone());
|
||||||
|
Ok(file)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Promote a config to the active slot. Pushes the current promotion
|
||||||
|
/// (if any) onto the history stack. Persists before returning — the
|
||||||
|
/// config is durable by the time this call completes.
|
||||||
|
pub async fn promote(
|
||||||
|
&self,
|
||||||
|
index_name: &str,
|
||||||
|
entry: PromotionEntry,
|
||||||
|
) -> Result<PromotionFile, String> {
|
||||||
|
let mut file = self.load(index_name).await?;
|
||||||
|
if let Some(prior) = file.current.take() {
|
||||||
|
file.history.push(prior);
|
||||||
|
// Cap history to something sensible so this file doesn't grow
|
||||||
|
// unbounded. 50 entries = 50 promotions — way more than any
|
||||||
|
// sane workflow needs.
|
||||||
|
const HISTORY_CAP: usize = 50;
|
||||||
|
if file.history.len() > HISTORY_CAP {
|
||||||
|
let drop = file.history.len() - HISTORY_CAP;
|
||||||
|
file.history.drain(0..drop);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
file.current = Some(entry);
|
||||||
|
file.index_name = index_name.to_string();
|
||||||
|
|
||||||
|
let store = self.store_for(index_name).await?;
|
||||||
|
let key = Self::key(index_name);
|
||||||
|
let json = serde_json::to_vec_pretty(&file).map_err(|e| e.to_string())?;
|
||||||
|
ops::put(&store, &key, json.into()).await?;
|
||||||
|
|
||||||
|
self.cache.write().await.insert(index_name.to_string(), file.clone());
|
||||||
|
tracing::info!(
|
||||||
|
"promoted '{}' to config {:?} (trial={})",
|
||||||
|
index_name, file.current.as_ref().unwrap().config, file.current.as_ref().unwrap().trial_id,
|
||||||
|
);
|
||||||
|
Ok(file)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pop the latest promotion back onto the current slot (if any
|
||||||
|
/// history exists). If current is set but history is empty, the
|
||||||
|
/// current promotion is cleared — the index falls back to defaults.
|
||||||
|
pub async fn rollback(&self, index_name: &str) -> Result<PromotionFile, String> {
|
||||||
|
let mut file = self.load(index_name).await?;
|
||||||
|
match file.history.pop() {
|
||||||
|
Some(prev) => {
|
||||||
|
file.current = Some(prev);
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
if file.current.is_none() {
|
||||||
|
return Err(format!("no promotion to rollback for '{index_name}'"));
|
||||||
|
}
|
||||||
|
file.current = None;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let store = self.store_for(index_name).await?;
|
||||||
|
let key = Self::key(index_name);
|
||||||
|
let json = serde_json::to_vec_pretty(&file).map_err(|e| e.to_string())?;
|
||||||
|
ops::put(&store, &key, json.into()).await?;
|
||||||
|
self.cache.write().await.insert(index_name.to_string(), file.clone());
|
||||||
|
tracing::info!("rolled back promotion for '{}'", index_name);
|
||||||
|
Ok(file)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the currently-promoted config (if any). Callers use this to
|
||||||
|
/// pick the right HnswConfig at build time.
|
||||||
|
pub async fn get_current(&self, index_name: &str) -> Option<PromotionEntry> {
|
||||||
|
self.load(index_name).await.ok().and_then(|f| f.current)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Convenience: return the promoted config or the provided default.
|
||||||
|
pub async fn config_or(&self, index_name: &str, default: HnswConfig) -> HnswConfig {
|
||||||
|
match self.get_current(index_name).await {
|
||||||
|
Some(entry) => entry.config,
|
||||||
|
None => default,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List every index that has a promotion recorded (for operator UI).
|
||||||
|
///
|
||||||
|
/// Federation: scans EVERY registered bucket for promotion files.
|
||||||
|
/// Per-profile buckets each have their own `_hnsw_promotions/` so we
|
||||||
|
/// aggregate across them. Dedups by index_name — if the same index
|
||||||
|
/// somehow has promotion files in multiple buckets, the one from the
|
||||||
|
/// bucket recorded in IndexMeta wins.
|
||||||
|
pub async fn list_all(&self) -> Result<Vec<PromotionFile>, String> {
|
||||||
|
let bucket_infos = self.buckets.list().await;
|
||||||
|
let mut by_name: HashMap<String, PromotionFile> = HashMap::new();
|
||||||
|
|
||||||
|
for b in &bucket_infos {
|
||||||
|
let store = match self.buckets.get(&b.name) {
|
||||||
|
Ok(s) => s,
|
||||||
|
Err(_) => continue,
|
||||||
|
};
|
||||||
|
let keys = ops::list(&store, Some(&format!("{PROMOTION_PREFIX}/")))
|
||||||
|
.await.unwrap_or_default();
|
||||||
|
for key in keys {
|
||||||
|
if !key.ends_with(".json") { continue; }
|
||||||
|
let bytes = match ops::get(&store, &key).await {
|
||||||
|
Ok(b) => b,
|
||||||
|
Err(_) => continue,
|
||||||
|
};
|
||||||
|
if let Ok(f) = serde_json::from_slice::<PromotionFile>(&bytes) {
|
||||||
|
by_name.insert(f.index_name.clone(), f);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(by_name.into_values().collect())
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,84 @@
|
|||||||
|
/// RAG pipeline: question → embed → search → retrieve → generate answer.
|
||||||
|
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use aibridge::client::{AiClient, EmbedRequest, GenerateRequest};
|
||||||
|
use crate::search::{self, SearchResult};
|
||||||
|
use crate::store;
|
||||||
|
|
||||||
|
/// Full RAG answer with provenance.
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct RagResponse {
|
||||||
|
pub answer: String,
|
||||||
|
pub model: String,
|
||||||
|
pub sources: Vec<SearchResult>,
|
||||||
|
pub tokens_generated: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute full RAG: embed question → search index → retrieve context → generate answer.
|
||||||
|
pub async fn query(
|
||||||
|
question: &str,
|
||||||
|
index_name: &str,
|
||||||
|
top_k: usize,
|
||||||
|
object_store: &Arc<dyn ObjectStore>,
|
||||||
|
ai_client: &AiClient,
|
||||||
|
) -> Result<RagResponse, String> {
|
||||||
|
// 1. Embed the question
|
||||||
|
tracing::info!("RAG: embedding question");
|
||||||
|
let embed_resp = ai_client.embed(EmbedRequest {
|
||||||
|
texts: vec![question.to_string()],
|
||||||
|
model: None,
|
||||||
|
}).await?;
|
||||||
|
|
||||||
|
if embed_resp.embeddings.is_empty() {
|
||||||
|
return Err("no embedding returned for question".into());
|
||||||
|
}
|
||||||
|
|
||||||
|
let query_vec: Vec<f32> = embed_resp.embeddings[0].iter().map(|&x| x as f32).collect();
|
||||||
|
|
||||||
|
// 2. Load index and search
|
||||||
|
tracing::info!("RAG: searching index '{index_name}'");
|
||||||
|
let embeddings = store::load_embeddings(object_store, index_name).await?;
|
||||||
|
let results = search::search(&query_vec, &embeddings, top_k);
|
||||||
|
|
||||||
|
if results.is_empty() {
|
||||||
|
return Ok(RagResponse {
|
||||||
|
answer: "No relevant information found.".into(),
|
||||||
|
model: String::new(),
|
||||||
|
sources: vec![],
|
||||||
|
tokens_generated: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. Build context from retrieved chunks
|
||||||
|
let context: String = results.iter().enumerate().map(|(i, r)| {
|
||||||
|
format!("[{}] (source: {}, doc: {}) {}", i + 1, r.source, r.doc_id, r.chunk_text)
|
||||||
|
}).collect::<Vec<_>>().join("\n\n");
|
||||||
|
|
||||||
|
// 4. Generate answer
|
||||||
|
tracing::info!("RAG: generating answer from {} chunks", results.len());
|
||||||
|
let prompt = format!(
|
||||||
|
"You are a helpful assistant answering questions based on retrieved documents from a data system.\n\n\
|
||||||
|
Use ONLY the following context to answer. If the context doesn't contain enough information, say so.\n\
|
||||||
|
Cite sources by their number [1], [2], etc.\n\n\
|
||||||
|
Context:\n{context}\n\n\
|
||||||
|
Question: {question}\n\n\
|
||||||
|
Answer:"
|
||||||
|
);
|
||||||
|
|
||||||
|
let gen_resp = ai_client.generate(GenerateRequest {
|
||||||
|
prompt,
|
||||||
|
model: None,
|
||||||
|
system: None,
|
||||||
|
temperature: Some(0.2),
|
||||||
|
max_tokens: Some(512),
|
||||||
|
}).await?;
|
||||||
|
|
||||||
|
Ok(RagResponse {
|
||||||
|
answer: gen_resp.text.trim().to_string(),
|
||||||
|
model: gen_resp.model,
|
||||||
|
sources: results,
|
||||||
|
tokens_generated: gen_resp.tokens_generated,
|
||||||
|
})
|
||||||
|
}
|
||||||
@@ -0,0 +1,312 @@
|
|||||||
|
//! Phase C: Decoupled embedding refresh.
|
||||||
|
//!
|
||||||
|
//! When a dataset's row-level data changes, the vector index it feeds is
|
||||||
|
//! stale. Historically we coupled ingest and embedding — writing new rows
|
||||||
|
//! also re-embedded every row, which is fine at 500 rows but blows up at
|
||||||
|
//! 100K+. The llms3.com architecture calls out "asynchronous vector
|
||||||
|
//! refresh cycles independent from transactional mutations" as the right
|
||||||
|
//! pattern.
|
||||||
|
//!
|
||||||
|
//! This module implements the refresh side. The ingest path marks
|
||||||
|
//! embeddings stale (see catalogd::registry::mark_embeddings_stale); this
|
||||||
|
//! code clears that staleness by embedding only rows whose `doc_id` isn't
|
||||||
|
//! already in the existing vector index.
|
||||||
|
//!
|
||||||
|
//! Scope — MVP:
|
||||||
|
//! - Reads the dataset's Parquet, extracts (doc_id, text) pairs from named
|
||||||
|
//! columns
|
||||||
|
//! - Loads existing embeddings via EmbeddingCache
|
||||||
|
//! - Filters to rows whose doc_id is NOT in the existing set
|
||||||
|
//! - Chunks, embeds via Ollama, appends to the index parquet
|
||||||
|
//! - Clears stale flag on success
|
||||||
|
//!
|
||||||
|
//! Not in MVP:
|
||||||
|
//! - UPDATE semantics (same doc_id, new content) — would need content-hash
|
||||||
|
//! comparison per row
|
||||||
|
//! - Large-scale resilience (batching with checkpoints like the
|
||||||
|
//! supervisor) — MVP does it inline
|
||||||
|
//! - Lance backend — ADR-019 makes this straightforward later; MVP stays
|
||||||
|
//! on Parquet sidecar indexes
|
||||||
|
|
||||||
|
use std::collections::HashSet;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use aibridge::client::{AiClient, EmbedRequest};
|
||||||
|
use arrow::array::{Array, Int32Array, Int64Array, StringArray};
|
||||||
|
use catalogd::registry::Registry;
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
|
||||||
|
use crate::chunker::{self, TextChunk};
|
||||||
|
use crate::embedding_cache::EmbeddingCache;
|
||||||
|
use crate::index_registry::IndexRegistry;
|
||||||
|
use crate::store::{self, StoredEmbedding};
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, serde::Deserialize)]
|
||||||
|
pub struct RefreshRequest {
|
||||||
|
pub index_name: String,
|
||||||
|
/// Column with the document id (row identity).
|
||||||
|
pub id_column: String,
|
||||||
|
/// Column with the text to embed.
|
||||||
|
pub text_column: String,
|
||||||
|
/// Chunk size (chars). Defaults to 500.
|
||||||
|
#[serde(default)]
|
||||||
|
pub chunk_size: Option<usize>,
|
||||||
|
/// Overlap (chars). Defaults to 50.
|
||||||
|
#[serde(default)]
|
||||||
|
pub overlap: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct RefreshResult {
|
||||||
|
pub index_name: String,
|
||||||
|
pub dataset_name: String,
|
||||||
|
pub pre_existing_docs: usize,
|
||||||
|
pub dataset_docs: usize,
|
||||||
|
pub new_docs_embedded: usize,
|
||||||
|
pub new_chunks_embedded: usize,
|
||||||
|
pub total_embeddings_after: usize,
|
||||||
|
pub duration_secs: f32,
|
||||||
|
pub stale_cleared: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Full refresh pipeline. Takes dataset_name as URL param, body has the
|
||||||
|
/// column selectors.
|
||||||
|
pub async fn refresh_index(
|
||||||
|
dataset_name: &str,
|
||||||
|
req: &RefreshRequest,
|
||||||
|
store_: &Arc<dyn ObjectStore>,
|
||||||
|
registry: &Registry,
|
||||||
|
ai_client: &AiClient,
|
||||||
|
embedding_cache: &EmbeddingCache,
|
||||||
|
index_registry: &IndexRegistry,
|
||||||
|
) -> Result<RefreshResult, String> {
|
||||||
|
let t0 = std::time::Instant::now();
|
||||||
|
|
||||||
|
// 1. Find the dataset manifest, pull its object storage key
|
||||||
|
let manifest = registry
|
||||||
|
.get_by_name(dataset_name)
|
||||||
|
.await
|
||||||
|
.ok_or_else(|| format!("dataset not found: {dataset_name}"))?;
|
||||||
|
|
||||||
|
if manifest.objects.is_empty() {
|
||||||
|
return Err(format!("dataset '{dataset_name}' has no object references"));
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. Read dataset rows — extract (doc_id, text) pairs from the
|
||||||
|
// specified columns. Multi-object datasets: concat across all.
|
||||||
|
let mut doc_id_to_text: Vec<(String, String)> = Vec::new();
|
||||||
|
for obj in &manifest.objects {
|
||||||
|
let data = storaged::ops::get(store_, &obj.key).await
|
||||||
|
.map_err(|e| format!("read {}: {e}", obj.key))?;
|
||||||
|
let (schema, batches) = shared::arrow_helpers::parquet_to_record_batches(&data)
|
||||||
|
.map_err(|e| format!("parse {}: {e}", obj.key))?;
|
||||||
|
|
||||||
|
let id_idx = schema.index_of(&req.id_column)
|
||||||
|
.map_err(|_| format!("id column '{}' not in dataset schema", req.id_column))?;
|
||||||
|
let text_idx = schema.index_of(&req.text_column)
|
||||||
|
.map_err(|_| format!("text column '{}' not in dataset schema", req.text_column))?;
|
||||||
|
|
||||||
|
for batch in &batches {
|
||||||
|
let text_col = batch
|
||||||
|
.column(text_idx)
|
||||||
|
.as_any()
|
||||||
|
.downcast_ref::<StringArray>()
|
||||||
|
.ok_or_else(|| format!("text column '{}' is not Utf8", req.text_column))?;
|
||||||
|
|
||||||
|
// Accept Utf8, Int32, or Int64 as id — that covers CSV-derived
|
||||||
|
// and Postgres-imported schemas without forcing upstream casts.
|
||||||
|
let id_reader: Box<dyn Fn(usize) -> Option<String>> = {
|
||||||
|
let col = batch.column(id_idx);
|
||||||
|
if let Some(s) = col.as_any().downcast_ref::<StringArray>() {
|
||||||
|
let s = s.clone();
|
||||||
|
Box::new(move |row| if s.is_null(row) { None } else { Some(s.value(row).to_string()) })
|
||||||
|
} else if let Some(a) = col.as_any().downcast_ref::<Int32Array>() {
|
||||||
|
let a = a.clone();
|
||||||
|
Box::new(move |row| if a.is_null(row) { None } else { Some(a.value(row).to_string()) })
|
||||||
|
} else if let Some(a) = col.as_any().downcast_ref::<Int64Array>() {
|
||||||
|
let a = a.clone();
|
||||||
|
Box::new(move |row| if a.is_null(row) { None } else { Some(a.value(row).to_string()) })
|
||||||
|
} else {
|
||||||
|
return Err(format!(
|
||||||
|
"id column '{}' must be Utf8, Int32, or Int64 — got {}",
|
||||||
|
req.id_column,
|
||||||
|
col.data_type(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
for row in 0..batch.num_rows() {
|
||||||
|
if text_col.is_null(row) { continue; }
|
||||||
|
let Some(id) = id_reader(row) else { continue; };
|
||||||
|
let text = text_col.value(row).to_string();
|
||||||
|
if text.trim().is_empty() { continue; }
|
||||||
|
doc_id_to_text.push((id, text));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let dataset_docs = doc_id_to_text.len();
|
||||||
|
tracing::info!("refresh '{}': dataset has {dataset_docs} rows", dataset_name);
|
||||||
|
|
||||||
|
// 3. Load existing embeddings (empty if no index yet)
|
||||||
|
let existing: Vec<StoredEmbedding> = match embedding_cache.get_or_load(&req.index_name).await {
|
||||||
|
Ok(arc) => arc.as_ref().clone(),
|
||||||
|
Err(_) => Vec::new(), // first-time index build
|
||||||
|
};
|
||||||
|
let pre_existing_docs: HashSet<String> = existing
|
||||||
|
.iter()
|
||||||
|
.map(|e| e.doc_id.clone())
|
||||||
|
.collect();
|
||||||
|
let pre_existing_count = pre_existing_docs.len();
|
||||||
|
|
||||||
|
// 4. Delta — rows whose doc_id isn't already indexed
|
||||||
|
let new_rows: Vec<(String, String)> = doc_id_to_text
|
||||||
|
.into_iter()
|
||||||
|
.filter(|(id, _)| !pre_existing_docs.contains(id))
|
||||||
|
.collect();
|
||||||
|
let new_docs = new_rows.len();
|
||||||
|
|
||||||
|
if new_docs == 0 {
|
||||||
|
tracing::info!("refresh '{}': no new docs to embed", dataset_name);
|
||||||
|
// Even on no-op, make sure index metadata is registered so
|
||||||
|
// downstream discovery (profile activation, A/B search) sees
|
||||||
|
// indexes built by earlier refresh calls that predated the
|
||||||
|
// auto-register behavior.
|
||||||
|
let _ = try_update_index_meta(index_registry, &req.index_name, existing.len()).await;
|
||||||
|
registry.clear_embeddings_stale(dataset_name).await?;
|
||||||
|
return Ok(RefreshResult {
|
||||||
|
index_name: req.index_name.clone(),
|
||||||
|
dataset_name: dataset_name.to_string(),
|
||||||
|
pre_existing_docs: pre_existing_count,
|
||||||
|
dataset_docs,
|
||||||
|
new_docs_embedded: 0,
|
||||||
|
new_chunks_embedded: 0,
|
||||||
|
total_embeddings_after: existing.len(),
|
||||||
|
duration_secs: t0.elapsed().as_secs_f32(),
|
||||||
|
stale_cleared: true,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// 5. Chunk the new rows
|
||||||
|
let chunk_size = req.chunk_size.unwrap_or(500);
|
||||||
|
let overlap = req.overlap.unwrap_or(50);
|
||||||
|
let doc_ids: Vec<String> = new_rows.iter().map(|(id, _)| id.clone()).collect();
|
||||||
|
let texts: Vec<String> = new_rows.iter().map(|(_, t)| t.clone()).collect();
|
||||||
|
let chunks: Vec<TextChunk> = chunker::chunk_column(
|
||||||
|
dataset_name, &doc_ids, &texts, chunk_size, overlap,
|
||||||
|
);
|
||||||
|
let new_chunks = chunks.len();
|
||||||
|
tracing::info!("refresh '{}': {} new docs -> {} chunks", dataset_name, new_docs, new_chunks);
|
||||||
|
|
||||||
|
// 6. Embed via Ollama (batched)
|
||||||
|
let batch_size = 32;
|
||||||
|
let mut all_vectors: Vec<Vec<f64>> = Vec::with_capacity(new_chunks);
|
||||||
|
for batch in chunks.chunks(batch_size) {
|
||||||
|
let batch_texts: Vec<String> = batch.iter().map(|c| c.text.clone()).collect();
|
||||||
|
let resp = ai_client
|
||||||
|
.embed(EmbedRequest { texts: batch_texts, model: None })
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("embed: {e}"))?;
|
||||||
|
all_vectors.extend(resp.embeddings);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 7. Combine existing + new and write back as a single parquet
|
||||||
|
// (MVP — append-as-rewrite. ADR-019 points to Lance for true native
|
||||||
|
// append; for the Parquet sidecar we rewrite at refresh time.)
|
||||||
|
let mut new_stored: Vec<StoredEmbedding> = existing.clone();
|
||||||
|
for (chunk, vector) in chunks.into_iter().zip(all_vectors.iter()) {
|
||||||
|
new_stored.push(StoredEmbedding {
|
||||||
|
source: chunk.source,
|
||||||
|
doc_id: chunk.doc_id,
|
||||||
|
chunk_idx: chunk.chunk_idx,
|
||||||
|
chunk_text: chunk.text,
|
||||||
|
vector: vector.iter().map(|&x| x as f32).collect(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// We have to reconstruct chunks + vectors for store_embeddings — but
|
||||||
|
// store_embeddings takes &[TextChunk] and vectors separately. Convert
|
||||||
|
// the combined StoredEmbedding back to that shape.
|
||||||
|
let combined_chunks: Vec<TextChunk> = new_stored
|
||||||
|
.iter()
|
||||||
|
.map(|e| TextChunk {
|
||||||
|
source: e.source.clone(),
|
||||||
|
doc_id: e.doc_id.clone(),
|
||||||
|
chunk_idx: e.chunk_idx,
|
||||||
|
text: e.chunk_text.clone(),
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let combined_vectors: Vec<Vec<f64>> = new_stored
|
||||||
|
.iter()
|
||||||
|
.map(|e| e.vector.iter().map(|&f| f as f64).collect())
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let _key = store::store_embeddings(
|
||||||
|
store_, &req.index_name, &combined_chunks, &combined_vectors,
|
||||||
|
).await?;
|
||||||
|
|
||||||
|
// 8. Evict embedding cache — next read will pick up the new file
|
||||||
|
let _ = embedding_cache.evict(&req.index_name).await;
|
||||||
|
|
||||||
|
// 9. Update index registry metadata (row/chunk counts; others unchanged
|
||||||
|
// so we don't disturb the existing entry if present)
|
||||||
|
let _ = try_update_index_meta(index_registry, &req.index_name, new_stored.len()).await;
|
||||||
|
|
||||||
|
// 10. Clear stale flag
|
||||||
|
registry.clear_embeddings_stale(dataset_name).await?;
|
||||||
|
|
||||||
|
let total = new_stored.len();
|
||||||
|
Ok(RefreshResult {
|
||||||
|
index_name: req.index_name.clone(),
|
||||||
|
dataset_name: dataset_name.to_string(),
|
||||||
|
pre_existing_docs: pre_existing_count,
|
||||||
|
dataset_docs,
|
||||||
|
new_docs_embedded: new_docs,
|
||||||
|
new_chunks_embedded: new_chunks,
|
||||||
|
total_embeddings_after: total,
|
||||||
|
duration_secs: t0.elapsed().as_secs_f32(),
|
||||||
|
stale_cleared: true,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Best-effort refresh of index registry metadata.
|
||||||
|
/// - If the index already exists, bump the chunk_count.
|
||||||
|
/// - If it's brand new (first-time build via refresh), REGISTER it so
|
||||||
|
/// downstream discovery (profile activation, A/B search) can find it.
|
||||||
|
async fn try_update_index_meta(
|
||||||
|
index_registry: &IndexRegistry,
|
||||||
|
index_name: &str,
|
||||||
|
chunk_count: usize,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
if let Some(mut meta) = index_registry.get(index_name).await {
|
||||||
|
meta.chunk_count = chunk_count;
|
||||||
|
return index_registry.register(meta).await;
|
||||||
|
}
|
||||||
|
// First-time registration — infer reasonable defaults. Convention:
|
||||||
|
// index name `{source}_v1` or `{source}_vN` implies a source dataset
|
||||||
|
// named by stripping the `_vN` suffix. Otherwise use the index name
|
||||||
|
// as the source and let the caller patch later.
|
||||||
|
let source = match index_name.rsplit_once('_') {
|
||||||
|
Some((base, suffix)) if suffix.starts_with('v') && suffix[1..].chars().all(|c| c.is_ascii_digit()) => {
|
||||||
|
base.to_string()
|
||||||
|
}
|
||||||
|
_ => index_name.to_string(),
|
||||||
|
};
|
||||||
|
let meta = crate::index_registry::IndexMeta {
|
||||||
|
index_name: index_name.to_string(),
|
||||||
|
source,
|
||||||
|
model_name: "nomic-embed-text".to_string(),
|
||||||
|
model_version: "latest".to_string(),
|
||||||
|
dimensions: 768,
|
||||||
|
chunk_count,
|
||||||
|
doc_count: chunk_count,
|
||||||
|
chunk_size: 500,
|
||||||
|
overlap: 50,
|
||||||
|
storage_key: format!("vectors/{index_name}.parquet"),
|
||||||
|
created_at: chrono::Utc::now(),
|
||||||
|
build_time_secs: 0.0,
|
||||||
|
chunks_per_sec: 0.0,
|
||||||
|
bucket: "primary".to_string(),
|
||||||
|
vector_backend: shared::types::VectorBackend::Parquet,
|
||||||
|
};
|
||||||
|
index_registry.register(meta).await
|
||||||
|
}
|
||||||
@@ -0,0 +1,107 @@
|
|||||||
|
/// Brute-force vector search with cosine similarity.
|
||||||
|
/// Works well up to ~100K vectors. HNSW index would go here for larger scale.
|
||||||
|
|
||||||
|
use crate::store::StoredEmbedding;
|
||||||
|
|
||||||
|
/// A search result with score.
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct SearchResult {
|
||||||
|
pub source: String,
|
||||||
|
pub doc_id: String,
|
||||||
|
pub chunk_idx: u32,
|
||||||
|
pub chunk_text: String,
|
||||||
|
pub score: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Search embeddings by cosine similarity. Returns top_k results.
|
||||||
|
pub fn search(
|
||||||
|
query_vector: &[f32],
|
||||||
|
embeddings: &[StoredEmbedding],
|
||||||
|
top_k: usize,
|
||||||
|
) -> Vec<SearchResult> {
|
||||||
|
let query_norm = norm(query_vector);
|
||||||
|
if query_norm == 0.0 {
|
||||||
|
return vec![];
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut scored: Vec<SearchResult> = embeddings.iter().map(|emb| {
|
||||||
|
let score = cosine_similarity(query_vector, &emb.vector, query_norm);
|
||||||
|
SearchResult {
|
||||||
|
source: emb.source.clone(),
|
||||||
|
doc_id: emb.doc_id.clone(),
|
||||||
|
chunk_idx: emb.chunk_idx,
|
||||||
|
chunk_text: emb.chunk_text.clone(),
|
||||||
|
score,
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
|
||||||
|
// Sort descending by score
|
||||||
|
scored.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
scored.truncate(top_k);
|
||||||
|
scored
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cosine_similarity(a: &[f32], b: &[f32], a_norm: f32) -> f32 {
|
||||||
|
if a.len() != b.len() {
|
||||||
|
return 0.0;
|
||||||
|
}
|
||||||
|
let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
|
||||||
|
let b_norm = norm(b);
|
||||||
|
if b_norm == 0.0 {
|
||||||
|
return 0.0;
|
||||||
|
}
|
||||||
|
dot / (a_norm * b_norm)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn norm(v: &[f32]) -> f32 {
|
||||||
|
v.iter().map(|x| x * x).sum::<f32>().sqrt()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn identical_vectors_score_1() {
|
||||||
|
let v = vec![1.0, 2.0, 3.0];
|
||||||
|
let emb = StoredEmbedding {
|
||||||
|
source: "test".into(),
|
||||||
|
doc_id: "1".into(),
|
||||||
|
chunk_idx: 0,
|
||||||
|
chunk_text: "hello".into(),
|
||||||
|
vector: v.clone(),
|
||||||
|
};
|
||||||
|
let results = search(&v, &[emb], 1);
|
||||||
|
assert!((results[0].score - 1.0).abs() < 0.001);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn orthogonal_vectors_score_0() {
|
||||||
|
let q = vec![1.0, 0.0];
|
||||||
|
let emb = StoredEmbedding {
|
||||||
|
source: "test".into(),
|
||||||
|
doc_id: "1".into(),
|
||||||
|
chunk_idx: 0,
|
||||||
|
chunk_text: "hello".into(),
|
||||||
|
vector: vec![0.0, 1.0],
|
||||||
|
};
|
||||||
|
let results = search(&q, &[emb], 1);
|
||||||
|
assert!(results[0].score.abs() < 0.001);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn returns_top_k() {
|
||||||
|
let q = vec![1.0, 0.0, 0.0];
|
||||||
|
let embs: Vec<StoredEmbedding> = (0..10).map(|i| StoredEmbedding {
|
||||||
|
source: "test".into(),
|
||||||
|
doc_id: format!("{i}"),
|
||||||
|
chunk_idx: 0,
|
||||||
|
chunk_text: format!("doc {i}"),
|
||||||
|
vector: vec![1.0 - i as f32 * 0.1, i as f32 * 0.1, 0.0],
|
||||||
|
}).collect();
|
||||||
|
let results = search(&q, &embs, 3);
|
||||||
|
assert_eq!(results.len(), 3);
|
||||||
|
assert!(results[0].score >= results[1].score);
|
||||||
|
assert!(results[1].score >= results[2].score);
|
||||||
|
}
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,116 @@
|
|||||||
|
/// Vector storage as Parquet files.
|
||||||
|
/// Each embedding index is stored as: source, doc_id, chunk_idx, chunk_text, vector (binary blob).
|
||||||
|
/// Vectors are stored as raw f32 bytes for compact storage and fast loading.
|
||||||
|
|
||||||
|
use arrow::array::{ArrayRef, BinaryArray, Float32Array, Int32Array, RecordBatch, StringArray};
|
||||||
|
use arrow::datatypes::{DataType, Field, Schema};
|
||||||
|
use bytes::Bytes;
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use storaged::ops;
|
||||||
|
|
||||||
|
use crate::chunker::TextChunk;
|
||||||
|
|
||||||
|
/// A stored embedding — chunk text + its vector.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct StoredEmbedding {
|
||||||
|
pub source: String,
|
||||||
|
pub doc_id: String,
|
||||||
|
pub chunk_idx: u32,
|
||||||
|
pub chunk_text: String,
|
||||||
|
pub vector: Vec<f32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Store embeddings as a Parquet file in object storage.
|
||||||
|
pub async fn store_embeddings(
|
||||||
|
store: &Arc<dyn ObjectStore>,
|
||||||
|
index_name: &str,
|
||||||
|
chunks: &[TextChunk],
|
||||||
|
vectors: &[Vec<f64>], // from embedding API (f64), we store as f32
|
||||||
|
) -> Result<String, String> {
|
||||||
|
if chunks.len() != vectors.len() {
|
||||||
|
return Err(format!("chunk count ({}) != vector count ({})", chunks.len(), vectors.len()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let n = chunks.len();
|
||||||
|
let sources: Vec<&str> = chunks.iter().map(|c| c.source.as_str()).collect();
|
||||||
|
let doc_ids: Vec<&str> = chunks.iter().map(|c| c.doc_id.as_str()).collect();
|
||||||
|
let chunk_idxs: Vec<i32> = chunks.iter().map(|c| c.chunk_idx as i32).collect();
|
||||||
|
let texts: Vec<&str> = chunks.iter().map(|c| c.text.as_str()).collect();
|
||||||
|
|
||||||
|
// Store vectors as raw f32 bytes (compact binary blob)
|
||||||
|
let vector_bytes: Vec<Vec<u8>> = vectors.iter().map(|v| {
|
||||||
|
v.iter().map(|&x| x as f32).flat_map(|f| f.to_le_bytes()).collect()
|
||||||
|
}).collect();
|
||||||
|
let vector_refs: Vec<&[u8]> = vector_bytes.iter().map(|v| v.as_slice()).collect();
|
||||||
|
|
||||||
|
let schema = Arc::new(Schema::new(vec![
|
||||||
|
Field::new("source", DataType::Utf8, false),
|
||||||
|
Field::new("doc_id", DataType::Utf8, false),
|
||||||
|
Field::new("chunk_idx", DataType::Int32, false),
|
||||||
|
Field::new("chunk_text", DataType::Utf8, false),
|
||||||
|
Field::new("vector", DataType::Binary, false),
|
||||||
|
]));
|
||||||
|
|
||||||
|
let arrays: Vec<ArrayRef> = vec![
|
||||||
|
Arc::new(StringArray::from(sources)),
|
||||||
|
Arc::new(StringArray::from(doc_ids)),
|
||||||
|
Arc::new(Int32Array::from(chunk_idxs)),
|
||||||
|
Arc::new(StringArray::from(texts)),
|
||||||
|
Arc::new(BinaryArray::from(vector_refs)),
|
||||||
|
];
|
||||||
|
|
||||||
|
let batch = RecordBatch::try_new(schema, arrays)
|
||||||
|
.map_err(|e| format!("RecordBatch error: {e}"))?;
|
||||||
|
|
||||||
|
let parquet = shared::arrow_helpers::record_batch_to_parquet(&batch)?;
|
||||||
|
let key = format!("vectors/{index_name}.parquet");
|
||||||
|
ops::put(store, &key, parquet).await?;
|
||||||
|
|
||||||
|
tracing::info!("stored {n} embeddings in {key}");
|
||||||
|
Ok(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Load all embeddings from a vector index file.
|
||||||
|
pub async fn load_embeddings(
|
||||||
|
store: &Arc<dyn ObjectStore>,
|
||||||
|
index_name: &str,
|
||||||
|
) -> Result<Vec<StoredEmbedding>, String> {
|
||||||
|
let key = format!("vectors/{index_name}.parquet");
|
||||||
|
let data = ops::get(store, &key).await?;
|
||||||
|
|
||||||
|
let (_, batches) = shared::arrow_helpers::parquet_to_record_batches(&data)?;
|
||||||
|
|
||||||
|
let mut embeddings = Vec::new();
|
||||||
|
for batch in &batches {
|
||||||
|
let sources = batch.column(0).as_any().downcast_ref::<StringArray>()
|
||||||
|
.ok_or("source column not string")?;
|
||||||
|
let doc_ids = batch.column(1).as_any().downcast_ref::<StringArray>()
|
||||||
|
.ok_or("doc_id column not string")?;
|
||||||
|
let chunk_idxs = batch.column(2).as_any().downcast_ref::<Int32Array>()
|
||||||
|
.ok_or("chunk_idx column not int")?;
|
||||||
|
let texts = batch.column(3).as_any().downcast_ref::<StringArray>()
|
||||||
|
.ok_or("chunk_text column not string")?;
|
||||||
|
let vectors = batch.column(4).as_any().downcast_ref::<BinaryArray>()
|
||||||
|
.ok_or("vector column not binary")?;
|
||||||
|
|
||||||
|
for i in 0..batch.num_rows() {
|
||||||
|
let vec_bytes = vectors.value(i);
|
||||||
|
let vector: Vec<f32> = vec_bytes.chunks_exact(4)
|
||||||
|
.map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
embeddings.push(StoredEmbedding {
|
||||||
|
source: sources.value(i).to_string(),
|
||||||
|
doc_id: doc_ids.value(i).to_string(),
|
||||||
|
chunk_idx: chunk_idxs.value(i) as u32,
|
||||||
|
chunk_text: texts.value(i).to_string(),
|
||||||
|
vector,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!("loaded {} embeddings from {key}", embeddings.len());
|
||||||
|
Ok(embeddings)
|
||||||
|
}
|
||||||
@@ -0,0 +1,286 @@
|
|||||||
|
/// Dual-pipeline supervisor for embedding ingestion.
|
||||||
|
/// Splits work into ranges, runs parallel pipelines, handles failures
|
||||||
|
/// with round-robin retry and checkpointing.
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
use aibridge::client::{AiClient, EmbedRequest};
|
||||||
|
use crate::chunker::TextChunk;
|
||||||
|
use crate::jobs::JobTracker;
|
||||||
|
use crate::store;
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
|
||||||
|
/// A range of chunks to embed.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
struct ChunkRange {
|
||||||
|
start: usize,
|
||||||
|
end: usize,
|
||||||
|
attempts: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Checkpoint: tracks which ranges are done, persisted for crash recovery.
|
||||||
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct Checkpoint {
|
||||||
|
pub job_id: String,
|
||||||
|
pub index_name: String,
|
||||||
|
pub total_chunks: usize,
|
||||||
|
pub completed_ranges: Vec<(usize, usize)>, // (start, end) pairs
|
||||||
|
pub failed_ranges: Vec<(usize, usize, String)>, // (start, end, error)
|
||||||
|
pub embedded_count: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pipeline status for monitoring.
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct PipelineStatus {
|
||||||
|
pub pipeline_id: String,
|
||||||
|
pub current_range: Option<(usize, usize)>,
|
||||||
|
pub chunks_done: usize,
|
||||||
|
pub rate: f32,
|
||||||
|
pub alive: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Supervisor configuration.
|
||||||
|
pub struct SupervisorConfig {
|
||||||
|
pub num_pipelines: usize, // default 2
|
||||||
|
pub batch_size: usize, // embeddings per API call (default 32)
|
||||||
|
pub range_size: usize, // chunks per range (default 2000)
|
||||||
|
pub max_retries: u32, // per range (default 3)
|
||||||
|
pub checkpoint_interval: usize, // checkpoint every N chunks (default 1000)
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for SupervisorConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
num_pipelines: 4, // i9 + A4000 can handle parallel embedding
|
||||||
|
batch_size: 64, // larger batches for GPU throughput
|
||||||
|
range_size: 2500, // bigger ranges, fewer coordination overhead
|
||||||
|
max_retries: 3,
|
||||||
|
checkpoint_interval: 1000,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run supervised dual-pipeline embedding.
|
||||||
|
pub async fn run_supervised(
|
||||||
|
job_id: &str,
|
||||||
|
index_name: &str,
|
||||||
|
chunks: Vec<TextChunk>,
|
||||||
|
ai_client: &AiClient,
|
||||||
|
obj_store: &Arc<dyn ObjectStore>,
|
||||||
|
tracker: &JobTracker,
|
||||||
|
config: SupervisorConfig,
|
||||||
|
) -> Result<String, String> {
|
||||||
|
let total = chunks.len();
|
||||||
|
let chunks = Arc::new(chunks);
|
||||||
|
|
||||||
|
tracing::info!("supervisor: starting {} pipelines for {} chunks (range_size={})",
|
||||||
|
config.num_pipelines, total, config.range_size);
|
||||||
|
|
||||||
|
// Split into ranges
|
||||||
|
let mut ranges: Vec<ChunkRange> = Vec::new();
|
||||||
|
let mut start = 0;
|
||||||
|
while start < total {
|
||||||
|
let end = (start + config.range_size).min(total);
|
||||||
|
ranges.push(ChunkRange { start, end, attempts: 0 });
|
||||||
|
start = end;
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!("supervisor: {} ranges created", ranges.len());
|
||||||
|
|
||||||
|
// Shared state
|
||||||
|
let work_queue: Arc<RwLock<Vec<ChunkRange>>> = Arc::new(RwLock::new(ranges));
|
||||||
|
let all_vectors: Arc<RwLock<HashMap<usize, Vec<Vec<f64>>>>> = Arc::new(RwLock::new(HashMap::new()));
|
||||||
|
let dead_letter: Arc<RwLock<Vec<(usize, usize, String)>>> = Arc::new(RwLock::new(Vec::new()));
|
||||||
|
let global_count: Arc<std::sync::atomic::AtomicUsize> = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||||
|
let start_time = std::time::Instant::now();
|
||||||
|
|
||||||
|
// Checkpoint state
|
||||||
|
let checkpoint: Arc<RwLock<Checkpoint>> = Arc::new(RwLock::new(Checkpoint {
|
||||||
|
job_id: job_id.to_string(),
|
||||||
|
index_name: index_name.to_string(),
|
||||||
|
total_chunks: total,
|
||||||
|
completed_ranges: vec![],
|
||||||
|
failed_ranges: vec![],
|
||||||
|
embedded_count: 0,
|
||||||
|
}));
|
||||||
|
|
||||||
|
// Try to load existing checkpoint (crash recovery)
|
||||||
|
if let Ok(data) = storaged::ops::get(obj_store, &format!("checkpoints/{job_id}.json")).await {
|
||||||
|
if let Ok(saved) = serde_json::from_slice::<Checkpoint>(&data) {
|
||||||
|
tracing::info!("supervisor: recovered checkpoint — {}/{} already done",
|
||||||
|
saved.embedded_count, saved.total_chunks);
|
||||||
|
|
||||||
|
// Remove already-completed ranges from work queue
|
||||||
|
let mut queue = work_queue.write().await;
|
||||||
|
queue.retain(|r| !saved.completed_ranges.iter().any(|(s, e)| *s == r.start && *e == r.end));
|
||||||
|
global_count.store(saved.embedded_count, std::sync::atomic::Ordering::Relaxed);
|
||||||
|
|
||||||
|
*checkpoint.write().await = saved;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Spawn pipeline workers
|
||||||
|
let mut handles = Vec::new();
|
||||||
|
for pipeline_id in 0..config.num_pipelines {
|
||||||
|
let queue = work_queue.clone();
|
||||||
|
let vectors = all_vectors.clone();
|
||||||
|
let dead = dead_letter.clone();
|
||||||
|
let count = global_count.clone();
|
||||||
|
let ckpt = checkpoint.clone();
|
||||||
|
let client = ai_client.clone();
|
||||||
|
let store = obj_store.clone();
|
||||||
|
let chunks = chunks.clone();
|
||||||
|
let tracker = tracker.clone();
|
||||||
|
let jid = job_id.to_string();
|
||||||
|
let batch_size = config.batch_size;
|
||||||
|
let max_retries = config.max_retries;
|
||||||
|
let ckpt_interval = config.checkpoint_interval;
|
||||||
|
|
||||||
|
let handle = tokio::spawn(async move {
|
||||||
|
let pid = format!("pipeline-{pipeline_id}");
|
||||||
|
tracing::info!("{pid}: started");
|
||||||
|
|
||||||
|
loop {
|
||||||
|
// Grab next range from queue
|
||||||
|
let range = {
|
||||||
|
let mut q = queue.write().await;
|
||||||
|
q.pop()
|
||||||
|
};
|
||||||
|
|
||||||
|
let range = match range {
|
||||||
|
Some(r) => r,
|
||||||
|
None => {
|
||||||
|
tracing::info!("{pid}: no more work, shutting down");
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
tracing::debug!("{pid}: processing range [{}, {})", range.start, range.end);
|
||||||
|
|
||||||
|
// Embed the range
|
||||||
|
match embed_range(&chunks[range.start..range.end], &client, batch_size).await {
|
||||||
|
Ok(range_vectors) => {
|
||||||
|
let n = range_vectors.len();
|
||||||
|
vectors.write().await.insert(range.start, range_vectors);
|
||||||
|
let prev = count.fetch_add(n, std::sync::atomic::Ordering::Relaxed);
|
||||||
|
let done = prev + n;
|
||||||
|
|
||||||
|
// Update job tracker
|
||||||
|
let elapsed = start_time.elapsed().as_secs_f32();
|
||||||
|
let rate = if elapsed > 0.0 { done as f32 / elapsed } else { 0.0 };
|
||||||
|
tracker.update_progress(&jid, done, rate).await;
|
||||||
|
|
||||||
|
// Update checkpoint
|
||||||
|
let mut ckpt = ckpt.write().await;
|
||||||
|
ckpt.completed_ranges.push((range.start, range.end));
|
||||||
|
ckpt.embedded_count = done;
|
||||||
|
|
||||||
|
// Persist checkpoint periodically
|
||||||
|
if done % ckpt_interval < n {
|
||||||
|
let json = serde_json::to_vec(&*ckpt).unwrap_or_default();
|
||||||
|
let _ = storaged::ops::put(&store, &format!("checkpoints/{jid}.json"), json.into()).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!("{pid}: range [{},{}) done — {done}/{} total",
|
||||||
|
range.start, range.end, chunks.len());
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
let attempt = range.attempts + 1;
|
||||||
|
tracing::warn!("{pid}: range [{},{}) failed (attempt {attempt}/{max_retries}): {e}",
|
||||||
|
range.start, range.end);
|
||||||
|
|
||||||
|
if attempt < max_retries {
|
||||||
|
// Push back to queue for retry (round-robin to other pipeline)
|
||||||
|
queue.write().await.push(ChunkRange {
|
||||||
|
start: range.start,
|
||||||
|
end: range.end,
|
||||||
|
attempts: attempt,
|
||||||
|
});
|
||||||
|
} else {
|
||||||
|
// Dead letter
|
||||||
|
tracing::error!("{pid}: range [{},{}) dead-lettered after {max_retries} attempts",
|
||||||
|
range.start, range.end);
|
||||||
|
dead.write().await.push((range.start, range.end, e));
|
||||||
|
|
||||||
|
let mut ckpt = ckpt.write().await;
|
||||||
|
ckpt.failed_ranges.push((range.start, range.end, format!("max retries exceeded")));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
handles.push(handle);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wait for all pipelines
|
||||||
|
for handle in handles {
|
||||||
|
let _ = handle.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check for dead letters
|
||||||
|
let dead = dead_letter.read().await;
|
||||||
|
if !dead.is_empty() {
|
||||||
|
tracing::warn!("supervisor: {} ranges failed permanently", dead.len());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Assemble vectors in order
|
||||||
|
let vectors_map = all_vectors.read().await;
|
||||||
|
let mut sorted_starts: Vec<usize> = vectors_map.keys().cloned().collect();
|
||||||
|
sorted_starts.sort();
|
||||||
|
|
||||||
|
let mut final_vectors: Vec<Vec<f64>> = Vec::with_capacity(total);
|
||||||
|
for start in sorted_starts {
|
||||||
|
final_vectors.extend(vectors_map[&start].clone());
|
||||||
|
}
|
||||||
|
|
||||||
|
let embedded_count = final_vectors.len();
|
||||||
|
tracing::info!("supervisor: {embedded_count}/{total} chunks embedded, storing index");
|
||||||
|
|
||||||
|
// Store — only the successfully embedded chunks
|
||||||
|
// We need to match chunks to vectors
|
||||||
|
let successful_chunks: Vec<TextChunk> = {
|
||||||
|
let ckpt = checkpoint.read().await;
|
||||||
|
let mut result = Vec::new();
|
||||||
|
for (s, e) in &ckpt.completed_ranges {
|
||||||
|
result.extend(chunks[*s..*e].iter().cloned());
|
||||||
|
}
|
||||||
|
result
|
||||||
|
};
|
||||||
|
|
||||||
|
if successful_chunks.is_empty() {
|
||||||
|
return Err("no chunks were successfully embedded".into());
|
||||||
|
}
|
||||||
|
|
||||||
|
let key = store::store_embeddings(obj_store, index_name, &successful_chunks, &final_vectors).await?;
|
||||||
|
|
||||||
|
// Clean up checkpoint
|
||||||
|
let _ = storaged::ops::delete(obj_store, &format!("checkpoints/{job_id}.json")).await;
|
||||||
|
|
||||||
|
let elapsed = start_time.elapsed().as_secs_f32();
|
||||||
|
let rate = embedded_count as f32 / elapsed;
|
||||||
|
tracing::info!("supervisor: completed — {embedded_count} vectors in {elapsed:.0}s ({rate:.0}/sec)");
|
||||||
|
|
||||||
|
Ok(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Embed a range of chunks. Returns vectors or error.
|
||||||
|
async fn embed_range(
|
||||||
|
chunks: &[TextChunk],
|
||||||
|
ai_client: &AiClient,
|
||||||
|
batch_size: usize,
|
||||||
|
) -> Result<Vec<Vec<f64>>, String> {
|
||||||
|
let mut vectors = Vec::with_capacity(chunks.len());
|
||||||
|
|
||||||
|
for batch in chunks.chunks(batch_size) {
|
||||||
|
let texts: Vec<String> = batch.iter().map(|c| c.text.clone()).collect();
|
||||||
|
let resp = ai_client.embed(EmbedRequest {
|
||||||
|
texts,
|
||||||
|
model: None,
|
||||||
|
}).await.map_err(|e| format!("embed error: {e}"))?;
|
||||||
|
vectors.extend(resp.embeddings);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(vectors)
|
||||||
|
}
|
||||||
@@ -0,0 +1,240 @@
|
|||||||
|
/// Trial journal for HNSW parameter tuning.
|
||||||
|
///
|
||||||
|
/// Every HNSW build+eval is recorded as a Trial. The journal is append-only
|
||||||
|
/// and stored under `_hnsw_trials/{index_name}/` as batched JSONL files —
|
||||||
|
/// an AI agent iterating on configs reads prior trials to decide what to
|
||||||
|
/// try next, and writes a new trial on each attempt.
|
||||||
|
///
|
||||||
|
/// Storage uses the shared `storaged::append_log::AppendLog` so appends are
|
||||||
|
/// write-once (new file per batch) rather than rewriting a single growing
|
||||||
|
/// JSONL on every event. See `append_log.rs` for the full rationale.
|
||||||
|
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use object_store::ObjectStore;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
use storaged::append_log::AppendLog;
|
||||||
|
use storaged::registry::BucketRegistry;
|
||||||
|
|
||||||
|
use crate::index_registry::IndexRegistry;
|
||||||
|
|
||||||
|
/// HNSW build/search parameters the agent can tune.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct HnswConfig {
|
||||||
|
pub ef_construction: usize,
|
||||||
|
pub ef_search: usize,
|
||||||
|
#[serde(default)]
|
||||||
|
pub seed: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for HnswConfig {
|
||||||
|
/// Production default, locked in 2026-04-16 based on trial grid against
|
||||||
|
/// resumes_100k_v2 (100K vectors, 20 queries, recall@10):
|
||||||
|
/// ec=20 es=30 → recall 0.960, p50 509us, build 8s
|
||||||
|
/// ec=80 es=30 → recall 1.000, p50 873us, build 230s ← sweet spot
|
||||||
|
/// ec=200 es=30 → recall 1.000, p50 874us, build 106s (no recall gain)
|
||||||
|
///
|
||||||
|
/// `ec=80` is the smallest value that reaches 100% recall. Higher values
|
||||||
|
/// waste build time. `es=30` gives faster search than `es=100` with no
|
||||||
|
/// recall penalty at this scale.
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
ef_construction: 80,
|
||||||
|
ef_search: 30,
|
||||||
|
seed: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Metrics collected on every trial. All latencies in microseconds.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct TrialMetrics {
|
||||||
|
pub build_time_secs: f32,
|
||||||
|
pub search_latency_p50_us: f32,
|
||||||
|
pub search_latency_p95_us: f32,
|
||||||
|
pub search_latency_p99_us: f32,
|
||||||
|
pub recall_at_k: f32,
|
||||||
|
pub memory_bytes: u64,
|
||||||
|
pub vectors: usize,
|
||||||
|
pub eval_queries: usize,
|
||||||
|
/// Brute-force latency for comparison — how much speedup did HNSW buy us?
|
||||||
|
pub brute_force_latency_us: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A single tuning attempt.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct Trial {
|
||||||
|
pub id: String,
|
||||||
|
pub index_name: String,
|
||||||
|
pub eval_set: String,
|
||||||
|
pub config: HnswConfig,
|
||||||
|
pub metrics: TrialMetrics,
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
/// Free-form note — the agent can record why it tried this config.
|
||||||
|
#[serde(default)]
|
||||||
|
pub note: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Trial {
|
||||||
|
pub fn new_id() -> String {
|
||||||
|
format!(
|
||||||
|
"trial-{}-{}",
|
||||||
|
Utc::now().timestamp_millis(),
|
||||||
|
&uuid::Uuid::new_v4().to_string()[..8]
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Per-index append log, lazy-created on first write.
|
||||||
|
///
|
||||||
|
/// Federation layer 2: the journal resolves each index's bucket from the
|
||||||
|
/// index registry and writes its JSONL batches to THAT bucket, not
|
||||||
|
/// primary. Back-compat is preserved by `IndexMeta::bucket` defaulting
|
||||||
|
/// to "primary" for pre-federation indexes. Indexes the registry has
|
||||||
|
/// never heard of (edge case — trials run before first register) fall
|
||||||
|
/// through to primary as well.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct TrialJournal {
|
||||||
|
buckets: Arc<BucketRegistry>,
|
||||||
|
index_registry: IndexRegistry,
|
||||||
|
/// Cache per (bucket, index) AppendLog so the in-memory buffer persists
|
||||||
|
/// across calls. Keyed by `(bucket, index_name)` so moving an index
|
||||||
|
/// between buckets is clean — the old journal stays intact.
|
||||||
|
logs: Arc<RwLock<HashMap<(String, String), Arc<AppendLog>>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TrialJournal {
|
||||||
|
pub fn new(buckets: Arc<BucketRegistry>, index_registry: IndexRegistry) -> Self {
|
||||||
|
Self {
|
||||||
|
buckets,
|
||||||
|
index_registry,
|
||||||
|
logs: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn prefix(index_name: &str) -> String {
|
||||||
|
format!("_hnsw_trials/{}", index_name)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolve which bucket holds this index's trial artifacts.
|
||||||
|
/// Falls back to primary for indexes without recorded metadata.
|
||||||
|
async fn bucket_for(&self, index_name: &str) -> String {
|
||||||
|
self.index_registry
|
||||||
|
.get(index_name)
|
||||||
|
.await
|
||||||
|
.map(|m| m.bucket)
|
||||||
|
.unwrap_or_else(|| "primary".to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn log_for(&self, index_name: &str) -> Result<Arc<AppendLog>, String> {
|
||||||
|
let bucket = self.bucket_for(index_name).await;
|
||||||
|
let key = (bucket.clone(), index_name.to_string());
|
||||||
|
|
||||||
|
if let Some(log) = self.logs.read().await.get(&key) {
|
||||||
|
return Ok(log.clone());
|
||||||
|
}
|
||||||
|
let mut guard = self.logs.write().await;
|
||||||
|
if let Some(log) = guard.get(&key) {
|
||||||
|
return Ok(log.clone());
|
||||||
|
}
|
||||||
|
let store = self.buckets.get(&bucket)?;
|
||||||
|
// Trials arrive one at a time during human/agent iteration — a low
|
||||||
|
// threshold gives "hit /trials and see my latest attempt" immediacy
|
||||||
|
// without creating one file per event.
|
||||||
|
let log = Arc::new(
|
||||||
|
AppendLog::new(store, Self::prefix(index_name))
|
||||||
|
.with_flush_threshold(4),
|
||||||
|
);
|
||||||
|
guard.insert(key, log.clone());
|
||||||
|
Ok(log)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append a trial record. In-memory buffered; persisted in batches.
|
||||||
|
pub async fn append(&self, trial: &Trial) -> Result<(), String> {
|
||||||
|
let line = serde_json::to_vec(trial).map_err(|e| e.to_string())?;
|
||||||
|
let log = self.log_for(&trial.index_name).await?;
|
||||||
|
log.append(line).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Read all trials for an index (flushed batches + unflushed buffer).
|
||||||
|
pub async fn list(&self, index_name: &str) -> Result<Vec<Trial>, String> {
|
||||||
|
let log = self.log_for(index_name).await?;
|
||||||
|
let lines = log.read_all().await?;
|
||||||
|
let mut trials = Vec::with_capacity(lines.len());
|
||||||
|
for line in lines {
|
||||||
|
match serde_json::from_slice::<Trial>(&line) {
|
||||||
|
Ok(t) => trials.push(t),
|
||||||
|
Err(e) => tracing::warn!("trial journal: skip malformed line: {e}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(trials)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Explicit flush for callers that want write-through semantics
|
||||||
|
/// (e.g. an agent that wants to commit a trial before querying stats).
|
||||||
|
pub async fn flush(&self, index_name: &str) -> Result<(), String> {
|
||||||
|
let log = self.log_for(index_name).await?;
|
||||||
|
log.flush().await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Compact all batch files for an index into one.
|
||||||
|
pub async fn compact(&self, index_name: &str) -> Result<storaged::append_log::CompactStats, String> {
|
||||||
|
let log = self.log_for(index_name).await?;
|
||||||
|
log.compact().await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Current champion for an index by the named metric.
|
||||||
|
/// Valid metrics: `recall`, `latency`, `pareto`.
|
||||||
|
///
|
||||||
|
/// The `pareto` strategy is a placeholder — J should tune the scoring
|
||||||
|
/// function to match what matters in production. Right now it's a simple
|
||||||
|
/// weighted sum.
|
||||||
|
pub async fn best(
|
||||||
|
&self,
|
||||||
|
index_name: &str,
|
||||||
|
metric: &str,
|
||||||
|
) -> Result<Option<Trial>, String> {
|
||||||
|
let trials = self.list(index_name).await?;
|
||||||
|
if trials.is_empty() {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
|
||||||
|
let best = match metric {
|
||||||
|
"recall" => trials
|
||||||
|
.into_iter()
|
||||||
|
.max_by(|a, b| {
|
||||||
|
a.metrics
|
||||||
|
.recall_at_k
|
||||||
|
.partial_cmp(&b.metrics.recall_at_k)
|
||||||
|
.unwrap_or(std::cmp::Ordering::Equal)
|
||||||
|
})
|
||||||
|
.unwrap(),
|
||||||
|
"latency" => trials
|
||||||
|
.into_iter()
|
||||||
|
.min_by(|a, b| {
|
||||||
|
a.metrics
|
||||||
|
.search_latency_p95_us
|
||||||
|
.partial_cmp(&b.metrics.search_latency_p95_us)
|
||||||
|
.unwrap_or(std::cmp::Ordering::Equal)
|
||||||
|
})
|
||||||
|
.unwrap(),
|
||||||
|
"pareto" | _ => trials
|
||||||
|
.into_iter()
|
||||||
|
.max_by(|a, b| pareto_score(a).partial_cmp(&pareto_score(b)).unwrap())
|
||||||
|
.unwrap(),
|
||||||
|
};
|
||||||
|
Ok(Some(best))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Simple Pareto-style score: reward recall, penalize p95 latency.
|
||||||
|
/// Tunable — J should swap this in production to match what matters.
|
||||||
|
fn pareto_score(t: &Trial) -> f32 {
|
||||||
|
// Recall is [0, 1]. Latency is us — assume 100us baseline.
|
||||||
|
let recall = t.metrics.recall_at_k;
|
||||||
|
let latency_penalty = (t.metrics.search_latency_p95_us / 1000.0).min(1.0); // cap at 1ms
|
||||||
|
recall - 0.2 * latency_penalty
|
||||||
|
}
|
||||||
@@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"id": "0fd78303-9ad4-45fd-90d7-db95607d9ab1",
|
||||||
|
"name": "timesheets",
|
||||||
|
"schema_fingerprint": "auto",
|
||||||
|
"objects": [
|
||||||
|
{
|
||||||
|
"bucket": "data",
|
||||||
|
"key": "datasets/timesheets.parquet",
|
||||||
|
"size_bytes": 17539932,
|
||||||
|
"created_at": "2026-03-27T14:42:43.922019299Z"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"created_at": "2026-03-27T14:42:43.922025703Z",
|
||||||
|
"updated_at": "2026-03-27T14:42:43.922025703Z",
|
||||||
|
"description": "",
|
||||||
|
"owner": "",
|
||||||
|
"sensitivity": null,
|
||||||
|
"columns": [],
|
||||||
|
"lineage": null,
|
||||||
|
"freshness": null,
|
||||||
|
"tags": [],
|
||||||
|
"row_count": null
|
||||||
|
}
|
||||||
@@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"id": "1339f3d6-7677-47fb-8182-5f8e43f27cde",
|
||||||
|
"name": "job_orders",
|
||||||
|
"schema_fingerprint": "auto",
|
||||||
|
"objects": [
|
||||||
|
{
|
||||||
|
"bucket": "data",
|
||||||
|
"key": "datasets/job_orders.parquet",
|
||||||
|
"size_bytes": 905534,
|
||||||
|
"created_at": "2026-03-27T14:42:38.935718195Z"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"created_at": "2026-03-27T14:42:38.935724058Z",
|
||||||
|
"updated_at": "2026-03-27T14:42:38.935724058Z",
|
||||||
|
"description": "",
|
||||||
|
"owner": "",
|
||||||
|
"sensitivity": null,
|
||||||
|
"columns": [],
|
||||||
|
"lineage": null,
|
||||||
|
"freshness": null,
|
||||||
|
"tags": [],
|
||||||
|
"row_count": null
|
||||||
|
}
|
||||||
@@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"id": "1d8a065e-59c1-45ce-967b-398bc8370cbb",
|
||||||
|
"name": "candidates",
|
||||||
|
"schema_fingerprint": "auto",
|
||||||
|
"objects": [
|
||||||
|
{
|
||||||
|
"bucket": "data",
|
||||||
|
"key": "datasets/candidates.parquet",
|
||||||
|
"size_bytes": 10592165,
|
||||||
|
"created_at": "2026-03-27T14:42:38.823368759Z"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"created_at": "2026-03-27T14:42:38.823374843Z",
|
||||||
|
"updated_at": "2026-03-27T14:42:38.823374843Z",
|
||||||
|
"description": "",
|
||||||
|
"owner": "",
|
||||||
|
"sensitivity": null,
|
||||||
|
"columns": [],
|
||||||
|
"lineage": null,
|
||||||
|
"freshness": null,
|
||||||
|
"tags": [],
|
||||||
|
"row_count": null
|
||||||
|
}
|
||||||
@@ -1,15 +0,0 @@
|
|||||||
{
|
|
||||||
"id": "1f81445a-404f-48ea-bd72-00f6721ee18b",
|
|
||||||
"name": "employees",
|
|
||||||
"schema_fingerprint": "auto",
|
|
||||||
"objects": [
|
|
||||||
{
|
|
||||||
"bucket": "data",
|
|
||||||
"key": "datasets/employees.parquet",
|
|
||||||
"size_bytes": 1424,
|
|
||||||
"created_at": "2026-03-27T11:55:55.013996293Z"
|
|
||||||
}
|
|
||||||
],
|
|
||||||
"created_at": "2026-03-27T11:55:55.014010258Z",
|
|
||||||
"updated_at": "2026-03-27T11:55:55.014010258Z"
|
|
||||||
}
|
|
||||||
@@ -0,0 +1,65 @@
|
|||||||
|
{
|
||||||
|
"id": "3b12ae24-17d4-4325-92bf-f3155982f3bf",
|
||||||
|
"name": "users",
|
||||||
|
"schema_fingerprint": "30c0e31f0963e6f4af02131bbb9ea246fbbd068b849b833565a4b28211fbc90b",
|
||||||
|
"objects": [
|
||||||
|
{
|
||||||
|
"bucket": "data",
|
||||||
|
"key": "datasets/users.parquet",
|
||||||
|
"size_bytes": 2012,
|
||||||
|
"created_at": "2026-03-28T01:38:59.904968123Z"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"created_at": "2026-03-28T01:38:59.904968680Z",
|
||||||
|
"updated_at": "2026-03-28T01:38:59.905221340Z",
|
||||||
|
"description": "",
|
||||||
|
"owner": "",
|
||||||
|
"sensitivity": null,
|
||||||
|
"columns": [
|
||||||
|
{
|
||||||
|
"name": "id",
|
||||||
|
"data_type": "Int32",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "username",
|
||||||
|
"data_type": "Utf8",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "password_hash",
|
||||||
|
"data_type": "Utf8",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "role",
|
||||||
|
"data_type": "Utf8",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "created_at",
|
||||||
|
"data_type": "Utf8",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"lineage": {
|
||||||
|
"source_system": "postgresql",
|
||||||
|
"source_file": "127.0.0.1:5432/knowledge_base.users",
|
||||||
|
"ingest_job": "pg-import-1774661939904",
|
||||||
|
"ingest_timestamp": "2026-03-28T01:38:59.904968123Z",
|
||||||
|
"parent_datasets": []
|
||||||
|
},
|
||||||
|
"freshness": null,
|
||||||
|
"tags": [],
|
||||||
|
"row_count": 1
|
||||||
|
}
|
||||||
@@ -0,0 +1,93 @@
|
|||||||
|
{
|
||||||
|
"id": "3c2579b4-f3f3-4875-95fa-58d8b49ad94c",
|
||||||
|
"name": "meta_runs",
|
||||||
|
"schema_fingerprint": "68f2c0d7a3ceb0aaa3c17c64900704519c72d213161bc9e5179c42ee53f6d0df",
|
||||||
|
"objects": [
|
||||||
|
{
|
||||||
|
"bucket": "data",
|
||||||
|
"key": "datasets/meta_runs.parquet",
|
||||||
|
"size_bytes": 729773,
|
||||||
|
"created_at": "2026-03-28T01:38:57.380576453Z"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"created_at": "2026-03-28T01:38:57.380577270Z",
|
||||||
|
"updated_at": "2026-03-28T01:38:57.380846224Z",
|
||||||
|
"description": "",
|
||||||
|
"owner": "",
|
||||||
|
"sensitivity": null,
|
||||||
|
"columns": [
|
||||||
|
{
|
||||||
|
"name": "id",
|
||||||
|
"data_type": "Int32",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "pipeline_id",
|
||||||
|
"data_type": "Int32",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "iteration",
|
||||||
|
"data_type": "Int32",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "stage_results",
|
||||||
|
"data_type": "Utf8",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "final_output",
|
||||||
|
"data_type": "Utf8",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "score",
|
||||||
|
"data_type": "Float64",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "model_config",
|
||||||
|
"data_type": "Utf8",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "duration_ms",
|
||||||
|
"data_type": "Int32",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "created_at",
|
||||||
|
"data_type": "Utf8",
|
||||||
|
"sensitivity": null,
|
||||||
|
"description": "",
|
||||||
|
"is_pii": false
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"lineage": {
|
||||||
|
"source_system": "postgresql",
|
||||||
|
"source_file": "127.0.0.1:5432/knowledge_base.meta_runs",
|
||||||
|
"ingest_job": "pg-import-1774661937380",
|
||||||
|
"ingest_timestamp": "2026-03-28T01:38:57.380576453Z",
|
||||||
|
"parent_datasets": []
|
||||||
|
},
|
||||||
|
"freshness": null,
|
||||||
|
"tags": [],
|
||||||
|
"row_count": 38
|
||||||
|
}
|
||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user