truss-db-mcp 1.0.0

package/src/lib/ai-query.ts

@@ -0,0 +1,274 @@
+import { getConfig } from './config.js';
+import { listTables, executeQuery } from './connection.js';
+import type { NaturalLanguageQueryResult, TableInfo } from '../types.js';
+
+// ── Schema Summary ──────────────────────────────────────────────────
+
+function buildSchemaContext(filterTables?: string[]): string {
+  const tables = listTables();
+  const relevant = filterTables
+    ? tables.filter(t => filterTables.includes(t.name))
+    : tables;
+
+  return relevant.map(t => formatTableSchema(t)).join('\n\n');
+}
+
+function formatTableSchema(t: TableInfo): string {
+  const cols = t.columns.map(c => {
+    const parts = [`  ${c.name} ${c.type}`];
+    if (c.primaryKey) parts.push('PRIMARY KEY');
+    if (!c.nullable) parts.push('NOT NULL');
+    if (c.defaultValue !== null) parts.push(`DEFAULT ${c.defaultValue}`);
+    return parts.join(' ');
+  }).join(',\n');
+
+  return `TABLE ${t.name} (${t.rowCount} rows):\n${cols}`;
+}
+
+// ── AI Providers ────────────────────────────────────────────────────
+
+interface AiResponse {
+  sql: string;
+  explanation: string;
+}
+
+async function callAnthropic(prompt: string, apiKey: string): Promise<AiResponse> {
+  const response = await fetch('https://api.anthropic.com/v1/messages', {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      'x-api-key': apiKey,
+      'anthropic-version': '2023-06-01',
+    },
+    body: JSON.stringify({
+      model: 'claude-sonnet-4-20250514',
+      max_tokens: 1024,
+      messages: [{
+        role: 'user',
+        content: prompt,
+      }],
+    }),
+    signal: AbortSignal.timeout(30000),
+  });
+
+  if (!response.ok) {
+    const body = await response.text();
+    throw new Error(`Anthropic API error (${response.status}): ${body}`);
+  }
+
+  const data = await response.json() as {
+    content: Array<{ type: string; text: string }>;
+  };
+
+  const text = data.content.find(c => c.type === 'text')?.text ?? '';
+  return parseAiResponse(text);
+}
+
+async function callOpenAI(prompt: string, apiKey: string): Promise<AiResponse> {
+  const response = await fetch('https://api.openai.com/v1/chat/completions', {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      'Authorization': `Bearer ${apiKey}`,
+    },
+    body: JSON.stringify({
+      model: 'gpt-4o-mini',
+      messages: [{
+        role: 'user',
+        content: prompt,
+      }],
+      max_tokens: 1024,
+    }),
+    signal: AbortSignal.timeout(30000),
+  });
+
+  if (!response.ok) {
+    const body = await response.text();
+    throw new Error(`OpenAI API error (${response.status}): ${body}`);
+  }
+
+  const data = await response.json() as {
+    choices: Array<{ message: { content: string } }>;
+  };
+
+  const text = data.choices[0]?.message?.content ?? '';
+  return parseAiResponse(text);
+}
+
+function parseAiResponse(text: string): AiResponse {
+  // Extract SQL from code blocks or raw text
+  const sqlMatch = text.match(/```sql\n([\s\S]*?)```/);
+  const sql = sqlMatch ? sqlMatch[1].trim() : extractSqlFromText(text);
+
+  // Extract explanation (everything that is not the SQL block)
+  const explanation = text
+    .replace(/```sql[\s\S]*?```/, '')
+    .replace(/^SQL:\s*/mi, '')
+    .trim()
+    || 'Query generated from natural language.';
+
+  if (!sql) {
+    throw new Error('AI did not produce a valid SQL query. Try rephrasing your question.');
+  }
+
+  return { sql, explanation };
+}
+
+function extractSqlFromText(text: string): string {
+  // Try to find a SELECT/WITH statement in the text
+  const match = text.match(/((?:WITH|SELECT)\b[\s\S]*?;)/i);
+  return match ? match[1].trim() : '';
+}
+
+// ── Public API ──────────────────────────────────────────────────────
+
+export async function naturalLanguageToSql(
+  question: string,
+  filterTables?: string[]
+): Promise<NaturalLanguageQueryResult> {
+  const config = getConfig();
+
+  if (!config.anthropicApiKey && !config.openaiApiKey) {
+    throw new Error(
+      'Natural language queries require an AI API key. ' +
+      'Set ANTHROPIC_API_KEY or OPENAI_API_KEY in your environment.'
+    );
+  }
+
+  const schema = buildSchemaContext(filterTables);
+
+  const prompt = `You are a SQL query generator. Given the database schema below and a natural language question, generate the correct SQL query.
+
+DATABASE SCHEMA:
+${schema}
+
+QUESTION: ${question}
+
+Respond with:
+1. The SQL query in a \`\`\`sql code block
+2. A brief explanation of what the query does
+
+Important:
+- Use only tables and columns that exist in the schema above
+- Write SQLite-compatible SQL
+- Use appropriate JOINs when data spans multiple tables
+- Always include a LIMIT if the result could be large (default LIMIT 100)`;
+
+  let aiResult: AiResponse;
+
+  if (config.anthropicApiKey) {
+    aiResult = await callAnthropic(prompt, config.anthropicApiKey);
+  } else {
+    aiResult = await callOpenAI(prompt, config.openaiApiKey!);
+  }
+
+  // Execute the generated SQL
+  const results = executeQuery(aiResult.sql);
+
+  return {
+    sql: aiResult.sql,
+    explanation: aiResult.explanation,
+    results,
+  };
+}
+
+export async function generateMigrationSql(
+  description: string,
+  dialect: string = 'sqlite'
+): Promise<{ upSql: string; downSql: string }> {
+  const config = getConfig();
+
+  if (!config.anthropicApiKey && !config.openaiApiKey) {
+    throw new Error(
+      'Migration generation requires an AI API key. ' +
+      'Set ANTHROPIC_API_KEY or OPENAI_API_KEY in your environment.'
+    );
+  }
+
+  let schemaContext = '';
+  try {
+    schemaContext = buildSchemaContext();
+  } catch {
+    // No connection -- generate migration without existing schema context
+  }
+
+  const prompt = `You are a database migration generator. Generate SQL migration statements for the following change.
+
+${schemaContext ? `EXISTING SCHEMA:\n${schemaContext}\n\n` : ''}DIALECT: ${dialect}
+CHANGE DESCRIPTION: ${description}
+
+Respond with two SQL code blocks:
+
+1. UP migration (apply the change):
+\`\`\`sql
+-- up migration here
+\`\`\`
+
+2. DOWN migration (revert the change):
+\`\`\`sql
+-- down migration here
+\`\`\`
+
+Important:
+- Use IF EXISTS / IF NOT EXISTS where appropriate
+- Make migrations idempotent when possible
+- Use correct syntax for the ${dialect} dialect`;
+
+  let text: string;
+
+  if (config.anthropicApiKey) {
+    const response = await fetch('https://api.anthropic.com/v1/messages', {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'x-api-key': config.anthropicApiKey,
+        'anthropic-version': '2023-06-01',
+      },
+      body: JSON.stringify({
+        model: 'claude-sonnet-4-20250514',
+        max_tokens: 2048,
+        messages: [{ role: 'user', content: prompt }],
+      }),
+      signal: AbortSignal.timeout(30000),
+    });
+    if (!response.ok) throw new Error(`Anthropic API error: ${response.status}`);
+    const data = await response.json() as { content: Array<{ type: string; text: string }> };
+    text = data.content.find(c => c.type === 'text')?.text ?? '';
+  } else {
+    const response = await fetch('https://api.openai.com/v1/chat/completions', {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'Authorization': `Bearer ${config.openaiApiKey}`,
+      },
+      body: JSON.stringify({
+        model: 'gpt-4o-mini',
+        messages: [{ role: 'user', content: prompt }],
+        max_tokens: 2048,
+      }),
+      signal: AbortSignal.timeout(30000),
+    });
+    if (!response.ok) throw new Error(`OpenAI API error: ${response.status}`);
+    const data = await response.json() as { choices: Array<{ message: { content: string } }> };
+    text = data.choices[0]?.message?.content ?? '';
+  }
+
+  // Parse two SQL blocks from response
+  const sqlBlocks = [...text.matchAll(/```sql\n([\s\S]*?)```/g)].map(m => m[1].trim());
+
+  if (sqlBlocks.length < 2) {
+    // Try to split on common markers
+    const upMatch = text.match(/(?:UP|up|Up)[:\s]*\n*```sql\n([\s\S]*?)```/);
+    const downMatch = text.match(/(?:DOWN|down|Down)[:\s]*\n*```sql\n([\s\S]*?)```/);
+
+    return {
+      upSql: upMatch?.[1]?.trim() || sqlBlocks[0] || '-- Could not generate up migration',
+      downSql: downMatch?.[1]?.trim() || sqlBlocks[1] || '-- Could not generate down migration',
+    };
+  }
+
+  return {
+    upSql: sqlBlocks[0],
+    downSql: sqlBlocks[1],
+  };
+}

