use rmcp::model::{Content, IntoContents};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use super::ToolError;
use crate::{mcp::McpServerSqlite, traits::SqliteServerTool};
#[derive(
Clone,
Copy,
Debug,
PartialEq,
Eq,
PartialOrd,
Ord,
Hash,
Default,
Serialize,
Deserialize,
JsonSchema,
)]
pub struct ListForeignKeysTool;
impl SqliteServerTool for ListForeignKeysTool {
const NAME: &str = "list_foreign_keys";
type Context = McpServerSqlite;
type Error = ToolError<ListForeignKeysError>;
type Input = ListForeignKeysInput;
type Output = ListForeignKeysOutput;
fn handle(
ctx: &Self::Context,
input: Self::Input,
) -> Result<Self::Output, Self::Error> {
let conn = ctx
.connection()
.map_err(|source| ToolError::Connection { source })?;
let pragma_sql = format!(
"PRAGMA foreign_key_list({})",
enquote_identifier(&input.table_name),
);
let mut stmt = conn.prepare(&pragma_sql).map_err(|source| {
ToolError::Tool(ListForeignKeysError::Query { source })
})?;
let foreign_keys = stmt
.query_map([], |row| {
Ok(ForeignKeyInfo {
id: row.get(0)?,
from_column: row.get(3)?,
to_table: row.get(2)?,
to_column: row.get(4)?,
on_update: row.get(5)?,
on_delete: row.get(6)?,
})
})
.map_err(|source| {
ToolError::Tool(ListForeignKeysError::Query { source })
})?
.collect::<Result<Vec<_>, _>>()
.map_err(|source| {
ToolError::Tool(ListForeignKeysError::Query { source })
})?;
Ok(ListForeignKeysOutput { foreign_keys })
}
}
fn enquote_identifier(name: &str) -> String {
format!("\"{}\"", name.replace('"', "\"\""))
}
#[derive(
Clone,
Debug,
PartialEq,
Eq,
PartialOrd,
Ord,
Hash,
Serialize,
Deserialize,
schemars::JsonSchema,
)]
pub struct ListForeignKeysInput {
#[schemars(description = "The name of the table to list foreign keys for")]
pub table_name: String,
}
#[derive(
Clone,
Debug,
PartialEq,
Eq,
PartialOrd,
Ord,
Hash,
Serialize,
Deserialize,
schemars::JsonSchema,
)]
pub struct ListForeignKeysOutput {
pub foreign_keys: Vec<ForeignKeyInfo>,
}
#[derive(
Clone,
Debug,
PartialEq,
Eq,
PartialOrd,
Ord,
Hash,
Serialize,
Deserialize,
schemars::JsonSchema,
)]
pub struct ForeignKeyInfo {
pub id: i64,
pub from_column: String,
pub to_table: String,
pub to_column: String,
pub on_update: String,
pub on_delete: String,
}
#[derive(Debug, thiserror::Error)]
pub enum ListForeignKeysError {
#[error("failed to list foreign keys: {source}")]
Query {
source: rusqlite::Error,
},
}
impl IntoContents for ListForeignKeysError {
fn into_contents(self) -> Vec<Content> {
vec![Content::text(self.to_string())]
}
}