diff --git a/src/gram.ts b/src/gram.ts index 9459ebf..1d084fb 100644 --- a/src/gram.ts +++ b/src/gram.ts @@ -6,6 +6,8 @@ import { executeWriteQueryGram } from "./tools/execute-write-query.ts"; import { getInsightsGram } from "./tools/get-insights.ts"; import { listClusterSizesGram } from "./tools/list-cluster-sizes.ts"; import { searchDocumentationGram } from "./tools/search-documentation.ts"; +import { getBranchKeyspacesGram } from "./tools/get-branch-keyspaces.ts"; +import { getBranchTablesGram, getTableSchemaGram } from "./tools/get-branch-schema.ts"; const gram = new Gram({ envSchema: { @@ -25,6 +27,9 @@ const gram = new Gram({ .extend(executeWriteQueryGram) .extend(getInsightsGram) .extend(listClusterSizesGram) - .extend(searchDocumentationGram); + .extend(searchDocumentationGram) + .extend(getBranchKeyspacesGram) + .extend(getBranchTablesGram) + .extend(getTableSchemaGram); export default gram; diff --git a/src/lib/planetscale-api.ts b/src/lib/planetscale-api.ts index b5957fa..042a1de 100644 --- a/src/lib/planetscale-api.ts +++ b/src/lib/planetscale-api.ts @@ -271,6 +271,94 @@ export async function deletePostgresRole( ); } +export interface Keyspace { + id: string; + name: string; + shards: number; + sharded: boolean; + replicas: number; + extra_replicas: number; + created_at: string; + updated_at: string; + cluster_name: string; + cluster_display_name: string; + resizing: boolean; + resize_pending: boolean; + config_change_in_progress: boolean; + ready: boolean; + metal: boolean; + default: boolean; + imported: boolean; + vector_pool_allocation: number | null; + node_ttl_strategy: string; + replication_durability_constraints: { strategy: string }; + vreplication_flags: { + optimize_inserts: boolean; + allow_no_blob_binlog_row_image: boolean; + vplayer_batching: boolean; + }; + mysqld_options: Record; + vttablet_options: Record; +} + +interface PaginatedResponse { + type: string; + current_page: number; + next_page: number | null; + next_page_url: string | null; + prev_page: number | null; + prev_page_url: string | null; + data: T[]; +} + +/** + * List keyspaces for a database branch + */ +export async function listKeyspaces( + organization: string, + database: string, + branch: string, + authHeader: string, + options?: { page?: number; perPage?: number } +): Promise> { + const params = new URLSearchParams(); + if (options?.page) params.set("page", String(options.page)); + if (options?.perPage) params.set("per_page", String(options.perPage)); + const query = params.toString(); + + return apiRequest>( + `/organizations/${encodeURIComponent(organization)}/databases/${encodeURIComponent(database)}/branches/${encodeURIComponent(branch)}/keyspaces${query ? `?${query}` : ""}`, + authHeader + ); +} + +export interface SchemaTable { + name: string; + html: string; + raw: string; + annotated: boolean; +} + +/** + * Get the schema for a database branch, optionally filtered to a single keyspace. + */ +export async function getBranchSchema( + organization: string, + database: string, + branch: string, + authHeader: string, + options?: { keyspace?: string } +): Promise<{ data: SchemaTable[] }> { + const params = new URLSearchParams(); + if (options?.keyspace) params.set("keyspace", options.keyspace); + const query = params.toString(); + + return apiRequest<{ data: SchemaTable[] }>( + `/organizations/${encodeURIComponent(organization)}/databases/${encodeURIComponent(database)}/branches/${encodeURIComponent(branch)}/schema${query ? `?${query}` : ""}`, + authHeader + ); +} + /** * Delete Vitess password credentials. */ diff --git a/src/tools/get-branch-keyspaces.ts b/src/tools/get-branch-keyspaces.ts new file mode 100644 index 0000000..967d871 --- /dev/null +++ b/src/tools/get-branch-keyspaces.ts @@ -0,0 +1,95 @@ +import { Gram } from "@gram-ai/functions"; +import { z } from "zod"; +import { PlanetScaleAPIError, listKeyspaces } from "../lib/planetscale-api.ts"; +import type { Keyspace } from "../lib/planetscale-api.ts"; +import { getAuthToken, getAuthHeader } from "../lib/auth.ts"; + +// Fields to strip from keyspace responses for token efficiency +const STRIP_FIELDS = new Set([ + "type", + "cluster_rate_name", + "cluster_rate_display_name", +]); + +/** + * Filter a keyspace entry to remove redundant fields and null values + */ +function filterKeyspace(keyspace: Keyspace): Partial { + const filtered: Record = {}; + for (const [key, value] of Object.entries(keyspace)) { + if (STRIP_FIELDS.has(key)) continue; + if (value === null) continue; + filtered[key] = value; + } + return filtered as Partial; +} + +export const getBranchKeyspacesGram = new Gram().tool({ + name: "get_branch_keyspaces", + description: + "List keyspaces for a PlanetScale database branch. Returns keyspace configuration including shard count, cluster size, replica count, replication durability settings, and MySQL/VTTablet options. Useful for understanding the topology and sizing of a database.", + inputSchema: { + organization: z.string().describe("PlanetScale organization name"), + database: z.string().describe("Database name"), + branch: z.string().describe("Branch name (e.g., 'main')"), + page: z + .number() + .optional() + .describe("Page number for pagination (default: 1)"), + per_page: z + .number() + .optional() + .describe("Results per page (default: 25)"), + }, + async execute(ctx, input) { + try { + const env = + Object.keys(ctx.env).length > 0 + ? (ctx.env as Record) + : process.env; + + const auth = getAuthToken(env); + if (!auth) { + return ctx.text("Error: No PlanetScale authentication configured."); + } + + const { organization, database, branch } = input; + if (!organization || !database || !branch) { + return ctx.text( + "Error: organization, database, and branch are required" + ); + } + + const authHeader = getAuthHeader(env); + const response = await listKeyspaces( + organization, + database, + branch, + authHeader, + { page: input.page, perPage: input.per_page } + ); + + const keyspaces = response.data.map(filterKeyspace); + + return ctx.json({ + organization, + database, + branch, + total_keyspaces: keyspaces.length, + current_page: response.current_page, + next_page: response.next_page, + keyspaces, + }); + } catch (error) { + if (error instanceof PlanetScaleAPIError) { + return ctx.text( + `Error: ${error.message} (status: ${error.statusCode})` + ); + } + if (error instanceof Error) { + return ctx.text(`Error: ${error.message}`); + } + return ctx.text("Error: An unexpected error occurred"); + } + }, +}); diff --git a/src/tools/get-branch-schema.ts b/src/tools/get-branch-schema.ts new file mode 100644 index 0000000..da1ea1b --- /dev/null +++ b/src/tools/get-branch-schema.ts @@ -0,0 +1,117 @@ +import { Gram } from "@gram-ai/functions"; +import { z } from "zod"; +import { + PlanetScaleAPIError, + getBranchSchema, +} from "../lib/planetscale-api.ts"; +import type { SchemaTable } from "../lib/planetscale-api.ts"; +import { getAuthToken, getAuthHeader } from "../lib/auth.ts"; + +function filterTable(table: SchemaTable): { name: string; raw: string; annotated: boolean } { + return { name: table.name, raw: table.raw, annotated: table.annotated }; +} + +function fetchSchema(ctx: { env: Record }, input: { organization: string; database: string; branch: string; keyspace?: string }) { + const env = + Object.keys(ctx.env).length > 0 + ? (ctx.env as Record) + : process.env; + + const auth = getAuthToken(env); + if (!auth) { + throw new Error("No PlanetScale authentication configured."); + } + + const authHeader = getAuthHeader(env); + return getBranchSchema( + input.organization, + input.database, + input.branch, + authHeader, + { keyspace: input.keyspace } + ); +} + +const branchInputSchema = { + organization: z.string().describe("PlanetScale organization name"), + database: z.string().describe("Database name"), + branch: z.string().describe("Branch name (e.g., 'main')"), + keyspace: z + .string() + .optional() + .describe( + "Vitess keyspace to filter by. When omitted, only tables in the default keyspace are returned — tables in other keyspaces will not be visible. Use get_branch_keyspaces to discover available keyspaces." + ), +}; + +export const getBranchTablesGram = new Gram().tool({ + name: "get_branch_tables", + description: + "List table names for a PlanetScale database branch. Returns only the table names without DDL — use get_table_schema to fetch the full CREATE TABLE statement for a specific table. Use get_branch_keyspaces first to discover available keyspaces.", + inputSchema: branchInputSchema, + async execute(ctx, input) { + try { + const response = await fetchSchema(ctx, input); + + return ctx.json({ + organization: input.organization, + database: input.database, + branch: input.branch, + keyspace: input.keyspace ?? null, + total_tables: response.data.length, + tables: response.data.map((t) => t.name), + }); + } catch (error) { + if (error instanceof PlanetScaleAPIError) { + return ctx.text( + `Error: ${error.message} (status: ${error.statusCode})` + ); + } + if (error instanceof Error) { + return ctx.text(`Error: ${error.message}`); + } + return ctx.text("Error: An unexpected error occurred"); + } + }, +}); + +export const getTableSchemaGram = new Gram().tool({ + name: "get_table_schema", + description: + "Get the CREATE TABLE DDL for a specific table on a PlanetScale database branch. Use get_branch_tables first to discover available table names.", + inputSchema: { + ...branchInputSchema, + table: z.string().describe("Table name to get the schema for"), + }, + async execute(ctx, input) { + try { + const response = await fetchSchema(ctx, input); + + const match = response.data.find((t) => t.name === input.table); + if (!match) { + const keyspaceLabel = input.keyspace ? `keyspace '${input.keyspace}'` : "this branch"; + return ctx.text( + `Error: Table '${input.table}' not found in ${keyspaceLabel}. Use get_branch_tables to list available tables.` + ); + } + + return ctx.json({ + organization: input.organization, + database: input.database, + branch: input.branch, + keyspace: input.keyspace ?? null, + table: filterTable(match), + }); + } catch (error) { + if (error instanceof PlanetScaleAPIError) { + return ctx.text( + `Error: ${error.message} (status: ${error.statusCode})` + ); + } + if (error instanceof Error) { + return ctx.text(`Error: ${error.message}`); + } + return ctx.text("Error: An unexpected error occurred"); + } + }, +}); diff --git a/src/tools/get-insights.ts b/src/tools/get-insights.ts index f133d4d..1b21d22 100644 --- a/src/tools/get-insights.ts +++ b/src/tools/get-insights.ts @@ -37,6 +37,7 @@ const RESULT_FIELDS = [ "max_egress_bytes", "max_shard_queries", "tables", + "qualified_tables", "index_usages", "keyspace", "last_run_at", @@ -61,6 +62,7 @@ export interface InsightsEntry { max_egress_bytes?: number; max_shard_queries?: number; tables?: string[]; + qualified_tables?: string[]; index_usages?: unknown[]; keyspace?: string; last_run_at?: string;