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,28 @@
|
||||
[package]
|
||||
name = "ingestd"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
shared = { path = "../shared" }
|
||||
storaged = { path = "../storaged" }
|
||||
catalogd = { path = "../catalogd" }
|
||||
vectord = { path = "../vectord" }
|
||||
journald = { path = "../journald" }
|
||||
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 }
|
||||
croner = "3"
|
||||
@@ -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,725 @@
|
||||
//! 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:
|
||||
//! - Sub-minute cron precision. 5-field Unix cron is supported; seconds
|
||||
//! are pinned to 0. Intervals cover sub-minute cases; 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 },
|
||||
/// 5-field Unix cron expression: "min hour dom month dow".
|
||||
/// Seconds are implicitly 0 (minute-granularity scheduling).
|
||||
/// Examples: "15 14 * * *" = 14:15 UTC daily, "0 */6 * * *" = every
|
||||
/// 6 hours on the hour, "0 9 * * 1-5" = 09:00 UTC weekdays.
|
||||
/// Timezone: UTC always — the scheduler's clock is UTC and we don't
|
||||
/// carry a per-schedule timezone field.
|
||||
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.
|
||||
//
|
||||
pub fn compute_next_run_at(
|
||||
trigger: &ScheduleTrigger,
|
||||
completed_at: DateTime<Utc>,
|
||||
_previous_next: DateTime<Utc>,
|
||||
) -> DateTime<Utc> {
|
||||
match trigger {
|
||||
ScheduleTrigger::Interval { secs } => {
|
||||
completed_at + Duration::seconds(*secs as i64)
|
||||
}
|
||||
ScheduleTrigger::Cron { expr } => {
|
||||
next_cron_fire(expr, completed_at)
|
||||
// Creation-time validation (see `validate_trigger`) rejects
|
||||
// unparseable expressions, so this fallback is only reached
|
||||
// if the schedule JSON was hand-edited on disk. One-hour
|
||||
// window keeps a bad expr from fire-looping the system.
|
||||
.unwrap_or(completed_at + Duration::hours(1))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a Vixie/POSIX cron expression and return the next fire time
|
||||
/// strictly after `after`. Accepts 5-field (`min hour dom month dow`) and
|
||||
/// 6-field (with leading `sec`) forms natively — croner handles both.
|
||||
/// Day-of-week follows Unix convention: 0=Sun, 1=Mon, …, 6=Sat (7=Sun as
|
||||
/// synonym). Returns None if the expression fails to parse or produces no
|
||||
/// upcoming match (which can happen for year-bounded patterns in the past).
|
||||
fn next_cron_fire(expr: &str, after: DateTime<Utc>) -> Option<DateTime<Utc>> {
|
||||
use std::str::FromStr;
|
||||
let cron = croner::Cron::from_str(expr).ok()?;
|
||||
cron.find_next_occurrence(&after, false).ok()
|
||||
}
|
||||
|
||||
/// Reject a ScheduleTrigger whose content can't be interpreted — used by
|
||||
/// the HTTP handlers to fail fast at create/patch time rather than falling
|
||||
/// back silently at fire time. Interval triggers get a sanity gate
|
||||
/// (no zero secs); cron triggers get a full parse through croner.
|
||||
pub fn validate_trigger(trigger: &ScheduleTrigger) -> Result<(), String> {
|
||||
match trigger {
|
||||
ScheduleTrigger::Interval { secs } => {
|
||||
if *secs == 0 {
|
||||
return Err("interval secs must be > 0".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
ScheduleTrigger::Cron { expr } => {
|
||||
use std::str::FromStr;
|
||||
croner::Cron::from_str(expr)
|
||||
.map_err(|e| format!("invalid cron expression '{expr}': {e}"))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// =================== 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::*;
|
||||
use chrono::Timelike;
|
||||
|
||||
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_every_minute_fires_within_sixty_seconds() {
|
||||
let trig = ScheduleTrigger::Cron { expr: "* * * * *".into() };
|
||||
let now = Utc::now();
|
||||
let next = compute_next_run_at(&trig, now, now);
|
||||
let delta = next - now;
|
||||
assert!(delta.num_seconds() > 0 && delta.num_seconds() <= 60,
|
||||
"expected next fire within 60s, got {}s", delta.num_seconds());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cron_daily_at_1415_utc_is_within_24h() {
|
||||
// 14:15 UTC daily — whether we're before or after 14:15 today, the
|
||||
// next fire is at most 24h out.
|
||||
let trig = ScheduleTrigger::Cron { expr: "15 14 * * *".into() };
|
||||
let now: DateTime<Utc> = "2026-04-20T10:00:00Z".parse().unwrap();
|
||||
let next = compute_next_run_at(&trig, now, now);
|
||||
assert!(next > now);
|
||||
assert!((next - now).num_hours() <= 24);
|
||||
// And the minute-of-hour is 15.
|
||||
assert_eq!(next.minute(), 15);
|
||||
assert_eq!(next.hour(), 14);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cron_weekday_skips_weekend() {
|
||||
// 09:00 UTC weekdays only. On a Saturday, next fire is Monday.
|
||||
let trig = ScheduleTrigger::Cron { expr: "0 9 * * 1-5".into() };
|
||||
// 2026-04-18 is a Saturday.
|
||||
let sat: DateTime<Utc> = "2026-04-18T10:00:00Z".parse().unwrap();
|
||||
let next = compute_next_run_at(&trig, sat, sat);
|
||||
// Monday 2026-04-20 at 09:00 UTC.
|
||||
assert_eq!(next, "2026-04-20T09:00:00Z".parse::<DateTime<Utc>>().unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cron_six_field_seconds_granularity() {
|
||||
// 6-field (seconds included): fire every minute at 30s past the minute.
|
||||
let trig = ScheduleTrigger::Cron { expr: "30 * * * * *".into() };
|
||||
let now = Utc::now();
|
||||
let next = compute_next_run_at(&trig, now, now);
|
||||
assert_eq!(next.second(), 30);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_rejects_bad_cron() {
|
||||
let bad = ScheduleTrigger::Cron { expr: "not a cron".into() };
|
||||
let err = validate_trigger(&bad).unwrap_err();
|
||||
assert!(err.contains("cron"), "error should mention cron, got: {err}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_rejects_wrong_field_count() {
|
||||
// 4 fields — neither 5 nor 6.
|
||||
let bad = ScheduleTrigger::Cron { expr: "* * * *".into() };
|
||||
assert!(validate_trigger(&bad).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_rejects_zero_interval() {
|
||||
let bad = ScheduleTrigger::Interval { secs: 0 };
|
||||
assert!(validate_trigger(&bad).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_accepts_good_cron() {
|
||||
let good = ScheduleTrigger::Cron { expr: "0 */6 * * *".into() };
|
||||
assert!(validate_trigger(&good).is_ok());
|
||||
}
|
||||
|
||||
#[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, SchemaRef};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// 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, Schema};
|
||||
use std::sync::Arc;
|
||||
|
||||
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,758 @@
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::{Multipart, Path, Query, State},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::IntoResponse,
|
||||
routing::{get, 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,
|
||||
/// Event journal for ADR-012 mutation history. Optional for back-compat
|
||||
/// with callers (like scheduled ingest tests) that don't wire it yet.
|
||||
/// When present, successful ingests emit a record_ingest event — closes
|
||||
/// P9-001 on the file-upload path. (2026-04-23)
|
||||
pub journal: Option<journald::journal::Journal>,
|
||||
}
|
||||
|
||||
/// 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;
|
||||
// P9-001 fix (2026-04-23): emit a mutation event on every
|
||||
// non-deduplicated ingest. Dedup no-ops don't need events
|
||||
// (ADR-020 register() is already idempotent on same fingerprint).
|
||||
if let Some(ref journal) = state.journal {
|
||||
if let Err(e) = journal.record_ingest(
|
||||
&result.dataset_name,
|
||||
result.rows as usize,
|
||||
"ingest_api",
|
||||
&filename,
|
||||
).await {
|
||||
tracing::warn!(
|
||||
"journal record_ingest failed for '{}': {}",
|
||||
result.dataset_name, e,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
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)));
|
||||
}
|
||||
if let Err(e) = schedule::validate_trigger(&req.trigger) {
|
||||
return Err((StatusCode::BAD_REQUEST, e));
|
||||
}
|
||||
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(t) = &req.trigger {
|
||||
if let Err(e) = schedule::validate_trigger(t) {
|
||||
return Err((StatusCode::BAD_REQUEST, e));
|
||||
}
|
||||
}
|
||||
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))
|
||||
}
|
||||
|
||||
// ─── Tests ───
|
||||
|
||||
#[cfg(test)]
|
||||
mod journal_integration_tests {
|
||||
//! P9-001 integration test: prove that a successful ingest produces a
|
||||
//! journal.record_ingest event. Block 2 on PR #10 was "journal event
|
||||
//! verified live" being unbacked by the diff. This test makes the
|
||||
//! verification committed and reproducible.
|
||||
|
||||
use journald::journal::{Event, Journal};
|
||||
use object_store::memory::InMemory;
|
||||
use std::sync::Arc;
|
||||
|
||||
// Helper: build a bare Journal against an in-memory object store.
|
||||
// Flush threshold 1 so every recorded event is persisted immediately.
|
||||
fn test_journal() -> Journal {
|
||||
let store: Arc<dyn object_store::ObjectStore> = Arc::new(InMemory::new());
|
||||
Journal::new(store, 1)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn journal_record_ingest_increments_counter() {
|
||||
// Arrange — fresh journal, counter starts at zero.
|
||||
let journal = test_journal();
|
||||
let stats0 = journal.stats().await;
|
||||
assert_eq!(stats0.total_events_created, 0);
|
||||
assert_eq!(stats0.buffer_events, 0);
|
||||
|
||||
// Act — simulate what the /ingest/file success path does.
|
||||
journal
|
||||
.record_ingest("test_dataset", 42, "ingest_api", "probe.csv")
|
||||
.await
|
||||
.expect("record_ingest should succeed");
|
||||
|
||||
// Assert — counter advanced, event exists. With threshold=1 the
|
||||
// event flushed to store; with threshold>N it would be in-buffer.
|
||||
let stats1 = journal.stats().await;
|
||||
assert_eq!(stats1.total_events_created, 1, "counter should reflect one recorded event");
|
||||
|
||||
// Assert — the event is retrievable by entity.
|
||||
let history = journal
|
||||
.get_entity_history("batch:42")
|
||||
.await
|
||||
.expect("history lookup");
|
||||
assert_eq!(history.len(), 1, "one event should be visible in history");
|
||||
let ev = &history[0];
|
||||
assert_eq!(ev.action, "ingest");
|
||||
assert_eq!(ev.entity_type, "test_dataset");
|
||||
assert_eq!(ev.actor, "ingest_api");
|
||||
assert!(
|
||||
ev.new_value.contains("probe.csv"),
|
||||
"new_value should carry source filename, got: {}",
|
||||
ev.new_value
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn optional_journal_field_none_is_valid_back_compat() {
|
||||
// IngestState.journal is Option<Journal>. Back-compat path: when
|
||||
// the field is None, the ingest handler MUST still succeed — the
|
||||
// journal call is fire-and-forget, never load-bearing.
|
||||
//
|
||||
// This test asserts the type shape: Option<Journal> is what we
|
||||
// expect. If a refactor makes it mandatory, this test forces an
|
||||
// explicit re-consideration.
|
||||
let none_journal: Option<Journal> = None;
|
||||
assert!(none_journal.is_none());
|
||||
|
||||
let some_journal: Option<Journal> = Some(test_journal());
|
||||
assert!(some_journal.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn journal_record_event_fields_match_adr_012_schema() {
|
||||
// ADR-012 locks the event schema: entity_type, entity_id, field,
|
||||
// action, old_value, new_value, actor, source, workspace_id plus
|
||||
// the auto-assigned event_id + timestamp. This test pins the
|
||||
// field names so a future refactor can't silently drop one.
|
||||
let journal = test_journal();
|
||||
let base = Event {
|
||||
event_id: String::new(),
|
||||
timestamp: chrono::Utc::now(),
|
||||
entity_type: "candidate".into(),
|
||||
entity_id: "CAND-0001".into(),
|
||||
field: "phone".into(),
|
||||
action: "update".into(),
|
||||
old_value: "555-0000".into(),
|
||||
new_value: "555-9999".into(),
|
||||
actor: "recruiter".into(),
|
||||
source: "api".into(),
|
||||
workspace_id: "ws-x".into(),
|
||||
};
|
||||
journal.record(base).await.expect("record should accept full-schema event");
|
||||
let h = journal
|
||||
.get_entity_history("CAND-0001")
|
||||
.await
|
||||
.expect("lookup");
|
||||
assert_eq!(h.len(), 1);
|
||||
assert_eq!(h[0].field, "phone");
|
||||
assert_eq!(h[0].old_value, "555-0000");
|
||||
assert_eq!(h[0].new_value, "555-9999");
|
||||
assert_eq!(h[0].workspace_id, "ws-x");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
/// 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. Bind filename once via
|
||||
// let-else so the subsequent use is unwrap-free — previous
|
||||
// version relied on a map_or guard above + an .unwrap() here
|
||||
// being consistent, which is a fragile invariant.
|
||||
if path.is_dir() { continue; }
|
||||
let Some(fn_os) = path.file_name() else { continue; };
|
||||
let filename = fn_os.to_string_lossy().to_string();
|
||||
if filename.starts_with('.') { continue; }
|
||||
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());
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user