package/src/lib/config.ts

@@ -0,0 +1,33 @@
+import { homedir } from 'node:os';
+import { join } from 'node:path';
+import { mkdirSync } from 'node:fs';
+import type { TrussConfig } from '../types.js';
+
+const DATA_DIR = join(homedir(), '.truss');
+
+// Ensure data directory exists on import
+try {
+  mkdirSync(DATA_DIR, { recursive: true });
+} catch {
+  // ignore -- directory already exists
+}
+
+export function getConfig(): TrussConfig {
+  return {
+    licenseKey: process.env.TRUSS_LICENSE_KEY,
+    databaseUrl: process.env.DATABASE_URL,
+    dbHost: process.env.DB_HOST,
+    dbPort: process.env.DB_PORT ? parseInt(process.env.DB_PORT, 10) : undefined,
+    dbName: process.env.DB_NAME,
+    dbUser: process.env.DB_USER,
+    dbPassword: process.env.DB_PASSWORD,
+    anthropicApiKey: process.env.ANTHROPIC_API_KEY,
+    openaiApiKey: process.env.OPENAI_API_KEY,
+    trussApiBaseUrl: process.env.TRUSS_API_BASE_URL || 'https://api.truss.dev',
+    dataDir: DATA_DIR,
+  };
+}
+
+export function getDataDir(): string {
+  return DATA_DIR;
+}

