google-cloud-mcp/dist/services/spanner.js

Model Context Protocol server for Google Cloud services

/**
 * Google Cloud Spanner integration for MCP
 */
import { ResourceTemplate } from '@modelcontextprotocol/sdk/server/mcp.js';
import { Spanner } from '@google-cloud/spanner';
import { z } from 'zod';
import { getProjectId } from '../utils/auth.js';
import { GcpMcpError } from '../utils/error.js';
/**
 * Initialises the Google Cloud Spanner client
 *
 * @returns A configured Spanner client
 */
function getSpannerClient() {
    return new Spanner({
        projectId: process.env.GOOGLE_CLOUD_PROJECT
    });
}
/**
 * Gets the Spanner instance and database from environment variables or parameters
 *
 * @param instanceId Optional instance ID (defaults to environment variable)
 * @param databaseId Optional database ID (defaults to environment variable)
 * @returns The instance and database IDs
 */
async function getSpannerConfig(instanceId, databaseId) {
    const instance = instanceId || process.env.SPANNER_INSTANCE;
    const database = databaseId || process.env.SPANNER_DATABASE;
    if (!instance) {
        throw new GcpMcpError('Spanner instance ID not provided. Set SPANNER_INSTANCE environment variable or provide instanceId parameter.', 'INVALID_ARGUMENT', 400);
    }
    if (!database) {
        throw new GcpMcpError('Spanner database ID not provided. Set SPANNER_DATABASE environment variable or provide databaseId parameter.', 'INVALID_ARGUMENT', 400);
    }
    return { instanceId: instance, databaseId: database };
}
/**
 * Extracts schema information from Spanner database
 *
 * @param instanceId Spanner instance ID
 * @param databaseId Spanner database ID
 * @returns Schema information for the database
 */
async function getSpannerSchema(instanceId, databaseId) {
    const spanner = getSpannerClient();
    const instance = spanner.instance(instanceId);
    const database = instance.database(databaseId);
    // Query for tables
    const [tablesResult] = await database.run({
        sql: `SELECT t.table_name
          FROM information_schema.tables t
          WHERE t.table_catalog = '' AND t.table_schema = ''
          ORDER BY t.table_name`
    });
    if (!tablesResult || tablesResult.length === 0) {
        return { tables: [] };
    }
    const tables = [];
    // Process each table to get columns
    for (const tableRow of tablesResult) {
        const tableName = tableRow.table_name;
        // Get columns for this table
        const [columnsResult] = await database.run({
            sql: `SELECT column_name, spanner_type, is_nullable
            FROM information_schema.columns
            WHERE table_catalog = '' AND table_schema = '' AND table_name = @tableName
            ORDER BY ordinal_position`,
            params: { tableName }
        });
        const columns = columnsResult.map((col) => ({
            name: col.column_name,
            type: col.spanner_type,
            nullable: col.is_nullable === 'YES'
        }));
        // Get indexes for this table
        const [indexesResult] = await database.run({
            sql: `SELECT i.index_name, ic.column_name, i.is_unique
            FROM information_schema.indexes i
            JOIN information_schema.index_columns ic
              ON i.table_catalog = ic.table_catalog
              AND i.table_schema = ic.table_schema
              AND i.table_name = ic.table_name
              AND i.index_name = ic.index_name
            WHERE i.table_catalog = '' AND i.table_schema = '' AND i.table_name = @tableName
            ORDER BY i.index_name, ic.ordinal_position`,
            params: { tableName }
        });
        // Group index columns by index name
        const indexMap = new Map();
        for (const idx of indexesResult) {
            const indexName = idx.index_name;
            const columnName = idx.column_name;
            const isUnique = idx.is_unique;
            if (!indexMap.has(indexName)) {
                indexMap.set(indexName, { name: indexName, columns: [], unique: isUnique });
            }
            indexMap.get(indexName)?.columns.push(columnName);
        }
        const indexes = Array.from(indexMap.values());
        tables.push({
            name: tableName,
            columns,
            indexes: indexes.length > 0 ? indexes : undefined
        });
    }
    return { tables };
}
/**
 * Formats schema information as markdown
 *
 * @param schema The schema to format
 * @returns A markdown string representation of the schema
 */
