ai-sql-knex/cli.js

AI-powered SQL query generator with advanced features: natural language to SQL conversion, schema exploration, query explanation, optimization suggestions, and interactive chat

#!/usr/bin/env node
import chalk from 'chalk'
import {Command} from 'commander'
import fs from 'fs'
import path from 'path'
import os from 'os'
import readline from 'readline/promises'
import {stdin, stdout} from 'process'
import OpenAI from 'openai'
import knex from 'knex'
import {URL} from 'url'

// --- Config ---
const CONFIG_PATH =
  process.env.AI_SQL_KNEX_CONFIG || path.join(os.homedir(), '.ai-sql-knex.json')
const HISTORY_PATH = path.join(os.homedir(), '.ai-sql-knex-history.json')
const configDir = path.dirname(CONFIG_PATH)
if (!fs.existsSync(configDir)) fs.mkdirSync(configDir, {recursive: true})

function loadConfig(filePath = CONFIG_PATH) {
  if (!fs.existsSync(filePath)) return {}
  try {
    return JSON.parse(fs.readFileSync(filePath, 'utf-8'))
  } catch (err) {
    console.error(chalk.red('❌ Failed to parse config JSON:'), err)
    return {}
  }
}

function saveConfig(config, filePath = CONFIG_PATH) {
  fs.writeFileSync(filePath, JSON.stringify(config, null, 2), 'utf-8')
  console.log(chalk.green(`✅ Configuration saved to ${filePath}`))
}

// --- History Management ---
function loadHistory(filePath = HISTORY_PATH) {
  if (!fs.existsSync(filePath)) return []
  try {
    return JSON.parse(fs.readFileSync(filePath, 'utf-8'))
  } catch (err) {
    console.error(chalk.red('❌ Failed to parse history JSON:'), err)
    return []
  }
}

function saveHistory(history, filePath = HISTORY_PATH) {
  fs.writeFileSync(filePath, JSON.stringify(history, null, 2), 'utf-8')
}

function addToHistory(nlQuery, sql, executed = false, timestamp = new Date().toISOString()) {
  const history = loadHistory()
  const entry = {
    id: Date.now(),
    timestamp,
    nlQuery,
    sql,
    executed
  }
  history.unshift(entry) // Add to beginning
  // Keep only last 20 entries
  if (history.length > 20) {
    history.splice(20)
  }
  saveHistory(history)
  return entry
}

async function promptIfMissing(key, value, promptText, secret = false, config) {
  if (value) return value
  const rl = readline.createInterface({input: stdin, output: stdout})
  const answer = secret
    ? await rl.question(`${promptText}: `, {hideEchoBack: true})
    : await rl.question(`${promptText}: `)
  rl.close()
  const trimmed = answer.trim()
  if (trimmed) {
    config[key] = trimmed
    saveConfig(config)
  }
  return trimmed
}