package/src/lib/connection.ts

@@ -0,0 +1,263 @@
+import { existsSync } from 'node:fs';
+import { resolve } from 'node:path';
+import Database from 'better-sqlite3';
+import type {
+  DatabaseDialect,
+  ConnectionInfo,
+  TableInfo,
+  TableSchema,
+  ColumnInfo,
+  IndexInfo,
+  ForeignKeyInfo,
+  QueryResult,
+} from '../types.js';
+
+// ── Connection State ────────────────────────────────────────────────
+
+let currentDb: Database.Database | null = null;
+let currentDialect: DatabaseDialect = 'sqlite';
+let currentPath: string | null = null;
+
+// ── SQLite Connection ───────────────────────────────────────────────
+
+export function connectSqlite(dbPath: string): ConnectionInfo {
+  const absPath = resolve(dbPath);
+
+  if (!existsSync(absPath)) {
+    throw new Error(`Database file not found: ${absPath}`);
+  }
+
+  // Close existing connection
+  if (currentDb) {
+    try { currentDb.close(); } catch { /* ignore */ }
+  }
+
+  currentDb = new Database(absPath, { readonly: false });
+  currentDialect = 'sqlite';
+  currentPath = absPath;
+
+  // Enable WAL mode for better concurrency
+  currentDb.pragma('journal_mode = WAL');
+
+  const tables = listTableNames();
+  const version = currentDb.pragma('user_version') as Array<{ user_version: number }>;
+  const sqliteVersion = currentDb.prepare('SELECT sqlite_version() as v').get() as { v: string };
+
+  return {
+    connected: true,
+    dialect: 'sqlite',
+    tables,
+    version: `SQLite ${sqliteVersion.v}`,
+    path: absPath,
+  };
+}
+
+// ── Connection Accessors ────────────────────────────────────────────
+
+export function getDb(): Database.Database {
+  if (!currentDb) {
+    throw new Error('No database connected. Use the connect tool first.');
+  }
+  return currentDb;
+}
+
+export function getDialect(): DatabaseDialect {
+  return currentDialect;
+}
+
+export function isConnected(): boolean {
+  return currentDb !== null;
+}
+
+export function closeConnection(): void {
+  if (currentDb) {
+    try { currentDb.close(); } catch { /* ignore */ }
+    currentDb = null;
+    currentPath = null;
+  }
+}
+
+// ── Table Operations ────────────────────────────────────────────────
+
+function listTableNames(): string[] {
+  const db = getDb();
+  const rows = db.prepare(
+    "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name"
+  ).all() as Array<{ name: string }>;
+  return rows.map(r => r.name);
+}
+
+export function listTables(): TableInfo[] {
+  const db = getDb();
+  const tableNames = listTableNames();
+  const results: TableInfo[] = [];
+
+  for (const tableName of tableNames) {
+    const columns = getColumnInfo(tableName);
+    const countRow = db.prepare(`SELECT COUNT(*) as cnt FROM "${tableName}"`).get() as { cnt: number };
+
+    results.push({
+      name: tableName,
+      columns,
+      rowCount: countRow.cnt,
+    });
+  }
+
+  return results;
+}
+
+export function describeTable(tableName: string): TableSchema {
+  const db = getDb();
+
+  // Verify table exists
+  const exists = db.prepare(
+    "SELECT name FROM sqlite_master WHERE type='table' AND name = ?"
+  ).get(tableName) as { name: string } | undefined;
+
+  if (!exists) {
+    throw new Error(`Table "${tableName}" not found.`);
+  }
+
+  const columns = getColumnInfo(tableName);
+  const indexes = getIndexInfo(tableName);
+  const foreignKeys = getForeignKeyInfo(tableName);
+  const countRow = db.prepare(`SELECT COUNT(*) as cnt FROM "${tableName}"`).get() as { cnt: number };
+
+  // Get CREATE TABLE statement
+  const createRow = db.prepare(
+    "SELECT sql FROM sqlite_master WHERE type='table' AND name = ?"
+  ).get(tableName) as { sql: string };
+
+  return {
+    name: tableName,
+    columns,
+    rowCount: countRow.cnt,
+    indexes,
+    foreignKeys,
+    createSql: createRow.sql,
+  };
+}
+
+function getColumnInfo(tableName: string): ColumnInfo[] {
+  const db = getDb();
+  const rows = db.pragma(`table_info("${tableName}")`) as Array<{
+    cid: number;
+    name: string;
+    type: string;
+    notnull: number;
+    dflt_value: string | null;
+    pk: number;
+  }>;
+
+  return rows.map(r => ({
+    name: r.name,
+    type: r.type || 'ANY',
+    nullable: r.notnull === 0,
+    defaultValue: r.dflt_value,
+    primaryKey: r.pk > 0,
+  }));
+}
+
+function getIndexInfo(tableName: string): IndexInfo[] {
+  const db = getDb();
+  const indexes = db.pragma(`index_list("${tableName}")`) as Array<{
+    seq: number;
+    name: string;
+    unique: number;
+    origin: string;
+    partial: number;
+  }>;
+
+  return indexes.map(idx => {
+    const cols = db.pragma(`index_info("${idx.name}")`) as Array<{
+      seqno: number;
+      cid: number;
+      name: string;
+    }>;
+
+    return {
+      name: idx.name,
+      columns: cols.map(c => c.name),
+      unique: idx.unique === 1,
+    };
+  });
+}
+
+function getForeignKeyInfo(tableName: string): ForeignKeyInfo[] {
+  const db = getDb();
+  const fks = db.pragma(`foreign_key_list("${tableName}")`) as Array<{
+    id: number;
+    seq: number;
+    table: string;
+    from: string;
+    to: string;
+    on_update: string;
+    on_delete: string;
+    match: string;
+  }>;
+
+  return fks.map(fk => ({
+    column: fk.from,
+    referencedTable: fk.table,
+    referencedColumn: fk.to,
+  }));
+}
+
+// ── Query Execution ─────────────────────────────────────────────────
+
+const WRITE_PATTERN = /^\s*(INSERT|UPDATE|DELETE|DROP|ALTER|CREATE|TRUNCATE|REPLACE|MERGE)\b/i;
+
+export function isWriteQuery(sql: string): boolean {
+  return WRITE_PATTERN.test(sql);
+}
+
+export function executeQuery(sql: string, params: unknown[] = []): QueryResult {
+  const db = getDb();
+  const start = performance.now();
+
+  const stmt = db.prepare(sql);
+  let rows: Record<string, unknown>[];
+
+  if (isWriteQuery(sql)) {
+    const info = stmt.run(...params);
+    rows = [{ changes: info.changes, lastInsertRowid: Number(info.lastInsertRowid) }];
+  } else {
+    rows = stmt.all(...params) as Record<string, unknown>[];
+  }
+
+  const executionTimeMs = Math.round((performance.now() - start) * 100) / 100;
+
+  // Extract column names
+  const columns = rows.length > 0 ? Object.keys(rows[0]) : [];
+
+  return {
+    rows,
+    columns,
+    rowCount: rows.length,
+    executionTimeMs,
+  };
+}
+
+export function executeReadOnlyQuery(sql: string, params: unknown[] = []): QueryResult {
+  if (isWriteQuery(sql)) {
+    throw new Error(
+      'Write operations (INSERT, UPDATE, DELETE, DROP, ALTER, CREATE, TRUNCATE) are not allowed in free tier. ' +
+      'Upgrade to Pro for full write access: https://truss.dev/pricing'
+    );
+  }
+  return executeQuery(sql, params);
+}
+
+// ── EXPLAIN ─────────────────────────────────────────────────────────
+
+export function explainQuery(sql: string): string {
+  const db = getDb();
+  const rows = db.prepare(`EXPLAIN QUERY PLAN ${sql}`).all() as Array<{
+    id: number;
+    parent: number;
+    notused: number;
+    detail: string;
+  }>;
+
+  return rows.map(r => `${' '.repeat(r.id * 2)}${r.detail}`).join('\n');
+}