function formatSchemaAsMarkdown(schema) {
    if (schema.tables.length === 0) {
        return 'No tables found in the database.';
    }
    let markdown = '# Database Schema\n\n';
    for (const table of schema.tables) {
        markdown += `## Table: ${table.name}\n\n`;
        // Columns
        markdown += '### Columns\n\n';
        markdown += '| Column | Type | Nullable |\n';
        markdown += '|--------|------|----------|\n';
        for (const column of table.columns) {
            markdown += `| ${column.name} | ${column.type} | ${column.nullable ? 'YES' : 'NO'} |\n`;
        }
        markdown += '\n';
        // Indexes
        if (table.indexes && table.indexes.length > 0) {
            markdown += '### Indexes\n\n';
            markdown += '| Name | Columns | Unique |\n';
            markdown += '|------|---------|--------|\n';
            for (const index of table.indexes) {
                markdown += `| ${index.name} | ${index.columns.join(', ')} | ${index.unique ? 'YES' : 'NO'} |\n`;
            }
            markdown += '\n';
        }
        markdown += '---\n\n';
    }
    return markdown;
}
/**
 * Registers Google Cloud Spanner resources with the MCP server
 *
 * @param server The MCP server instance
 */
export function registerSpannerResources(server) {
    // Register a resource for database schema
    server.resource('spanner-schema', new ResourceTemplate('gcp-spanner://{projectId}/{instanceId}/{databaseId}/schema', { list: undefined }), async (uri, { projectId, instanceId, databaseId }, _extra) => {
        try {
            const actualProjectId = projectId || await getProjectId();
            const config = await getSpannerConfig(Array.isArray(instanceId) ? instanceId[0] : instanceId, Array.isArray(databaseId) ? databaseId[0] : databaseId);
            const schema = await getSpannerSchema(config.instanceId, config.databaseId);
            const markdown = formatSchemaAsMarkdown(schema);
            return {
                contents: [{
                        uri: uri.href,
                        text: `# Spanner Database Schema\n\nProject: ${actualProjectId}\nInstance: ${config.instanceId}\nDatabase: ${config.databaseId}\n\n${markdown}`
                    }]
            };
        }
        catch (error) {
            console.error('Error fetching Spanner schema:', error);
            throw error;
        }
    });
    // Register a resource for table data preview
    server.resource('table-preview', new ResourceTemplate('gcp-spanner://{projectId}/{instanceId}/{databaseId}/tables/{tableName}/preview', { list: undefined }), async (uri, { projectId, instanceId, databaseId, tableName }, _extra) => {
        try {
            const actualProjectId = projectId || await getProjectId();
            const config = await getSpannerConfig(Array.isArray(instanceId) ? instanceId[0] : instanceId, Array.isArray(databaseId) ? databaseId[0] : databaseId);
            if (!tableName) {
                throw new GcpMcpError('Table name is required', 'INVALID_ARGUMENT', 400);
            }
            const spanner = getSpannerClient();
            const instance = spanner.instance(config.instanceId);
            const database = instance.database(config.databaseId);
            // Get a preview of the table data (first 10 rows)
            const [result] = await database.run({
                sql: `SELECT * FROM ${tableName} LIMIT 10`
            });
            if (!result || result.length === 0) {
                return {
                    contents: [{
                            uri: uri.href,
                            text: `# Table Preview: ${tableName}\n\nNo data found in the table.`
                        }]
                };
            }
            // Convert to markdown table
            const columns = Object.keys(result[0]);
            let markdown = `# Table Preview: ${tableName}\n\n`;
            // Table header
            markdown += '| ' + columns.join(' | ') + ' |\n';
            markdown += '| ' + columns.map(() => '---').join(' | ') + ' |\n';
            // Table rows
            for (const row of result) {
                const rowValues = columns.map(col => {
                    const value = row[col];
                    if (value === null || value === undefined)
                        return 'NULL';
                    if (typeof value === 'object')
                        return JSON.stringify(value);
                    return String(value);
                });
                markdown += '| ' + rowValues.join(' | ') + ' |\n';
            }
            return {
                contents: [{
                        uri: uri.href,
                        text: `# Table Preview: ${tableName}\n\nProject: ${actualProjectId}\nInstance: ${config.instanceId}\nDatabase: ${config.databaseId}\n\n${markdown}`
                    }]
            };
        }
        catch (error) {
            console.error('Error fetching table preview:', error);
            throw error;
        }
    });
}
/**
 * Registers Google Cloud Spanner tools with the MCP server
 *
 * @param server The MCP server instance
 */
