use crate::engine::{is_internal_table, Engine};
use crate::error::{ErrorCode, McpError};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub const TABLE_CATALOG_TABLE: &str = "_table_catalog";
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct CatalogEntry {
pub table_name: String,
pub source_url: Option<String>,
pub source_description: Option<String>,
pub purpose: Option<String>,
pub load_tool: Option<String>,
pub load_params: Option<String>,
pub license: Option<String>,
pub loaded_at: DateTime<Utc>,
pub last_refreshed_at: DateTime<Utc>,
pub row_count: Option<i64>,
pub notes: Option<String>,
}
impl CatalogEntry {
#[must_use]
pub fn to_json(&self) -> Value {
serde_json::json!({
"table_name": self.table_name,
"source_url": self.source_url,
"source_description": self.source_description,
"purpose": self.purpose,
"load_tool": self.load_tool,
"load_params": self.load_params,
"license": self.license,
"loaded_at": self.loaded_at.to_rfc3339(),
"last_refreshed_at": self.last_refreshed_at.to_rfc3339(),
"row_count": self.row_count,
"notes": self.notes,
})
}
}
#[derive(Debug, Default, Clone)]
pub struct MetadataFields {
pub source_url: Option<String>,
pub source_description: Option<String>,
pub purpose: Option<String>,
pub license: Option<String>,
pub notes: Option<String>,
}
impl MetadataFields {
#[must_use]
pub fn is_empty(&self) -> bool {
self.source_url.is_none()
&& self.source_description.is_none()
&& self.purpose.is_none()
&& self.license.is_none()
&& self.notes.is_none()
}
}
const CATALOG_COLUMNS: &str = "(\
table_name TEXT NOT NULL, \
source_url TEXT, \
source_description TEXT, \
purpose TEXT, \
load_tool TEXT, \
load_params TEXT, \
license TEXT, \
loaded_at TIMESTAMP NOT NULL, \
last_refreshed_at TIMESTAMP NOT NULL, \
row_count BIGINT, \
notes TEXT\
)";
pub fn ensure_exists(engine: &Engine) -> Result<(), McpError> {
let ddl = format!("CREATE TABLE IF NOT EXISTS \"{TABLE_CATALOG_TABLE}\" {CATALOG_COLUMNS}");
engine.execute_command(&ddl)?;
Ok(())
}
pub fn ensure_exists_in_database(engine: &Engine, db_alias: &str) -> Result<(), McpError> {
let alias_esc = db_alias.replace('"', "\"\"");
let ddl = format!(
"CREATE TABLE IF NOT EXISTS \"{alias_esc}\".\"public\".\"{TABLE_CATALOG_TABLE}\" \
{CATALOG_COLUMNS}"
);
engine.execute_command(&ddl)?;
Ok(())
}
pub fn list(engine: &Engine) -> Result<Vec<CatalogEntry>, McpError> {
if !table_present(engine)? {
return Ok(Vec::new());
}
let sql = format!(
"SELECT table_name, source_url, source_description, purpose, \
load_tool, load_params, license, loaded_at, last_refreshed_at, \
row_count, notes \
FROM \"{TABLE_CATALOG_TABLE}\" ORDER BY table_name"
);
let rows = engine.execute_query_to_json(&sql)?;
rows.iter().map(row_to_entry).collect()
}
pub fn get(engine: &Engine, table_name: &str) -> Result<Option<CatalogEntry>, McpError> {
if !table_present(engine)? {
return Ok(None);
}
let sql = format!(
"SELECT table_name, source_url, source_description, purpose, \
load_tool, load_params, license, loaded_at, last_refreshed_at, \
row_count, notes \
FROM \"{TABLE_CATALOG_TABLE}\" WHERE table_name = {}",
sql_literal(table_name)
);
let rows = engine.execute_query_to_json(&sql)?;
match rows.first() {
Some(row) => Ok(Some(row_to_entry(row)?)),
None => Ok(None),
}
}
pub fn upsert_stub(
engine: &Engine,
table_name: &str,
load_tool: &str,
load_params: Option<&str>,
row_count: Option<i64>,
bump_refresh: bool,
) -> Result<(), McpError> {
ensure_exists(engine)?;
let existing = get(engine, table_name)?;
let now = Utc::now();
let loaded_at = existing.as_ref().map_or(now, |e| e.loaded_at);
let last_refreshed_at = if bump_refresh {
now
} else {
existing.as_ref().map_or(now, |e| e.last_refreshed_at)
};
let (source_url, source_description, purpose, license, notes) = match existing.as_ref() {
Some(e) => (
e.source_url.clone(),
e.source_description.clone(),
e.purpose.clone(),
e.license.clone(),
e.notes.clone(),
),
None => (None, None, None, None, None),
};
engine.execute_in_transaction(|engine| {
let delete_sql = format!(
"DELETE FROM \"{TABLE_CATALOG_TABLE}\" WHERE table_name = {}",
sql_literal(table_name)
);
engine.execute_command(&delete_sql)?;
let insert_sql = format!(
"INSERT INTO \"{TABLE_CATALOG_TABLE}\" \
(table_name, source_url, source_description, purpose, load_tool, \
load_params, license, loaded_at, last_refreshed_at, row_count, notes) \
VALUES ({name}, {source_url}, {source_description}, {purpose}, {load_tool}, \
{load_params}, {license}, TIMESTAMP {loaded_at}, TIMESTAMP {last_refreshed_at}, \
{row_count}, {notes})",
name = sql_literal(table_name),
source_url = opt_sql_literal(source_url.as_deref()),
source_description = opt_sql_literal(source_description.as_deref()),
purpose = opt_sql_literal(purpose.as_deref()),
load_tool = sql_literal(load_tool),
load_params = opt_sql_literal(load_params),
license = opt_sql_literal(license.as_deref()),
loaded_at = sql_literal(&format_timestamp(loaded_at)),
last_refreshed_at = sql_literal(&format_timestamp(last_refreshed_at)),
row_count = row_count.map_or_else(|| "NULL".into(), |n| n.to_string()),
notes = opt_sql_literal(notes.as_deref()),
);
engine.execute_command(&insert_sql)?;
Ok(())
})?;
Ok(())
}
pub fn set_metadata(
engine: &Engine,
table_name: &str,
fields: &MetadataFields,
) -> Result<CatalogEntry, McpError> {
ensure_exists(engine)?;
if fields.is_empty() {
return Err(McpError::new(
ErrorCode::EmptyData,
"set_table_metadata requires at least one of source_url, \
source_description, purpose, license, notes",
));
}
let existing = get(engine, table_name)?.ok_or_else(|| {
McpError::new(
ErrorCode::TableNotFound,
format!(
"No catalog entry for table '{table_name}'. Load the table first \
(load_file / load_data / execute CREATE TABLE) or create it and \
re-run; the catalog is refreshed automatically on those paths."
),
)
})?;
let mut assignments: Vec<String> = Vec::new();
if let Some(v) = &fields.source_url {
assignments.push(format!("source_url = {}", sql_literal_or_null_if_empty(v)));
}
if let Some(v) = &fields.source_description {
assignments.push(format!(
"source_description = {}",
sql_literal_or_null_if_empty(v)
));
}
if let Some(v) = &fields.purpose {
assignments.push(format!("purpose = {}", sql_literal_or_null_if_empty(v)));
}
if let Some(v) = &fields.license {
assignments.push(format!("license = {}", sql_literal_or_null_if_empty(v)));
}
if let Some(v) = &fields.notes {
assignments.push(format!("notes = {}", sql_literal_or_null_if_empty(v)));
}
let update_sql = format!(
"UPDATE \"{TABLE_CATALOG_TABLE}\" SET {assigns} WHERE table_name = {name}",
assigns = assignments.join(", "),
name = sql_literal(table_name),
);
engine.execute_command(&update_sql)?;
get(engine, table_name)?.map_or(Ok(existing), Ok)
}
pub fn delete_for(engine: &Engine, table_name: &str) -> Result<bool, McpError> {
if !table_present(engine)? {
return Ok(false);
}
let sql = format!(
"DELETE FROM \"{TABLE_CATALOG_TABLE}\" WHERE table_name = {}",
sql_literal(table_name)
);
let affected = engine.execute_command(&sql)?;
Ok(affected > 0)
}
pub fn reconcile(engine: &Engine) -> Result<(), McpError> {
ensure_exists(engine)?;
let tables = user_tables(engine)?;
let catalog_entries = list(engine)?;
let catalog_names: std::collections::HashSet<String> = catalog_entries
.iter()
.map(|e| e.table_name.clone())
.collect();
let live_tables: std::collections::HashSet<String> = tables.iter().cloned().collect();
for entry in &catalog_entries {
if !live_tables.contains(&entry.table_name) {
delete_for(engine, &entry.table_name)?;
}
}
for table in &tables {
let row_count = row_count_of(engine, table).ok();
if catalog_names.contains(table) {
refresh_row_count(engine, table, row_count)?;
} else {
upsert_stub(engine, table, "unknown", None, row_count, false)?;
}
}
Ok(())
}
fn user_tables(engine: &Engine) -> Result<Vec<String>, McpError> {
let describe = engine.describe_tables()?;
let mut names = Vec::new();
for table in describe {
if let Some(name) = table.get("name").and_then(|v| v.as_str()) {
if name == TABLE_CATALOG_TABLE || is_internal_table(name) {
continue;
}
names.push(name.to_string());
}
}
Ok(names)
}
fn table_present(engine: &Engine) -> Result<bool, McpError> {
let describe = engine.describe_tables()?;
Ok(describe.iter().any(|t| {
t.get("name")
.and_then(|v| v.as_str())
.is_some_and(|n| n == TABLE_CATALOG_TABLE)
}))
}
fn row_count_of(engine: &Engine, table_name: &str) -> Result<i64, McpError> {
let quoted = table_name.replace('"', "\"\"");
let sql = format!("SELECT COUNT(*) AS cnt FROM \"{quoted}\"");
let rows = engine.execute_query_to_json(&sql)?;
Ok(rows
.first()
.and_then(|r| r.get("cnt").and_then(serde_json::Value::as_i64))
.unwrap_or(0))
}
fn refresh_row_count(
engine: &Engine,
table_name: &str,
row_count: Option<i64>,
) -> Result<(), McpError> {
let sql = format!(
"UPDATE \"{TABLE_CATALOG_TABLE}\" SET row_count = {count} WHERE table_name = {name}",
count = row_count.map_or_else(|| "NULL".into(), |n| n.to_string()),
name = sql_literal(table_name),
);
engine.execute_command(&sql)?;
Ok(())
}
fn parse_timestamp(s: &str) -> Result<DateTime<Utc>, McpError> {
DateTime::parse_from_rfc3339(s)
.map(|d| d.with_timezone(&Utc))
.or_else(|_| {
chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f")
.or_else(|_| chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S"))
.map(|ndt| ndt.and_utc())
.map_err(|e| {
McpError::new(
ErrorCode::InternalError,
format!("Could not parse timestamp '{s}': {e}"),
)
})
})
}
fn format_timestamp(ts: DateTime<Utc>) -> String {
ts.format("%Y-%m-%d %H:%M:%S%.6f").to_string()
}
fn row_to_entry(row: &Value) -> Result<CatalogEntry, McpError> {
let table_name = row
.get("table_name")
.and_then(|v| v.as_str())
.ok_or_else(|| {
McpError::new(
ErrorCode::InternalError,
"_table_catalog row missing 'table_name'",
)
})?
.to_string();
let str_field = |name: &str| row.get(name).and_then(|v| v.as_str()).map(str::to_string);
let loaded_at = parse_timestamp(row.get("loaded_at").and_then(|v| v.as_str()).unwrap_or(""))?;
let last_refreshed_at = parse_timestamp(
row.get("last_refreshed_at")
.and_then(|v| v.as_str())
.unwrap_or(""),
)?;
let row_count = row.get("row_count").and_then(serde_json::Value::as_i64);
Ok(CatalogEntry {
table_name,
source_url: str_field("source_url"),
source_description: str_field("source_description"),
purpose: str_field("purpose"),
load_tool: str_field("load_tool"),
load_params: str_field("load_params"),
license: str_field("license"),
loaded_at,
last_refreshed_at,
row_count,
notes: str_field("notes"),
})
}
fn sql_literal(s: &str) -> String {
format!("'{}'", s.replace('\'', "''"))
}
fn sql_literal_or_null_if_empty(s: &str) -> String {
if s.is_empty() {
"NULL".into()
} else {
sql_literal(s)
}
}
fn opt_sql_literal(v: Option<&str>) -> String {
match v {
Some(s) => sql_literal(s),
None => "NULL".into(),
}
}