feat: add schema exploration tools (keyspaces, tables, DDL)

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;

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<string, string>;
+  vttablet_options: Record<string, string>;
+}
+
+interface PaginatedResponse<T> {
+  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<PaginatedResponse<Keyspace>> {
+  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<PaginatedResponse<Keyspace>>(
+    `/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.
  */

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<Keyspace> {
+  const filtered: Record<string, unknown> = {};
+  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<Keyspace>;
+}
+
+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<string, string | undefined>)
+          : 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");
+    }
+  },
+});

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<string, unknown> }, input: { organization: string; database: string; branch: string; keyspace?: string }) {
+  const env =
+    Object.keys(ctx.env).length > 0
+      ? (ctx.env as Record<string, string | undefined>)
+      : 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");
+    }
+  },
+});

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;