export function registerSpannerTools(server) {
    // Tool to execute SQL queries
    server.tool('execute-spanner-query', {
        sql: z.string().describe('The SQL query to execute'),
        instanceId: z.string().optional().describe('Spanner instance ID (defaults to SPANNER_INSTANCE env var)'),
        databaseId: z.string().optional().describe('Spanner database ID (defaults to SPANNER_DATABASE env var)'),
        params: z.record(z.any()).optional().describe('Query parameters')
    }, async ({ sql, instanceId, databaseId, params }, _extra) => {
        try {
            const projectId = await getProjectId();
            const config = await getSpannerConfig(Array.isArray(instanceId) ? instanceId[0] : instanceId, Array.isArray(databaseId) ? databaseId[0] : databaseId);
            const spanner = getSpannerClient();
            const instance = spanner.instance(config.instanceId);
            const database = instance.database(config.databaseId);
            // Execute the query
            const [result] = await database.run({
                sql,
                params: params || {}
            });
            if (!result || result.length === 0) {
                return {
                    content: [{
                            type: 'text',
                            text: `# Query Results\n\nProject: ${projectId}\nInstance: ${config.instanceId}\nDatabase: ${config.databaseId}\n\nQuery executed successfully. No results returned.`
                        }]
                };
            }
            // Convert to markdown table
            const columns = Object.keys(result[0]);
            let markdown = `# Query Results\n\nProject: ${projectId}\nInstance: ${config.instanceId}\nDatabase: ${config.databaseId}\n\n`;
            markdown += `SQL: \`${sql}\`\n\n`;
            markdown += `Rows: ${result.length}\n\n`;
            // Table header
            markdown += '| ' + columns.join(' | ') + ' |\n';
            markdown += '| ' + columns.map(() => '---').join(' | ') + ' |\n';
            // Table rows (limit to 100 rows for display)
            const displayRows = result.slice(0, 100);
            for (const row of displayRows) {
                const rowValues = columns.map(col => {
                    const value = row[col];
                    if (value === null || value === undefined)
                        return 'NULL';
                    if (typeof value === 'object')
                        return JSON.stringify(value);
                    return String(value);
                });
                markdown += '| ' + rowValues.join(' | ') + ' |\n';
            }
            if (result.length > 100) {
                markdown += '\n*Results truncated. Showing 100 of ' + result.length + ' rows.*';
            }
            return {
                content: [{
                        type: 'text',
                        text: markdown
                    }]
            };
        }
        catch (error) {
            console.error('Error executing Spanner query:', error);
            throw error;
        }
    });
    // Tool to list tables
    server.tool('list-spanner-tables', {
        instanceId: z.string().optional().describe('Spanner instance ID (defaults to SPANNER_INSTANCE env var)'),
        databaseId: z.string().optional().describe('Spanner database ID (defaults to SPANNER_DATABASE env var)')
    }, async ({ instanceId, databaseId }, _extra) => {
        try {
            const projectId = await getProjectId();
            const config = await getSpannerConfig(Array.isArray(instanceId) ? instanceId[0] : instanceId, Array.isArray(databaseId) ? databaseId[0] : databaseId);
            const spanner = getSpannerClient();
            const instance = spanner.instance(config.instanceId);
            const database = instance.database(config.databaseId);
            // Query for tables
            console.error('Executing Spanner query to list tables...');
            const [tablesResult] = await database.run({
                sql: `SELECT t.table_name, 
                    (SELECT COUNT(1) FROM information_schema.columns 
                     WHERE table_name = t.table_name) as column_count
              FROM information_schema.tables t
              WHERE t.table_catalog = '' AND t.table_schema = ''
              ORDER BY t.table_name`
            });
            console.error('Raw tablesResult:', tablesResult);
            if (tablesResult && tablesResult.length > 0) {
                console.error('First row keys:', Object.keys(tablesResult[0]));
                console.error('First row:', tablesResult[0]);
            }
            if (!tablesResult || tablesResult.length === 0) {
                return {
                    content: [{
                            type: 'text',
                            text: `# Spanner Tables\n\nProject: ${projectId}\nInstance: ${config.instanceId}\nDatabase: ${config.databaseId}\n\nNo tables found in the database.`
                        }]
                };
            }
            let markdown = `# Spanner Tables\n\nProject: ${projectId}\nInstance: ${config.instanceId}\nDatabase: ${config.databaseId}\n\n`;
            // Table header
            markdown += '| Table Name | Column Count |\n';
            markdown += '|------------|-------------|\n';
            // Table rows
            for (const row of tablesResult) {
                // Convert the row to a plain JavaScript object
                const rowObj = JSON.parse(JSON.stringify(row));
                // Extract table name and column count
                const tableName = rowObj.table_name || 'unknown';
                const columnCount = rowObj.column_count || '0';
                markdown += `| ${tableName} | ${columnCount} |\n`;
            }
            return {
                content: [{
                        type: 'text',
                        text: markdown
                    }]
            };
        }
        catch (error) {
            console.error('Error listing Spanner tables:', error);
            throw error;
        }
    });
}
//# sourceMappingURL=spanner.js.map