checkpoint: matrix-agent-validated (2026-04-25)
Architectural snapshot of the lakehouse codebase at the point where the
full matrix-driven agent loop with Mem0 versioning + deletion was
validated end-to-end.
WHAT THIS REPO IS
A clean single-commit snapshot of the lakehouse code. Heavy test data
(.parquet datasets, vector indexes) excluded — see REPLICATION.md for
regen path. Full lakehouse history at git.agentview.dev/profit/lakehouse.
WHAT WAS PROVEN
- Vector retrieval across multi-corpora matrix (chicago_permits + entity
briefs + sec_tickers + distilled procedural + llm_team runs)
- Observer hand-review (cloud + heuristic fallback) gating each candidate
- Local-model agent loop (qwen3.5:latest) with tool use + scratchpad
- Playbook seal on success → next-iter retrieval surfaces it as preamble
- Mem0 versioning + deletion in pathway_memory:
* UPSERT: ADD on new workflow, UPDATE bumps replay_count on identical
* REVISE: chains versions, parent.superseded_at + superseded_by stamped
* RETIRE: marks specific trace retired with reason, excluded from retrieval
* HISTORY: walks chain root→tip, cycle-safe
KEY DIRECTORIES
- crates/vectord/src/pathway_memory.rs — Mem0 ops live here
- crates/vectord/src/playbook_memory.rs — original Mem0 reference
- tests/agent_test/ — local-model agent harness + PRD + session archives
- scripts/dump_raw_corpus.sh — MinIO bucket dump (raw test corpus)
- scripts/vectorize_raw_corpus.ts — corpus → vector indexes
- scripts/analyze_chicago_contracts.ts — real inference pipeline
- scripts/seal_agent_playbook.ts — Mem0 upsert from agent traces
Replication: see REPLICATION.md for Debian 13 clean install + cloud-only
adaptation (no local Ollama).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
profit
@@ -0,0 +1,23 @@
|
||||
[package]
|
||||
name = "queryd"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
shared = { path = "../shared" }
|
||||
catalogd = { path = "../catalogd" }
|
||||
storaged = { path = "../storaged" }
|
||||
truth = { path = "../truth" }
|
||||
tokio = { workspace = true }
|
||||
axum = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
datafusion = { workspace = true }
|
||||
object_store = { workspace = true }
|
||||
arrow = { workspace = true }
|
||||
parquet = { workspace = true }
|
||||
bytes = { workspace = true }
|
||||
futures = { 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>,
|
||||
}
|
||||
@@ -0,0 +1,328 @@
|
||||
use catalogd::registry::Registry;
|
||||
use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl};
|
||||
use datafusion::datasource::MemTable;
|
||||
use datafusion::datasource::file_format::parquet::ParquetFormat;
|
||||
use datafusion::prelude::*;
|
||||
use object_store::ObjectStore;
|
||||
use std::sync::Arc;
|
||||
use storaged::registry::BucketRegistry;
|
||||
use url::Url;
|
||||
|
||||
use crate::cache::MemCache;
|
||||
use crate::delta;
|
||||
|
||||
const STORE_SCHEME: &str = "lakehouse";
|
||||
|
||||
/// 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)]
|
||||
pub struct QueryEngine {
|
||||
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>,
|
||||
cache: MemCache,
|
||||
}
|
||||
|
||||
impl QueryEngine {
|
||||
pub fn new(registry: Registry, buckets: Arc<BucketRegistry>, cache: MemCache) -> Self {
|
||||
let store = buckets.default_store();
|
||||
Self { registry, buckets, store, cache }
|
||||
}
|
||||
|
||||
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> {
|
||||
let ctx = self.build_context().await?;
|
||||
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}"))?;
|
||||
Ok(batches)
|
||||
}
|
||||
|
||||
/// 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> {
|
||||
let ctx = SessionContext::new();
|
||||
|
||||
// Phase E: snapshot tombstones by dataset before registering tables
|
||||
// so we can wrap tombstoned tables in a filter view. The underlying
|
||||
// base table is registered under an internal name `__raw__{dataset}`
|
||||
// and the public `{dataset}` name becomes the filtered view.
|
||||
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());
|
||||
|
||||
let datasets = self.registry.list().await;
|
||||
for dataset in &datasets {
|
||||
if dataset.objects.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// 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()
|
||||
.filter_map(|o| {
|
||||
let bucket = if o.bucket.is_empty() { "primary" } else { &o.bucket };
|
||||
ListingTableUrl::parse(&bucket_object_url(bucket, &o.key)).ok()
|
||||
})
|
||||
.collect();
|
||||
|
||||
if table_paths.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let schema = opts.infer_schema(&ctx.state(), &table_paths[0]).await
|
||||
.map_err(|e| format!("schema inference failed for {}: {e}", dataset.name))?;
|
||||
let config = ListingTableConfig::new_with_multi_paths(table_paths)
|
||||
.with_listing_options(opts)
|
||||
.with_schema(schema);
|
||||
let table = ListingTable::try_new(config)
|
||||
.map_err(|e| format!("table creation failed for {}: {e}", dataset.name))?;
|
||||
|
||||
// Decide the registration name: if this dataset has any
|
||||
// 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()
|
||||
};
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,272 @@
|
||||
/// 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();
|
||||
// Row counts captured before extend; previously base_rows subtracted delta_count (files) from rows — unit mismatch.
|
||||
let base_row_count: usize = base_batches.iter().map(|b| b.num_rows()).sum();
|
||||
let delta_row_count: usize = delta_batches.iter().map(|b| b.num_rows()).sum();
|
||||
|
||||
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_row_count,
|
||||
delta_rows: 0,
|
||||
final_rows: base_row_count,
|
||||
deltas_merged: 0,
|
||||
tombstones_applied: 0,
|
||||
rows_dropped_by_tombstones: 0,
|
||||
});
|
||||
}
|
||||
|
||||
base_batches.extend(delta_batches);
|
||||
let pre_filter_rows: usize = base_row_count + delta_row_count;
|
||||
|
||||
// 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: base_row_count,
|
||||
delta_rows: delta_row_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,
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
pub mod cache;
|
||||
pub mod context;
|
||||
pub mod delta;
|
||||
pub mod paged;
|
||||
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}"))
|
||||
}
|
||||
@@ -0,0 +1,332 @@
|
||||
use arrow::array::RecordBatch;
|
||||
use arrow::json::writer::{JsonArray, Writer as JsonWriter};
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::{Path, Query, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
routing::{get, post},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use std::sync::Arc;
|
||||
use truth::{RuleAction, TruthStore};
|
||||
|
||||
use crate::context::QueryEngine;
|
||||
use crate::delta;
|
||||
use crate::paged::ResultStore;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct QueryState {
|
||||
pub engine: QueryEngine,
|
||||
pub result_store: ResultStore,
|
||||
// Policy gate for incoming SQL. Every /sql and /paged request is
|
||||
// evaluated against this store before hitting DataFusion. Added for
|
||||
// P42-002 ("raw SQL forwarded without schema or policy gate") after
|
||||
// the scrum master's queryd/service.rs finding looped across iters
|
||||
// 3-5 without ever being reachable by the 6-line auto-applier.
|
||||
pub truth: Arc<TruthStore>,
|
||||
}
|
||||
|
||||
pub fn router(engine: QueryEngine) -> Router {
|
||||
router_with_truth(engine, Arc::new(truth::sql_query_guard_store()))
|
||||
}
|
||||
|
||||
/// Test/integration hook: construct the router with a caller-supplied
|
||||
/// TruthStore so tests can assert reject/pass behavior deterministically
|
||||
/// without depending on the default needle list.
|
||||
pub fn router_with_truth(engine: QueryEngine, truth: Arc<TruthStore>) -> Router {
|
||||
let state = QueryState {
|
||||
engine: engine.clone(),
|
||||
result_store: ResultStore::new(100, 50), // 100 rows/page, keep 50 results
|
||||
truth,
|
||||
};
|
||||
Router::new()
|
||||
.route("/health", get(health))
|
||||
.route("/sql", post(execute_query))
|
||||
.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 {
|
||||
"queryd ok"
|
||||
}
|
||||
|
||||
// --- SQL Query ---
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct QueryRequest {
|
||||
sql: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct QueryResponse {
|
||||
columns: Vec<ColumnInfo>,
|
||||
rows: serde_json::Value,
|
||||
row_count: usize,
|
||||
// Elapsed wall time from handler entry to response. Required for
|
||||
// audit-log parity — gateway's audit row previously stored null here.
|
||||
// Scrum iter 9 finding, populated from std::time::Instant captured
|
||||
// at the top of execute_query / paged_query.
|
||||
latency_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ColumnInfo {
|
||||
name: String,
|
||||
data_type: String,
|
||||
}
|
||||
|
||||
fn batches_to_json(batches: &[RecordBatch]) -> Result<serde_json::Value, String> {
|
||||
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 error: {e}"))?;
|
||||
}
|
||||
writer.finish().map_err(|e| format!("JSON finish error: {e}"))?;
|
||||
drop(writer);
|
||||
serde_json::from_slice(&buf).map_err(|e| format!("JSON parse error: {e}"))
|
||||
}
|
||||
|
||||
/// Evaluate the request SQL against the configured TruthStore. Returns
|
||||
/// the Reject/Block message on the first failing mandatory rule so the
|
||||
/// handler can short-circuit. Returns None when all rules pass (or when
|
||||
/// the failures' declared action is non-mandatory like Redact/Pass).
|
||||
fn sql_policy_check(truth: &TruthStore, sql: &str) -> Option<String> {
|
||||
let ctx = serde_json::json!({ "sql": sql });
|
||||
for outcome in truth.evaluate("sql_query", &ctx) {
|
||||
if !outcome.passed {
|
||||
// FieldEmpty / FieldContainsAny etc. are enforced only when
|
||||
// condition HOLDS (i.e. passed=true). Below means "passed=false",
|
||||
// so the rule condition did not hold — no enforcement.
|
||||
continue;
|
||||
}
|
||||
match &outcome.action {
|
||||
RuleAction::Reject { message } | RuleAction::Block { message } => {
|
||||
return Some(message.clone());
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
async fn execute_query(
|
||||
State(state): State<QueryState>,
|
||||
Json(req): Json<QueryRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let started = std::time::Instant::now();
|
||||
tracing::info!("executing query: {}", req.sql);
|
||||
|
||||
if let Some(reason) = sql_policy_check(&state.truth, &req.sql) {
|
||||
tracing::warn!("sql rejected by truth gate: {reason}");
|
||||
return Err((StatusCode::FORBIDDEN, reason));
|
||||
}
|
||||
|
||||
match state.engine.query(&req.sql).await {
|
||||
Ok(batches) => {
|
||||
if batches.is_empty() {
|
||||
return Ok(Json(QueryResponse {
|
||||
columns: vec![],
|
||||
rows: serde_json::Value::Array(vec![]),
|
||||
row_count: 0,
|
||||
latency_ms: started.elapsed().as_millis() as u64,
|
||||
}));
|
||||
}
|
||||
|
||||
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 rows = batches_to_json(&batches)
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||
|
||||
let row_count = rows.as_array().map(|a| a.len()).unwrap_or(0);
|
||||
|
||||
Ok(Json(QueryResponse {
|
||||
columns,
|
||||
rows,
|
||||
row_count,
|
||||
latency_ms: started.elapsed().as_millis() as u64,
|
||||
}))
|
||||
}
|
||||
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);
|
||||
if let Some(reason) = sql_policy_check(&state.truth, &req.sql) {
|
||||
tracing::warn!("paged sql rejected by truth gate: {reason}");
|
||||
return Err((StatusCode::FORBIDDEN, reason));
|
||||
}
|
||||
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)),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod sql_policy_tests {
|
||||
use super::*;
|
||||
use truth::sql_query_guard_store;
|
||||
|
||||
// These tests exercise the policy gate without spinning up a DataFusion
|
||||
// engine — they only need `TruthStore`. Purpose: prove the P42-002
|
||||
// enforcement point actually rejects destructive SQL. This is the
|
||||
// regression guard for the queryd/service.rs finding that looped
|
||||
// across scrum iters 3-5.
|
||||
|
||||
#[test]
|
||||
fn blocks_drop_table() {
|
||||
let store = sql_query_guard_store();
|
||||
let reason = sql_policy_check(&store, "DROP TABLE users").expect("must reject");
|
||||
assert!(reason.contains("destructive"), "reason: {reason}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blocks_delete_from() {
|
||||
let store = sql_query_guard_store();
|
||||
assert!(sql_policy_check(&store, "delete from t where 1=1").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blocks_truncate() {
|
||||
let store = sql_query_guard_store();
|
||||
assert!(sql_policy_check(&store, "TRUNCATE workers").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blocks_empty_sql() {
|
||||
let store = sql_query_guard_store();
|
||||
assert!(sql_policy_check(&store, "").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allows_benign_select() {
|
||||
let store = sql_query_guard_store();
|
||||
assert!(sql_policy_check(&store, "SELECT count(*) FROM workers").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allows_select_with_deleted_word_in_column() {
|
||||
// Substring match is narrow ("delete from", not "delete"), so a
|
||||
// column named `deleted_at` doesn't trip the guard. Important
|
||||
// check — false positives on benign queries would make the gate
|
||||
// unusable in practice.
|
||||
let store = sql_query_guard_store();
|
||||
assert!(
|
||||
sql_policy_check(&store, "SELECT deleted_at FROM t").is_none(),
|
||||
"column names containing 'delete' must not be rejected"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn case_insensitive_match_catches_mixed_case() {
|
||||
let store = sql_query_guard_store();
|
||||
assert!(sql_policy_check(&store, "Drop Table X").is_some());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,269 @@
|
||||
/// 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 chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
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)),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user