use std::collections::HashMap;
use base64::{Engine as _, engine::general_purpose::STANDARD as B64};
use nexql_index::{IndexOverrides, IndexStore, JoinEdge, JoinGraph, ObjectEntry};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
pub const PAGE_SIZE: usize = 200;
pub const ERR_INVALID_PARAMS: i32 = -32602;
pub const ERR_RESOURCE_NOT_FOUND: i32 = -32002;
#[derive(Debug, Error)]
pub enum ResourceError {
#[error("{0}")]
InvalidParams(String),
#[error("{0}")]
NotFound(String),
#[error("{0}")]
Internal(String),
}
impl ResourceError {
pub fn code(&self) -> i32 {
match self {
Self::InvalidParams(_) => ERR_INVALID_PARAMS,
Self::NotFound(_) => ERR_RESOURCE_NOT_FOUND,
Self::Internal(_) => -32603,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpResource {
pub uri: String,
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(rename = "mimeType")]
pub mime_type: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceListResult {
pub resources: Vec<McpResource>,
#[serde(rename = "nextCursor", skip_serializing_if = "Option::is_none")]
pub next_cursor: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceContents {
pub uri: String,
#[serde(rename = "mimeType")]
pub mime_type: String,
pub text: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceReadResult {
pub contents: Vec<ResourceContents>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceTemplate {
#[serde(rename = "uriTemplate")]
pub uri_template: String,
pub name: String,
pub description: String,
#[serde(rename = "mimeType")]
pub mime_type: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct CursorState {
db: usize,
offset: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ResourceKind {
Manifest,
JoinGraph,
Object { schema: String, name: String },
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ParsedUri {
pub connection_id: String,
pub database: String,
pub kind: ResourceKind,
}
pub struct ResourceProvider {
store: IndexStore,
}
impl ResourceProvider {
pub fn new(store: IndexStore) -> Self {
Self { store }
}
pub fn store(&self) -> &IndexStore {
&self.store
}
pub fn list(&self, cursor: Option<&str>) -> Result<ResourceListResult, ResourceError> {
let state = decode_cursor(cursor)?;
let databases = self
.store
.list_indexed_databases()
.map_err(|e| ResourceError::Internal(e.to_string()))?;
let mut resources = Vec::new();
let mut db_index = state.0;
let mut offset = state.1;
while db_index < databases.len() {
let (connection_id, database) = &databases[db_index];
let all = self.list_for_database(connection_id, database)?;
let take = PAGE_SIZE.saturating_sub(resources.len());
let end = (offset + take).min(all.len());
let page = &all[offset.min(all.len())..end];
resources.extend_from_slice(page);
if end < all.len() {
return Ok(ResourceListResult {
resources,
next_cursor: Some(encode_cursor(&CursorState {
db: db_index,
offset: end,
})),
});
}
if resources.len() >= PAGE_SIZE && db_index + 1 < databases.len() {
return Ok(ResourceListResult {
resources,
next_cursor: Some(encode_cursor(&CursorState {
db: db_index + 1,
offset: 0,
})),
});
}
db_index += 1;
offset = 0;
}
Ok(ResourceListResult {
resources,
next_cursor: None,
})
}
pub fn read(&self, uri: &str) -> Result<ResourceReadResult, ResourceError> {
let parsed = parse_uri(uri)?;
let base = self.store.base_dir(&parsed.connection_id, &parsed.database);
let manifest = self
.store
.read_manifest(&base)
.map_err(|e| ResourceError::Internal(e.to_string()))?
.ok_or_else(|| {
ResourceError::NotFound(format!(
"Resource not found: no index for {}/{}",
parsed.connection_id, parsed.database
))
})?;
let payload: Value = match &parsed.kind {
ResourceKind::Manifest => serde_json::to_value(&manifest)
.map_err(|e| ResourceError::Internal(e.to_string()))?,
ResourceKind::JoinGraph => {
let graph = self
.store
.read_join_graph(&base, &manifest)
.map_err(|e| ResourceError::Internal(e.to_string()))?
.ok_or_else(|| {
ResourceError::NotFound(format!(
"Resource not found: join graph missing for {}",
parsed.database
))
})?;
let overrides = self
.store
.read_overrides(&base)
.map_err(|e| ResourceError::Internal(e.to_string()))?;
let merged = merge_join_graph(graph, overrides.as_ref());
serde_json::to_value(&merged).map_err(|e| ResourceError::Internal(e.to_string()))?
}
ResourceKind::Object { schema, name } => {
let entry = self
.store
.get_object_entry(&base, &manifest, schema, name)
.map_err(|e| ResourceError::Internal(e.to_string()))?;
match entry {
Some(e) if e.excluded != Some(true) => serde_json::to_value(&e)
.map_err(|e| ResourceError::Internal(e.to_string()))?,
_ => {
return Err(ResourceError::NotFound(format!(
"Resource not found: {uri}"
)));
}
}
}
};
let text = serde_json::to_string_pretty(&payload)
.map_err(|e| ResourceError::Internal(e.to_string()))?;
Ok(ResourceReadResult {
contents: vec![ResourceContents {
uri: uri.to_owned(),
mime_type: "application/json".into(),
text,
}],
})
}
pub fn list_templates(&self) -> Vec<ResourceTemplate> {
vec![ResourceTemplate {
uri_template: "nexql://{connectionId}/{database}/object/{schema}/{name}".into(),
name: "Database object".into(),
description:
"Structural card for an indexed table, view, materialized view, or function \
(columns, keys, indexes, definition). Use the search_schema tool to discover refs."
.into(),
mime_type: "application/json".into(),
}]
}
fn list_for_database(
&self,
connection_id: &str,
database: &str,
) -> Result<Vec<McpResource>, ResourceError> {
let base = self.store.base_dir(connection_id, database);
let Some(manifest) = self
.store
.read_manifest(&base)
.map_err(|e| ResourceError::Internal(e.to_string()))?
else {
return Ok(Vec::new());
};
let overrides = self
.store
.read_overrides(&base)
.map_err(|e| ResourceError::Internal(e.to_string()))?;
let prefix = format!(
"nexql://{}/{}",
encode_uri_component(connection_id),
encode_uri_component(database)
);
let mut resources = vec![
McpResource {
uri: format!("{prefix}/manifest"),
name: format!("{database} index manifest"),
description: Some(format!(
"Index metadata for {database} (schemas, counts, fingerprint, build time)."
)),
mime_type: "application/json".into(),
},
McpResource {
uri: format!("{prefix}/joingraph"),
name: format!("{database} join graph"),
description: Some(format!(
"Declared and inferred foreign-key relationships between tables in {database}."
)),
mime_type: "application/json".into(),
},
];
let mut objects: Vec<(String, ObjectEntry)> = Vec::new();
for shard in &manifest.shards {
let Some(entries) = self
.store
.read_shard_entries(&base, &shard.file)
.map_err(|e| ResourceError::Internal(e.to_string()))?
else {
continue;
};
for (ref_, entry) in entries {
if entry.excluded == Some(true) || object_excluded(overrides.as_ref(), &ref_) {
continue;
}
objects.push((ref_, entry));
}
}
objects.sort_by(|a, b| a.0.cmp(&b.0));
for (ref_, entry) in objects {
let (schema, name) = split_ref(&ref_);
let description = match &entry.comment {
Some(c) if !c.is_empty() => format!("{} — {c}", entry.kind.as_str()),
_ => entry.kind.as_str().to_owned(),
};
resources.push(McpResource {
uri: format!(
"{prefix}/object/{}/{}",
encode_uri_component(&schema),
encode_uri_component(&name)
),
name: ref_,
description: Some(description),
mime_type: "application/json".into(),
});
}
Ok(resources)
}
}
fn object_excluded(overrides: Option<&IndexOverrides>, ref_: &str) -> bool {
overrides
.and_then(|o| o.objects.as_ref())
.and_then(|objs| objs.get(ref_))
.and_then(|obj| obj.excluded)
== Some(true)
}
fn split_ref(ref_: &str) -> (String, String) {
match ref_.split_once('.') {
Some((schema, name)) => (schema.to_owned(), name.to_owned()),
None => ("public".to_owned(), ref_.to_owned()),
}
}
fn merge_join_graph(mut graph: JoinGraph, overrides: Option<&IndexOverrides>) -> JoinGraph {
let Some(joins) = overrides.and_then(|o| o.joins.as_ref()) else {
return graph;
};
if joins.is_empty() {
return graph;
}
let mut override_map: HashMap<String, JoinEdge> = HashMap::new();
for edge in joins {
let key = format!("{}->{}:{}", edge.from, edge.to, edge.via);
override_map.insert(key, edge.clone());
}
let mut merged = Vec::new();
for base in graph.edges.drain(..) {
let key = format!("{}->{}:{}", base.from, base.to, base.via);
if let Some(over) = override_map.remove(&key) {
if over.disabled != Some(true) {
merged.push(over);
}
} else {
merged.push(base);
}
}
for edge in override_map.into_values() {
if edge.disabled != Some(true) {
merged.push(edge);
}
}
graph.edges = merged;
graph
}
pub fn encode_cursor_state(db: usize, offset: usize) -> String {
encode_cursor_inner(&CursorState { db, offset })
}
pub fn decode_cursor(cursor: Option<&str>) -> Result<(usize, usize), ResourceError> {
let state = decode_cursor_inner(cursor)?;
Ok((state.db, state.offset))
}
fn encode_cursor_inner(state: &CursorState) -> String {
let json = serde_json::to_string(state).unwrap_or_else(|_| "{}".into());
B64.encode(json.as_bytes())
}
fn decode_cursor_inner(cursor: Option<&str>) -> Result<CursorState, ResourceError> {
let Some(cursor) = cursor.filter(|c| !c.is_empty()) else {
return Ok(CursorState { db: 0, offset: 0 });
};
let bytes = B64
.decode(cursor.as_bytes())
.map_err(|_| ResourceError::InvalidParams("Invalid cursor".into()))?;
let state: CursorState = serde_json::from_slice(&bytes)
.map_err(|_| ResourceError::InvalidParams("Invalid cursor".into()))?;
Ok(state)
}
fn encode_cursor(state: &CursorState) -> String {
encode_cursor_inner(state)
}
pub fn parse_uri(uri: &str) -> Result<ParsedUri, ResourceError> {
let rest = uri
.strip_prefix("nexql://")
.ok_or_else(|| ResourceError::NotFound(format!("Resource not found: {uri}")))?;
let mut parts = rest.splitn(3, '/');
let connection_id = parts.next().filter(|s| !s.is_empty());
let database = parts.next().filter(|s| !s.is_empty());
let tail = parts.next().filter(|s| !s.is_empty());
let (Some(conn_raw), Some(db_raw), Some(tail)) = (connection_id, database, tail) else {
return Err(ResourceError::NotFound(format!(
"Resource not found: {uri}"
)));
};
let connection_id = decode_uri_component(conn_raw)
.ok_or_else(|| ResourceError::NotFound(format!("Resource not found: {uri}")))?;
let database = decode_uri_component(db_raw)
.ok_or_else(|| ResourceError::NotFound(format!("Resource not found: {uri}")))?;
if tail == "manifest" {
return Ok(ParsedUri {
connection_id,
database,
kind: ResourceKind::Manifest,
});
}
if tail == "joingraph" {
return Ok(ParsedUri {
connection_id,
database,
kind: ResourceKind::JoinGraph,
});
}
let mut obj_parts = tail.splitn(3, '/');
let kind = obj_parts.next();
let schema_raw = obj_parts.next();
let name_raw = obj_parts.next();
if kind != Some("object") || schema_raw.is_none() || name_raw.is_none() {
return Err(ResourceError::NotFound(format!(
"Resource not found: {uri}"
)));
}
let schema = decode_uri_component(schema_raw.unwrap())
.ok_or_else(|| ResourceError::NotFound(format!("Resource not found: {uri}")))?;
let name = decode_uri_component(name_raw.unwrap())
.ok_or_else(|| ResourceError::NotFound(format!("Resource not found: {uri}")))?;
if schema.is_empty() || name.is_empty() || name.contains('/') {
return Err(ResourceError::NotFound(format!(
"Resource not found: {uri}"
)));
}
Ok(ParsedUri {
connection_id,
database,
kind: ResourceKind::Object { schema, name },
})
}
pub fn encode_uri_component(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for b in s.bytes() {
match b {
b'A'..=b'Z'
| b'a'..=b'z'
| b'0'..=b'9'
| b'-'
| b'_'
| b'.'
| b'!'
| b'~'
| b'*'
| b'\''
| b'('
| b')' => out.push(b as char),
_ => out.push_str(&format!("%{b:02X}")),
}
}
out
}
pub fn decode_uri_component(s: &str) -> Option<String> {
let bytes = s.as_bytes();
let mut out = Vec::with_capacity(bytes.len());
let mut i = 0;
while i < bytes.len() {
match bytes[i] {
b'%' => {
if i + 2 >= bytes.len() {
return None;
}
let hi = from_hex(bytes[i + 1])?;
let lo = from_hex(bytes[i + 2])?;
out.push((hi << 4) | lo);
i += 3;
}
b'+' => {
out.push(b' ');
i += 1;
}
c => {
out.push(c);
i += 1;
}
}
}
String::from_utf8(out).ok()
}
fn from_hex(b: u8) -> Option<u8> {
match b {
b'0'..=b'9' => Some(b - b'0'),
b'a'..=b'f' => Some(b - b'a' + 10),
b'A'..=b'F' => Some(b - b'A' + 10),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
use nexql_index::{
BuildDepth, BuildMode, ColumnEntry, DbObjectKind, IndexCounts, IndexDerived, IndexManifest,
IndexScope, IndexStats, ObjectShard,
};
use tempfile::TempDir;
#[test]
fn cursor_round_trip() {
let encoded = encode_cursor_state(2, 150);
let (db, offset) = decode_cursor(Some(&encoded)).unwrap();
assert_eq!((db, offset), (2, 150));
}
#[test]
fn cursor_default_when_absent() {
assert_eq!(decode_cursor(None).unwrap(), (0, 0));
assert_eq!(decode_cursor(Some("")).unwrap(), (0, 0));
}
#[test]
fn invalid_cursor_is_invalid_params() {
let err = decode_cursor(Some("!!!not-base64-json!!!")).unwrap_err();
assert_eq!(err.code(), ERR_INVALID_PARAMS);
assert!(err.to_string().contains("Invalid cursor"));
}
#[test]
fn parse_manifest_joingraph_object_uris() {
let m = parse_uri("nexql://prod/appdb/manifest").unwrap();
assert_eq!(m.connection_id, "prod");
assert_eq!(m.database, "appdb");
assert_eq!(m.kind, ResourceKind::Manifest);
let j = parse_uri("nexql://prod/appdb/joingraph").unwrap();
assert_eq!(j.kind, ResourceKind::JoinGraph);
let o = parse_uri("nexql://prod/appdb/object/public/users").unwrap();
assert_eq!(
o.kind,
ResourceKind::Object {
schema: "public".into(),
name: "users".into()
}
);
}
#[test]
fn parse_uri_decodes_components() {
let o = parse_uri("nexql://my%20conn/db%2F1/object/public/my%20table").unwrap();
assert_eq!(o.connection_id, "my conn");
assert_eq!(o.database, "db/1");
assert_eq!(
o.kind,
ResourceKind::Object {
schema: "public".into(),
name: "my table".into()
}
);
}
#[test]
fn unknown_uri_is_not_found() {
let err = parse_uri("nexql://a/b/c").unwrap_err();
assert_eq!(err.code(), ERR_RESOURCE_NOT_FOUND);
let err = parse_uri("http://x").unwrap_err();
assert_eq!(err.code(), ERR_RESOURCE_NOT_FOUND);
}
#[test]
fn empty_index_lists_nothing() {
let tmp = TempDir::new().unwrap();
let provider = ResourceProvider::new(IndexStore::new(tmp.path()));
let result = provider.list(None).unwrap();
assert!(result.resources.is_empty());
assert!(result.next_cursor.is_none());
}
#[test]
fn list_and_read_manifest_from_fixture() {
let tmp = TempDir::new().unwrap();
let store = IndexStore::new(tmp.path());
let base = store.base_dir("conn1", "db1");
let manifest = IndexManifest {
format_version: 1,
connection_id: "conn1".into(),
database: "db1".into(),
indexed_at: "2026-01-01T00:00:00Z".into(),
build_mode: BuildMode::Guided,
build_depth: BuildDepth::Structure,
schema_fingerprint: "fp".into(),
pg_version: "16".into(),
environment: "development".into(),
scope: IndexScope {
included_schemas: vec!["public".into()],
excluded_objects: vec![],
pii_excluded_columns: vec![],
},
counts: IndexCounts {
tables: 1,
views: 0,
functions: 0,
enums: 0,
},
shards: vec![ObjectShard {
file: "objects_public.json".into(),
schema: "public".into(),
objects: 1,
bytes: 10,
hash: "h".into(),
}],
derived: IndexDerived {
tokens: "tokens.json".into(),
join_graph: "joingraph.json".into(),
values: None,
embeddings: None,
embeddings_meta: None,
},
stats: IndexStats {
build_ms: 1,
queries_run: 1,
warnings: vec![],
},
};
store.write_manifest(&base, &manifest).unwrap();
let mut entries = HashMap::new();
entries.insert(
"public.users".into(),
ObjectEntry {
kind: DbObjectKind::Table,
oid: 1,
object_hash: "x".into(),
comment: Some("users table".into()),
row_estimate: 10.0,
size_bytes: 8192,
columns: vec![ColumnEntry {
name: "id".into(),
type_name: "integer".into(),
not_null: true,
default_value: None,
comment: None,
ordinal: 1,
is_pk: Some(true),
profile: None,
pii: None,
}],
primary_key: Some(vec!["id".into()]),
foreign_keys: None,
indexes: None,
checks: None,
excluded: None,
definition: None,
signature: None,
language: None,
volatility: None,
body: None,
values: None,
base_type: None,
constraint: None,
},
);
store
.write_shard_entries(&base, "objects_public.json", &entries)
.unwrap();
let provider = ResourceProvider::new(store);
let listed = provider.list(None).unwrap();
assert_eq!(listed.resources.len(), 3);
assert!(
listed
.resources
.iter()
.any(|r| r.uri.ends_with("/manifest"))
);
assert!(
listed
.resources
.iter()
.any(|r| r.uri.ends_with("/object/public/users"))
);
let read = provider
.read("nexql://conn1/db1/object/public/users")
.unwrap();
assert_eq!(read.contents.len(), 1);
assert!(read.contents[0].text.contains("users table"));
}
}