package/src/lib/license.ts

@@ -0,0 +1,118 @@
+import { readFileSync, writeFileSync } from 'node:fs';
+import { join } from 'node:path';
+import { getConfig, getDataDir } from './config.js';
+import type { LicenseStatus, LicenseTier } from '../types.js';
+
+const CACHE_FILE = 'db-license-cache.json';
+const CACHE_TTL_MS = 7 * 24 * 60 * 60 * 1000; // 7 days
+
+interface LicenseCache {
+  key: string;
+  valid: boolean;
+  tier: LicenseTier;
+  expiresAt: string | null;
+  cachedAt: number;
+}
+
+function getCachePath(): string {
+  return join(getDataDir(), CACHE_FILE);
+}
+
+function readCache(): LicenseCache | null {
+  try {
+    const raw = readFileSync(getCachePath(), 'utf-8');
+    const cache = JSON.parse(raw) as LicenseCache;
+    if (Date.now() - cache.cachedAt < CACHE_TTL_MS) {
+      return cache;
+    }
+    return null; // expired
+  } catch {
+    return null;
+  }
+}
+
+function writeCache(cache: LicenseCache): void {
+  try {
+    writeFileSync(getCachePath(), JSON.stringify(cache, null, 2));
+  } catch {
+    // non-fatal
+  }
+}
+
+function isValidKeyFormat(key: string): boolean {
+  return /^truss_[0-9a-f]{32}$/.test(key);
+}
+
+async function validateRemote(key: string): Promise<{ valid: boolean; expiresAt: string | null }> {
+  const config = getConfig();
+  const url = `${config.trussApiBaseUrl}/validate/${key}`;
+
+  try {
+    const response = await fetch(url, {
+      method: 'GET',
+      headers: {
+        'Accept': 'application/json',
+        'User-Agent': 'truss-db-mcp/1.0.0',
+      },
+      signal: AbortSignal.timeout(5000),
+    });
+
+    if (!response.ok) {
+      return { valid: false, expiresAt: null };
+    }
+
+    const body = await response.json() as { valid: boolean; expires_at?: string };
+    return {
+      valid: body.valid === true,
+      expiresAt: body.expires_at ?? null,
+    };
+  } catch {
+    // Network error -- fall back to format check + cache
+    return { valid: isValidKeyFormat(key), expiresAt: null };
+  }
+}
+
+export async function getLicenseStatus(): Promise<LicenseStatus> {
+  const config = getConfig();
+  const key = config.licenseKey;
+
+  // No key -> free tier
+  if (!key) {
+    return { tier: 'free', valid: true, expiresAt: null };
+  }
+
+  // Format check
+  if (!isValidKeyFormat(key)) {
+    return { tier: 'free', valid: false, expiresAt: null };
+  }
+
+  // Check cache first
+  const cached = readCache();
+  if (cached && cached.key === key) {
+    return { tier: cached.tier, valid: cached.valid, expiresAt: cached.expiresAt };
+  }
+
+  // Remote validation
+  const result = await validateRemote(key);
+  const tier: LicenseTier = result.valid ? 'pro' : 'free';
+
+  writeCache({
+    key,
+    valid: result.valid,
+    tier,
+    expiresAt: result.expiresAt,
+    cachedAt: Date.now(),
+  });
+
+  return { tier, valid: result.valid, expiresAt: result.expiresAt };
+}
+
+export async function requirePro(): Promise<void> {
+  const status = await getLicenseStatus();
+  if (status.tier !== 'pro') {
+    throw new Error(
+      'This feature requires a TRUSS Pro license ($25/mo). ' +
+      'Get yours at https://truss.dev/pricing and set TRUSS_LICENSE_KEY.'
+    );
+  }
+}