use dbmcp_server::pagination::Pager;
use super::prelude::*;
use crate::types::{ListEntriesResponse, ListTablesRequest};
const NAME: &str = "listTables";
const TITLE: &str = "List Tables";
const DESCRIPTION: &str = include_str!("../../assets/tools/list_tables.md");
pub(crate) struct ListTablesTool;
impl ToolBase for ListTablesTool {
type Parameter = ListTablesRequest;
type Output = ListEntriesResponse;
type Error = ErrorData;
fn name() -> Cow<'static, str> {
NAME.into()
}
fn title() -> Option<String> {
Some(TITLE.into())
}
fn description() -> Option<Cow<'static, str>> {
Some(DESCRIPTION.into())
}
fn annotations() -> Option<ToolAnnotations> {
Some(
ToolAnnotations::new()
.read_only(true)
.destructive(false)
.idempotent(true)
.open_world(false),
)
}
fn input_schema() -> Option<Arc<JsonObject>> {
Some(input_schema::<Self::Parameter>(true))
}
fn output_schema() -> Option<Arc<JsonObject>> {
Some(output_schema::<Self::Output>())
}
}
impl AsyncTool<SqliteHandler> for ListTablesTool {
async fn invoke(handler: &SqliteHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
handler.list_tables(params).await
}
}
const BRIEF_SQL: &str = r"
SELECT tl.name
FROM pragma_table_list tl
WHERE tl.schema = 'main'
AND tl.type IN ('table', 'virtual')
AND tl.name NOT LIKE 'sqlite_%'
AND (?1 IS NULL OR tl.name LIKE '%' || ?1 || '%')
ORDER BY tl.name
LIMIT ?2 OFFSET ?3";
const DETAILED_SQL: &str = r#"
WITH table_info AS (
SELECT
tl.name AS table_name,
CASE tl.type
WHEN 'virtual' THEN 'VIRTUAL_TABLE'
ELSE 'TABLE'
END AS kind
FROM pragma_table_list tl
WHERE tl.schema = 'main'
AND tl.type IN ('table', 'virtual')
AND tl.name NOT LIKE 'sqlite_%'
AND (?1 IS NULL OR tl.name LIKE '%' || ?1 || '%')
ORDER BY tl.name
LIMIT ?2 OFFSET ?3
),
columns_info AS (
SELECT
ti.table_name,
c.cid AS cid,
json_object(
'name', c.name,
'dataType', c.type,
'ordinalPosition', c.cid + 1,
'nullable', json(CASE WHEN c."notnull" = 0 THEN 'true' ELSE 'false' END),
'default', c.dflt_value,
'comment', NULL
) AS column_json
FROM table_info ti, pragma_table_info(ti.table_name) c
),
pk_constraints AS (
SELECT
ti.table_name,
json_object(
'name', 'PRIMARY',
'type', 'PRIMARY KEY',
'columns', json_group_array(c.name),
'definition', 'PRIMARY KEY (' || group_concat('"' || c.name || '"', ', ') || ')'
) AS constraint_json
FROM table_info ti, pragma_table_info(ti.table_name) c
WHERE c.pk > 0
GROUP BY ti.table_name
HAVING COUNT(c.name) > 0
),
fk_constraints AS (
SELECT
ti.table_name,
json_object(
'name', 'fk_' || ti.table_name || '_' || f.id,
'type', 'FOREIGN KEY',
'columns', json_group_array(f."from"),
'definition', 'FOREIGN KEY (' || group_concat('"' || f."from" || '"', ', ')
|| ') REFERENCES "' || f."table" || '"('
|| group_concat('"' || f."to" || '"', ', ') || ')',
'referencedTable', f."table",
'referencedColumns', json_group_array(f."to")
) AS constraint_json
FROM table_info ti, pragma_foreign_key_list(ti.table_name) f
GROUP BY ti.table_name, f.id
),
unique_constraints AS (
SELECT
ti.table_name,
json_object(
'name', i.name,
'type', 'UNIQUE',
'columns', (SELECT json_group_array(ii.name)
FROM pragma_index_info(i.name) ii),
'definition', 'UNIQUE ('
|| (SELECT group_concat('"' || ii.name || '"', ', ')
FROM pragma_index_info(i.name) ii)
|| ')'
) AS constraint_json
FROM table_info ti, pragma_index_list(ti.table_name) i
WHERE i."unique" = 1 AND i.origin <> 'pk'
),
all_constraints AS (
SELECT table_name, constraint_json FROM pk_constraints
UNION ALL
SELECT table_name, constraint_json FROM fk_constraints
UNION ALL
SELECT table_name, constraint_json FROM unique_constraints
),
indexes_info AS (
SELECT
ti.table_name,
json_object(
'name', il.name,
'columns', (SELECT json_group_array(ii.name) FROM pragma_index_info(il.name) ii),
'unique', json(CASE WHEN il."unique" = 1 THEN 'true' ELSE 'false' END),
'primary', json(CASE WHEN il.origin = 'pk' THEN 'true' ELSE 'false' END),
'method', 'btree',
'definition', COALESCE(
(SELECT m2.sql FROM sqlite_master m2 WHERE m2.type = 'index' AND m2.name = il.name),
'CREATE ' || CASE il."unique" WHEN 1 THEN 'UNIQUE INDEX ' ELSE 'INDEX ' END
|| '"' || il.name || '" ON "' || ti.table_name || '"('
|| (SELECT group_concat('"' || ii.name || '"', ', ') FROM pragma_index_info(il.name) ii)
|| ')'
)
) AS index_json
FROM table_info ti, pragma_index_list(ti.table_name) il
),
triggers_info AS (
SELECT
m.tbl_name AS table_name,
json_object(
'name', m.name,
'definition', m.sql,
'enabled', json('true')
) AS trigger_json
FROM sqlite_master m
JOIN table_info ti ON ti.table_name = m.tbl_name
WHERE m.type = 'trigger'
)
SELECT
ti.table_name AS name,
json_object(
'schema', 'main',
'kind', ti.kind,
'owner', NULL,
'comment', NULL,
'columns', COALESCE(
(SELECT json_group_array(json(ci.column_json))
FROM (SELECT column_json
FROM columns_info
WHERE table_name = ti.table_name
ORDER BY cid) ci),
json('[]')),
'constraints', COALESCE(
(SELECT json_group_array(json(ac.constraint_json))
FROM all_constraints ac
WHERE ac.table_name = ti.table_name),
json('[]')),
'indexes', COALESCE(
(SELECT json_group_array(json(ii.index_json))
FROM indexes_info ii
WHERE ii.table_name = ti.table_name),
json('[]')),
'triggers', COALESCE(
(SELECT json_group_array(json(tg.trigger_json))
FROM triggers_info tg
WHERE tg.table_name = ti.table_name),
json('[]'))
) AS entry
FROM table_info ti
ORDER BY ti.table_name"#;
impl SqliteHandler {
pub async fn list_tables(
&self,
ListTablesRequest {
cursor,
search,
detailed,
}: ListTablesRequest,
) -> Result<ListEntriesResponse, ErrorData> {
let pattern = search.as_deref().map(str::trim).filter(|s| !s.is_empty());
let pager = Pager::new(cursor, self.config.page_size);
if detailed {
let rows: Vec<(String, sqlx::types::Json<serde_json::Value>)> = self
.connection
.fetch(
sqlx::query(DETAILED_SQL)
.bind(pattern)
.bind(pager.limit())
.bind(pager.offset()),
None,
)
.await?;
let (rows, next_cursor) = pager.paginate(rows);
return Ok(ListEntriesResponse::detailed(
rows.into_iter().map(|(name, json)| (name, json.0)).collect(),
next_cursor,
));
}
let rows: Vec<String> = self
.connection
.fetch_scalar(
sqlx::query(BRIEF_SQL)
.bind(pattern)
.bind(pager.limit())
.bind(pager.offset()),
None,
)
.await?;
let (tables, next_cursor) = pager.paginate(rows);
Ok(ListEntriesResponse::brief(tables, next_cursor))
}
}