// --- DB Init ---
async function initDB(databaseUrl) {
  if (!databaseUrl) throw new Error('Database URL is required')
  const url = new URL(databaseUrl)
  const connection = {
    host: url.hostname,
    port: url.port ? parseInt(url.port) : undefined,
    user: url.username,
    password: url.password,
    database: url.pathname.replace(/^\//, ''),
  }

  const client = url.protocol.startsWith('postgres') ? 'pg' : 'mysql2'
  const db = knex({client, connection, pool: {min: 2, max: 10}})

  try {
    await db.raw(client === 'pg' ? 'SELECT 1' : 'SELECT 1')
    return db
  } catch (err) {
    console.error(
      chalk.red('❌ Failed to connect to the database:'),
      err.message,
    )
    process.exit(1)
  }
}

// --- CLI ---
const program = new Command()
program.name('ai-sql-knex').description('AI SQL History Manager')

// --- Config command ---
program
  .command('config')
  .description('Set or update CLI configuration')
  .option('--openai-key <key>', 'OpenAI API key')
  .option('--database-url <url>', 'Database connection URL')
  .action(async (opts) => {
    const config = loadConfig()
    if (opts.openaiKey) config.OPENAI_API_KEY = opts.openaiKey
    if (opts.databaseUrl) config.DATABASE_URL = opts.databaseUrl

    if (!opts.openaiKey && !opts.databaseUrl) {
      console.log(
        chalk.yellow('No options provided, prompting interactively...'),
      )
      config.OPENAI_API_KEY = await promptIfMissing(
        'OPENAI_API_KEY',
        config.OPENAI_API_KEY,
        'Enter OpenAI API key',
        true,
        config,
      )
      config.DATABASE_URL = await promptIfMissing(
        'DATABASE_URL',
        config.DATABASE_URL,
        'Enter database URL',
        false,
        config,
      )
    }

    saveConfig(config)
  })

// --- History command ---
program
  .command('history')
  .description('View query history')
  .option('--limit <number>', 'Number of entries to show', '10')
  .option('--executed', 'Show only executed queries')
  .action(async (options) => {
    const history = loadHistory()
    let filteredHistory = history
    
    if (options.executed) {
      filteredHistory = history.filter(entry => entry.executed)
    }
    
    const limit = parseInt(options.limit)
    const entries = filteredHistory.slice(0, limit)
    
    if (entries.length === 0) {
      console.log(chalk.yellow('No history entries found.'))
      return
    }
    
    console.log(chalk.blue.bold(`\n📚 Query History (${entries.length} entries):\n`))
    
    entries.forEach((entry, index) => {
      const status = entry.executed ? chalk.green('✓') : chalk.gray('○')
      const date = new Date(entry.timestamp).toLocaleString()
      
      console.log(chalk.cyan(`${index + 1}. [${status}] ${date}`))
      console.log(chalk.white(`   Query: ${entry.nlQuery}`))
      console.log(chalk.gray(`   SQL: ${entry.sql}`))
      console.log('')
    })
  })

// --- Chat command ---
program
  .command('chat')
  .description('Start an interactive chat with AI')
  .option('--model <model>', 'OpenAI model', 'gpt-4')
  .action(async (options) => {
    const config = loadConfig()
    const openaiKey = config.OPENAI_API_KEY

    if (!openaiKey) {
      console.error(chalk.red('❌ Missing OpenAI API key. Run `config` first.'))
      process.exit(1)
    }

    const client = new OpenAI({apiKey: openaiKey})
    const rl = readline.createInterface({input: stdin, output: stdout})
    
    console.log(chalk.blue.bold('\n🤖 AI Chat Started! Type "exit" to quit.\n'))
    
    const conversationHistory = []
    
    while (true) {
      try {
        const userInput = await rl.question(chalk.cyan('You: '))
        
        if (userInput.toLowerCase().trim() === 'exit') {
          console.log(chalk.yellow('\n👋 Chat ended. Goodbye!'))
          break
        }
        
        if (!userInput.trim()) continue
        
        // Add user message to conversation
        conversationHistory.push({role: 'user', content: userInput})
        
        process.stdout.write(chalk.gray('AI: '))
        
        const response = await client.chat.completions.create({
          model: options.model,
          messages: conversationHistory,
          temperature: 0.7,
          max_tokens: 12000,
        })
        
        const aiResponse = response.choices[0].message.content
        console.log(chalk.white(aiResponse))
        
        // Add AI response to conversation
        conversationHistory.push({role: 'assistant', content: aiResponse})
        
        // Keep conversation history manageable (last 10 exchanges)
        if (conversationHistory.length > 20) {
          conversationHistory.splice(0, 2)
        }
        
        console.log('') // Empty line for readability
        
      } catch (err) {
        console.error(chalk.red('\n❌ Error in chat:'), err.message)
        console.log('')
      }
    }
    
    rl.close()
  })

// --- Query command ---
program
  .command('query <nl_query...>')
  .description('Convert natural language to SQL and optionally execute')
  .option('--model <model>', 'OpenAI model', 'gpt-4')
  .option('--no-run', 'Only generate SQL, do not execute')
  .option('--explain', 'Explain the generated SQL query')
  .option('--optimize', 'Suggest optimizations for the query')
  .action(async (nlQuery, options) => {
    const config = loadConfig()
    const databaseUrl = config.DATABASE_URL
    const openaiKey = config.OPENAI_API_KEY

    if (!databaseUrl || !openaiKey) {
      console.error(chalk.red('❌ Missing configuration. Run `config` first.'))
      process.exit(1)
    }

    const db = await initDB(databaseUrl)
    const client = new OpenAI({apiKey: openaiKey})
    const queryText = nlQuery.join(' ')

    // Get detailed table information
    const tablesRaw = await db('information_schema.tables')
      .select('table_name')
      .where(
        'table_schema',
        db.client.config.client === 'pg' ? 'public' : db.client.database(),
      )
    
    const tableList = tablesRaw.map((t) => t.table_name)
    
    // Get simplified schema information (only essential columns)
    let schemaInfo = ''
    const maxTables = Math.min(30, tableList.length) // Limit to 5 tables max
    
    for (const tableName of tableList.slice(0, maxTables)) {
      try {
        const columns = await db('information_schema.columns')
          .select('column_name', 'data_type')
          .where('table_name', tableName)
          .where('table_schema', db.client.config.client === 'pg' ? 'public' : db.client.database())
          .limit(8) // Limit to 8 columns per table
        
        schemaInfo += `\n${tableName}: `
        const columnList = columns.map(col => `${col.column_name}(${col.data_type})`).join(', ')
        schemaInfo += columnList
      } catch (err) {
        // If we can't get column info, just include table name
        schemaInfo += `\n${tableName}`
      }
    }
    
    // If we have more tables, add a summary
    if (tableList.length > maxTables) {
      schemaInfo += `\n... and ${tableList.length - maxTables} more tables`
    }

    // Get recent history for context
    const history = loadHistory().slice(0, 3) // Last 3 queries for better context
    const historyContext = history.length > 0 
      ? `\n\nRecent successful queries for reference:\n${history.filter(entry => entry.executed).map(entry => 
          `- "${entry.nlQuery}" → ${entry.sql}`
        ).join('\n')}`
      : ''

    // Generate SQL
    let resp
    try {
      resp = await client.chat.completions.create({
        model: options.model,
        messages: [
          {
            role: 'system',
            content: `Generate SQL queries from natural language. Available tables:${schemaInfo}${historyContext}

Rules: Use proper SQL syntax, include LIMIT for exploration, return only SQL (no explanations).`,
          },
          {role: 'user', content: queryText},
        ],
        temperature: 0.1,
        max_tokens: 1500,
      })
    } catch (err) {
      if (err.code === 'context_length_exceeded') {
        console.error(chalk.red('❌ Query too complex. Try simplifying your request or use --model gpt-3.5-turbo'))
        await db.destroy()
        return
      } else if (err.code === 'insufficient_quota') {
        console.error(chalk.red('❌ OpenAI API quota exceeded. Please check your account.'))
        await db.destroy()
        return
      } else {
        console.error(chalk.red('❌ Error generating SQL:'), err.message)
        await db.destroy()
        return
      }
    }

    const sql = resp.choices[0].message.content
      .replace(/```(?:sql)?\n?|```/g, '')
      .trim()
    console.log(chalk.green.bold('\n💡 Generated SQL:\n'))
    console.log(chalk.cyan(sql))
    
    // Add explanation if requested
    if (options.explain) {
      console.log(chalk.blue.bold('\n📖 Query Explanation:\n'))
      try {
        const explainResp = await client.chat.completions.create({
          model: options.model,
          messages: [
            {
              role: 'system',
              content: 'You are a SQL expert. Explain what this SQL query does in simple terms, breaking down each part of the query.'
            },
            {role: 'user', content: `Explain this SQL query: ${sql}`}
          ],
          temperature: 0.3,
          max_tokens: 500,
        })
        console.log(chalk.white(explainResp.choices[0].message.content))
      } catch (err) {
        console.log(chalk.red('❌ Failed to generate explanation'))
      }
    }
    
    // Add optimization suggestions if requested
    if (options.optimize) {
      console.log(chalk.yellow.bold('\n⚡ Optimization Suggestions:\n'))
      try {
        const optimizeResp = await client.chat.completions.create({
          model: options.model,
          messages: [
            {
              role: 'system',
              content: 'You are a SQL performance expert. Analyze this query and suggest specific optimizations for better performance, including index recommendations, query structure improvements, and best practices.'
            },
            {role: 'user', content: `Optimize this SQL query: ${sql}`}
          ],
          temperature: 0.3,
          max_tokens: 600,
        })
        console.log(chalk.white(optimizeResp.choices[0].message.content))
      } catch (err) {
        console.log(chalk.red('❌ Failed to generate optimization suggestions'))
      }
    }
    
    // Add to history (not executed yet)
    addToHistory(queryText, sql, false)

    if (!options.run) {
      console.log(chalk.yellow('\nℹ️  Use --run to execute the query'))
      await db.destroy()
      return
    }

    // Confirm execution
    const rl = readline.createInterface({input: stdin, output: stdout})
    const queryType = sql.trim().split(' ')[0].toUpperCase()
    const isSelect = queryType === 'SELECT'
    const answer = await rl.question(
      `⚠️ Execute this ${isSelect ? 'SELECT' : 'non-SELECT'} query? (y/N): `,
    )
    rl.close()
    if (!/^y/i.test(answer)) {
      console.log(chalk.red('❌ Query execution cancelled'))
      await db.destroy()
      return
    }

    try {
      const result = await db.raw(sql)

      // MySQL: result is [rows, fields], PostgreSQL: result.rows
      const rows = result.rows || result[0]

      if (rows && rows.length) {
        console.table(rows)
        console.log(chalk.green(`${rows.length} row(s) returned.`))
      } else {
        console.log(chalk.yellow('No results returned.'))
      }
      
      // Update history entry to mark as executed
      const history = loadHistory()
      const latestEntry = history.find(entry => entry.sql === sql && !entry.executed)
      if (latestEntry) {
        latestEntry.executed = true
        saveHistory(history)
      }
    } catch (err) {
      console.error(chalk.red('❌ Error executing query:'), err.message)
    }
    
    await db.destroy()
  })

// --- Schema command ---
program
  .command('schema')
  .description('Explore database schema and table information')
  .option('--table <table>', 'Show details for specific table')
  .option('--tables', 'List all tables')
  .action(async (options) => {
    const config = loadConfig()
    const databaseUrl = config.DATABASE_URL

    if (!databaseUrl) {
      console.error(chalk.red('❌ Missing database URL. Run `config` first.'))
      process.exit(1)
    }

    const db = await initDB(databaseUrl)
    
    try {
      if (options.tables || (!options.table && !options.tables)) {
        // List all tables
        const tables = await db('information_schema.tables')
          .select('table_name', 'table_type')
          .where('table_schema', db.client.config.client === 'pg' ? 'public' : db.client.database())
          .orderBy('table_name')
        
        console.log(chalk.blue.bold('\n📋 Database Tables:\n'))
        tables.forEach(table => {
          console.log(chalk.cyan(`• ${table.table_name}`) + chalk.gray(` (${table.table_type})`))
        })
        console.log(chalk.gray(`\nTotal: ${tables.length} tables`))
      }
      
      if (options.table) {
        // Show specific table details
        console.log(chalk.blue.bold(`\n🔍 Table: ${options.table}\n`))
        
        const columns = await db('information_schema.columns')
          .select('column_name', 'data_type', 'is_nullable', 'column_default', 'character_maximum_length')
          .where('table_name', options.table)
          .where('table_schema', db.client.config.client === 'pg' ? 'public' : db.client.database())
          .orderBy('ordinal_position')
        
        if (columns.length === 0) {
          console.log(chalk.red(`❌ Table '${options.table}' not found`))
        } else {
          console.log(chalk.white('Columns:'))
          columns.forEach(col => {
            const nullable = col.is_nullable === 'NO' ? chalk.red('NOT NULL') : chalk.green('NULL')
            const length = col.character_maximum_length ? `(${col.character_maximum_length})` : ''
            const defaultVal = col.column_default ? ` DEFAULT ${col.column_default}` : ''
            console.log(chalk.cyan(`  • ${col.column_name}`) + chalk.gray(` ${col.data_type}${length}`) + ` ${nullable}` + chalk.gray(defaultVal))
          })
          
          // Get row count
          try {
            const countResult = await db(options.table).count('* as count').first()
            const rowCount = countResult.count || countResult[Object.keys(countResult)[0]]
            console.log(chalk.gray(`\nRows: ${rowCount}`))
          } catch (err) {
            console.log(chalk.gray('\nRows: Unable to count'))
          }
        }
      }
      
    } catch (err) {
      console.error(chalk.red('❌ Error exploring schema:'), err.message)
    } finally {
      await db.destroy()
    }
  })

// --- Global error handling ---
process.on('unhandledRejection', (err) => {
  console.error('Unhandled rejection:', err)
  process.exit(1)
})
process.on('uncaughtException', (err) => {
  console.error('Uncaught exception:', err)
  process.exit(1)
})

// --- Run CLI ---
program.parse(process.argv)