use std::fmt::Write as _;
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";
fn qualified_catalog_in(engine: &Engine, target_db: Option<&str>) -> Option<String> {
let alias = match target_db {
None | Some("") => {
if engine.has_persistent() {
Engine::PERSISTENT_ALIAS.to_string()
} else {
return None;
}
}
Some(a) if a.eq_ignore_ascii_case(Engine::PERSISTENT_ALIAS) => {
if !engine.has_persistent() {
return None;
}
Engine::PERSISTENT_ALIAS.to_string()
}
Some(a) => a.to_ascii_lowercase(),
};
let alias_esc = alias.replace('"', "\"\"");
Some(format!(
"\"{alias_esc}\".\"public\".\"{TABLE_CATALOG_TABLE}\""
))
}
#[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>,
pub created_by: Option<String>,
pub last_modified_by: 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,
"created_by": self.created_by,
"last_modified_by": self.last_modified_by,
})
}
}
#[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, \
created_by TEXT, \
last_modified_by TEXT\
)";
pub fn ensure_exists_in(engine: &Engine, target_db: Option<&str>) -> Result<(), McpError> {
let Some(qualified) = qualified_catalog_in(engine, target_db) else {
return Ok(());
};
let ddl = format!("CREATE TABLE IF NOT EXISTS {qualified} {CATALOG_COLUMNS}");
engine.execute_command(&ddl)?;
for col in ["created_by TEXT", "last_modified_by TEXT"] {
let alter = format!("ALTER TABLE {qualified} ADD COLUMN {col}");
let _ = engine.execute_command(&alter);
}
if let Some(alias) = resolve_catalog_alias(engine, target_db) {
engine.mark_catalog_present_for(&alias);
}
Ok(())
}
pub fn ensure_exists(engine: &Engine) -> Result<(), McpError> {
ensure_exists_in(engine, None)
}
#[deprecated(since = "0.2.0", note = "use ensure_exists_in(engine, Some(db_alias))")]
pub fn ensure_exists_in_database(engine: &Engine, db_alias: &str) -> Result<(), McpError> {
ensure_exists_in(engine, Some(db_alias))
}
pub fn list_in(engine: &Engine, target_db: Option<&str>) -> Result<Vec<CatalogEntry>, McpError> {
let Some(qualified) = qualified_catalog_in(engine, target_db) else {
return Ok(Vec::new());
};
if !table_present_in(engine, target_db)? {
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, created_by, last_modified_by \
FROM {qualified} ORDER BY table_name"
);
let rows = engine.execute_query_to_json(&sql)?;
rows.iter().map(row_to_entry).collect()
}
pub fn list(engine: &Engine) -> Result<Vec<CatalogEntry>, McpError> {
list_in(engine, None)
}
pub fn get_in(
engine: &Engine,
table_name: &str,
target_db: Option<&str>,
) -> Result<Option<CatalogEntry>, McpError> {
let Some(qualified) = qualified_catalog_in(engine, target_db) else {
return Ok(None);
};
if !table_present_in(engine, target_db)? {
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, created_by, last_modified_by \
FROM {qualified} 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 get(engine: &Engine, table_name: &str) -> Result<Option<CatalogEntry>, McpError> {
get_in(engine, table_name, None)
}
pub fn upsert_stub_in(
engine: &Engine,
table_name: &str,
load_tool: &str,
load_params: Option<&str>,
row_count: Option<i64>,
bump_refresh: bool,
target_db: Option<&str>,
client_name: Option<&str>,
) -> Result<(), McpError> {
ensure_exists_in(engine, target_db)?;
let Some(qualified) = qualified_catalog_in(engine, target_db) else {
return Ok(());
};
let now = Utc::now();
let row_count_sql = row_count.map_or_else(|| "NULL".into(), |n: i64| n.to_string());
let client_sql = opt_sql_literal(client_name);
let set_clauses = if bump_refresh {
let mut s = format!(
"load_tool = {}, load_params = {}, row_count = {}, \
last_refreshed_at = TIMESTAMP {}",
sql_literal(load_tool),
opt_sql_literal(load_params),
row_count_sql,
sql_literal(&format_timestamp(now)),
);
if let Some(name) = client_name {
let _ = write!(s, ", last_modified_by = {}", sql_literal(name));
}
s
} else {
let mut s = format!(
"load_tool = {}, load_params = {}, row_count = {}",
sql_literal(load_tool),
opt_sql_literal(load_params),
row_count_sql,
);
if let Some(name) = client_name {
let _ = write!(s, ", last_modified_by = {}", sql_literal(name));
}
s
};
let update_sql = format!(
"UPDATE {qualified} SET {set_clauses} WHERE table_name = {}",
sql_literal(table_name),
);
let updated = engine.execute_command(&update_sql)?;
if updated == 0 {
let insert_sql = format!(
"INSERT INTO {qualified} \
(table_name, source_url, source_description, purpose, load_tool, \
load_params, license, loaded_at, last_refreshed_at, row_count, notes, \
created_by, last_modified_by) \
SELECT {name}, NULL, NULL, NULL, {load_tool}, \
{load_params}, NULL, TIMESTAMP {loaded_at}, TIMESTAMP {last_refreshed_at}, \
{row_count}, NULL, {created_by}, {modified_by} \
WHERE NOT EXISTS (SELECT 1 FROM {qualified} WHERE table_name = {name})",
name = sql_literal(table_name),
load_tool = sql_literal(load_tool),
load_params = opt_sql_literal(load_params),
loaded_at = sql_literal(&format_timestamp(now)),
last_refreshed_at = sql_literal(&format_timestamp(now)),
row_count = row_count_sql,
created_by = client_sql,
modified_by = client_sql,
);
engine.execute_command(&insert_sql)?;
}
Ok(())
}
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> {
upsert_stub_in(
engine,
table_name,
load_tool,
load_params,
row_count,
bump_refresh,
None,
None,
)
}
pub fn set_metadata_in(
engine: &Engine,
table_name: &str,
fields: &MetadataFields,
target_db: Option<&str>,
) -> Result<CatalogEntry, McpError> {
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_in(engine, table_name, target_db)?.ok_or_else(|| {
let where_clause = match target_db {
Some(alias) if !alias.eq_ignore_ascii_case(Engine::PERSISTENT_ALIAS) => {
format!(" in database '{alias}'")
}
_ => String::new(),
};
McpError::new(
ErrorCode::TableNotFound,
format!(
"No catalog entry for table '{table_name}'{where_clause}. 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 qualified = qualified_catalog_in(engine, target_db).ok_or_else(|| {
McpError::new(
ErrorCode::ReadOnlyViolation,
"set_table_metadata is unavailable in --ephemeral-only mode \
because the catalog has nowhere durable to live.",
)
})?;
let update_sql = format!(
"UPDATE {qualified} SET {assigns} WHERE table_name = {name}",
assigns = assignments.join(", "),
name = sql_literal(table_name),
);
engine.execute_command(&update_sql)?;
get_in(engine, table_name, target_db)?.map_or(Ok(existing), Ok)
}
pub fn set_metadata(
engine: &Engine,
table_name: &str,
fields: &MetadataFields,
) -> Result<CatalogEntry, McpError> {
set_metadata_in(engine, table_name, fields, None)
}
pub fn delete_for_in(
engine: &Engine,
table_name: &str,
target_db: Option<&str>,
) -> Result<bool, McpError> {
let Some(qualified) = qualified_catalog_in(engine, target_db) else {
return Ok(false);
};
if !table_present_in(engine, target_db)? {
return Ok(false);
}
let sql = format!(
"DELETE FROM {qualified} WHERE table_name = {}",
sql_literal(table_name)
);
let affected = engine.execute_command(&sql)?;
Ok(affected > 0)
}
pub fn delete_for(engine: &Engine, table_name: &str) -> Result<bool, McpError> {
delete_for_in(engine, table_name, None)
}
pub fn reconcile_in(engine: &Engine, target_db: Option<&str>) -> Result<(), McpError> {
ensure_exists_in(engine, target_db)?;
let tables = user_tables_in(engine, target_db)?;
let catalog_entries = list_in(engine, target_db)?;
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_in(engine, &entry.table_name, target_db)?;
}
}
for table in &tables {
let row_count = row_count_of_in(engine, table, target_db).ok();
if catalog_names.contains(table) {
refresh_row_count_in(engine, table, row_count, target_db)?;
} else {
upsert_stub_in(
engine, table, "unknown", None, row_count, false, target_db, None,
)?;
}
}
Ok(())
}
pub fn reconcile(engine: &Engine) -> Result<(), McpError> {
reconcile_in(engine, None)
}
fn user_tables_in(engine: &Engine, target_db: Option<&str>) -> Result<Vec<String>, McpError> {
let Some(alias) = resolve_catalog_alias(engine, target_db) else {
return Ok(Vec::new());
};
let alias_esc = alias.replace('"', "\"\"");
let sql = format!(
"SELECT tablename FROM \"{alias_esc}\".pg_catalog.pg_tables \
WHERE schemaname = 'public'"
);
let rows = engine.execute_query_to_json(&sql)?;
Ok(rows
.into_iter()
.filter_map(|r| {
r.get("tablename")
.and_then(|v| v.as_str())
.map(str::to_string)
})
.filter(|name| name != TABLE_CATALOG_TABLE && !is_internal_table(name))
.collect())
}
fn table_present_in(engine: &Engine, target_db: Option<&str>) -> Result<bool, McpError> {
let Some(alias) = resolve_catalog_alias(engine, target_db) else {
return Ok(false);
};
let alias_for_probe = alias.clone();
engine.catalog_present_in(&alias, move |engine| {
let alias_esc = alias_for_probe.replace('"', "\"\"");
let sql = format!(
"SELECT tablename FROM \"{alias_esc}\".pg_catalog.pg_tables \
WHERE schemaname = 'public' AND tablename = {}",
sql_literal(TABLE_CATALOG_TABLE)
);
let rows = engine.execute_query_to_json(&sql)?;
Ok(!rows.is_empty())
})
}
fn row_count_of_in(
engine: &Engine,
table_name: &str,
target_db: Option<&str>,
) -> Result<i64, McpError> {
let Some(alias) = resolve_catalog_alias(engine, target_db) else {
return Ok(0);
};
let alias_esc = alias.replace('"', "\"\"");
let quoted = table_name.replace('"', "\"\"");
let sql = format!("SELECT COUNT(*) AS cnt FROM \"{alias_esc}\".\"public\".\"{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_in(
engine: &Engine,
table_name: &str,
row_count: Option<i64>,
target_db: Option<&str>,
) -> Result<(), McpError> {
let Some(qualified) = qualified_catalog_in(engine, target_db) else {
return Ok(());
};
let sql = format!(
"UPDATE {qualified} 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 resolve_catalog_alias(engine: &Engine, target_db: Option<&str>) -> Option<String> {
match target_db {
None | Some("") => {
if engine.has_persistent() {
Some(Engine::PERSISTENT_ALIAS.to_string())
} else {
None
}
}
Some(a) if a.eq_ignore_ascii_case(Engine::PERSISTENT_ALIAS) => {
if engine.has_persistent() {
Some(Engine::PERSISTENT_ALIAS.to_string())
} else {
None
}
}
Some(a) => Some(a.to_ascii_lowercase()),
}
}
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"),
created_by: str_field("created_by"),
last_modified_by: str_field("last_modified_by"),
})
}
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(),
}
}