@hyqf98/easy_db_mcp_server 1.0.0 → 2.0.0

package/src/database/postgresql.ts

@@ -1,11 +1,25 @@
 import pg from 'pg';
+import fs from 'fs';
 import type {
   DatabaseAdapter,
   TableInfo,
   TableColumn,
   QueryResult,
+  TransactionResult,
+  BatchInsertResult,
+  BatchUpdateResult,
+  BatchUpdateOptions,
+  ExportResult,
+  ExportOptions,
+  CreateTableResult,
+  DropTableResult,
+  TableColumnDef,
+  TableStatsResult,
+  PreviewDataResult,
+  SampleDataResult,
 } from './base.js';
 import type { DatabaseConfig } from '../config.js';
+import { expandTilde, formatFileSize, ensureDirectoryExists, getDefaultFilePath } from '../utils/file.js';
 
 const { Pool } = pg;
 
@@ -87,12 +101,36 @@ export class PostgreSQLAdapter implements DatabaseAdapter {
   async executeQuery(sql: string, database?: string): Promise<QueryResult> {
     if (!this.pool) throw new Error('Not connected');
 
-    // Allow queries with fully qualified table names (schema.table)
     const trimmed = sql.trim().toUpperCase();
     if (!trimmed.startsWith('SELECT')) {
       throw new Error('Only SELECT queries allowed in executeQuery');
     }
 
+    // If database is specified and different from default, use it
+    const dbName = database || this.config.database;
+    if (dbName && dbName !== this.config.database) {
+      // For PostgreSQL, we need to connect to the specific database
+      // Create a temporary pool for the different database
+      const tempPool = new Pool({
+        host: this.config.host,
+        port: this.config.port || 5432,
+        user: this.config.user,
+        password: this.config.password,
+        database: dbName,
+        max: 1,
+      });
+
+      try {
+        const result = await tempPool.query(sql);
+        return {
+          rows: result.rows,
+          rowCount: result.rowCount || 0,
+        };
+      } finally {
+        await tempPool.end();
+      }
+    }
+
     const result = await this.pool.query(sql);
 
     return {
@@ -104,7 +142,36 @@ export class PostgreSQLAdapter implements DatabaseAdapter {
   async executeSQL(sql: string, database?: string): Promise<QueryResult | { affectedRows: number }> {
     if (!this.pool) throw new Error('Not connected');
 
-    // Allow SQL execution with database-specific queries
+    // If database is specified and different from default, use it
+    const dbName = database || this.config.database;
+    if (dbName && dbName !== this.config.database) {
+      // For PostgreSQL, we need to connect to the specific database
+      // Create a temporary pool for the different database
+      const tempPool = new Pool({
+        host: this.config.host,
+        port: this.config.port || 5432,
+        user: this.config.user,
+        password: this.config.password,
+        database: dbName,
+        max: 1,
+      });
+
+      try {
+        const result = await tempPool.query(sql);
+
+        if (result.rowCount !== null && result.command !== 'SELECT') {
+          return { affectedRows: result.rowCount };
+        }
+
+        return {
+          rows: result.rows,
+          rowCount: result.rowCount || 0,
+        };
+      } finally {
+        await tempPool.end();
+      }
+    }
+
     const result = await this.pool.query(sql);
 
     if (result.rowCount !== null && result.command !== 'SELECT') {
@@ -123,4 +190,489 @@ export class PostgreSQLAdapter implements DatabaseAdapter {
       this.pool = undefined;
     }
   }
+
+  async executeTransaction(queries: string[], database?: string): Promise<TransactionResult> {
+    if (!this.pool) throw new Error('Not connected');
+
+    let client: pg.PoolClient | null = null;
+
+    // If different database, use temporary connection
+    if (database && database !== this.config.database) {
+      const tempPool = new pg.Pool({
+        host: this.config.host,
+        port: this.config.port || 5432,
+        user: this.config.user,
+        password: this.config.password,
+        database,
+        max: 1,
+      });
+
+      try {
+        client = await tempPool.connect();
+        await client.query('BEGIN');
+
+        const results: unknown[] = [];
+        let totalAffectedRows = 0;
+
+        for (const query of queries) {
+          const result = await client.query(query);
+          if (result.rowCount !== null) {
+            totalAffectedRows += result.rowCount;
+          }
+          results.push(result);
+        }
+
+        await client.query('COMMIT');
+
+        return {
+          success: true,
+          affectedRows: totalAffectedRows,
+          results,
+        };
+      } catch (error) {
+        if (client) await client.query('ROLLBACK');
+        throw error;
+      } finally {
+        if (client) client.release();
+        await tempPool.end();
+      }
+    }
+
+    // Use default pool
+    client = await this.pool.connect();
+
+    try {
+      await client.query('BEGIN');
+
+      const results: unknown[] = [];
+      let totalAffectedRows = 0;
+
+      for (const query of queries) {
+        const result = await client.query(query);
+        if (result.rowCount !== null) {
+          totalAffectedRows += result.rowCount;
+        }
+        results.push(result);
+      }
+
+      await client.query('COMMIT');
+
+      return {
+        success: true,
+        affectedRows: totalAffectedRows,
+        results,
+      };
+    } catch (error) {
+      await client.query('ROLLBACK');
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+
+  async batchInsert(
+    table: string,
+    data: Record<string, unknown>[],
+    database?: string
+  ): Promise<BatchInsertResult> {
+    if (!this.pool) throw new Error('Not connected');
+
+    if (data.length === 0) {
+      return { insertedRows: 0, duplicateRows: 0 };
+    }
+
+    let pool: pg.Pool = this.pool;
+    let dbName = database || this.config.database;
+
+    // Use temporary pool for different database
+    if (database && database !== this.config.database) {
+      pool = new pg.Pool({
+        host: this.config.host,
+        port: this.config.port || 5432,
+        user: this.config.user,
+        password: this.config.password,
+        database,
+        max: 1,
+      });
+    }
+
+    const client = await pool.connect();
+
+    try {
+      await client.query('BEGIN');
+
+      const columns = Object.keys(data[0]);
+      const columnList = columns.map((c) => pg.escapeIdentifier(c)).join(', ');
+
+      let insertedRows = 0;
+      let duplicateRows = 0;
+
+      for (const row of data) {
+        const values = columns.map((col) => row[col]);
+        const placeholders = values.map((_, i) => `$${i + 1}`).join(', ');
+
+        try {
+          const query = `INSERT INTO ${pg.escapeIdentifier(table)} (${columnList}) VALUES (${placeholders})`;
+          await client.query(query, values);
+          insertedRows++;
+        } catch (error: any) {
+          if (error.code === '23505') { // unique_violation
+            duplicateRows++;
+          } else {
+            await client.query('ROLLBACK');
+            throw error;
+          }
+        }
+      }
+
+      await client.query('COMMIT');
+
+      return { insertedRows, duplicateRows };
+    } finally {
+      client.release();
+      if (pool !== this.pool) await pool.end();
+    }
+  }
+
+  async batchUpdate(
+    table: string,
+    updates: BatchUpdateOptions,
+    database?: string
+  ): Promise<BatchUpdateResult> {
+    if (!this.pool) throw new Error('Not connected');
+
+    let pool: pg.Pool = this.pool;
+
+    if (database && database !== this.config.database) {
+      pool = new pg.Pool({
+        host: this.config.host,
+        port: this.config.port || 5432,
+        user: this.config.user,
+        password: this.config.password,
+        database,
+        max: 1,
+      });
+    }
+
+    const client = await pool.connect();
+
+    try {
+      const setClause = Object.entries(updates.set)
+        .map(([key, value]) => `${pg.escapeIdentifier(key)} = ${this.escapeValue(value)}`)
+        .join(', ');
+
+      const query = `UPDATE ${pg.escapeIdentifier(table)} SET ${setClause} WHERE ${updates.where}`;
+      const result = await client.query(query);
+
+      return { affectedRows: result.rowCount || 0 };
+    } finally {
+      client.release();
+      if (pool !== this.pool) await pool.end();
+    }
+  }
+
+  private escapeValue(value: unknown): string {
+    if (value === null) return 'NULL';
+    if (typeof value === 'number') return String(value);
+    if (typeof value === 'boolean') return value ? 'TRUE' : 'FALSE';
+    return `'${String(value).replace(/'/g, "''")}'`;
+  }
+
+  async exportData(
+    table: string,
+    format: 'json' | 'csv',
+    filePath?: string,
+    options?: ExportOptions,
+    database?: string
+  ): Promise<ExportResult> {
+    if (!this.pool) throw new Error('Not connected');
+
+    let pool: pg.Pool = this.pool;
+
+    if (database && database !== this.config.database) {
+      pool = new pg.Pool({
+        host: this.config.host,
+        port: this.config.port || 5432,
+        user: this.config.user,
+        password: this.config.password,
+        database,
+        max: 1,
+      });
+    }
+
+    const client = await pool.connect();
+
+    try {
+      let query = `SELECT * FROM ${pg.escapeIdentifier(table)}`;
+      if (options?.where) {
+        query += ` WHERE ${options.where}`;
+      }
+      if (options?.limit) {
+        query += ` LIMIT ${options.limit}`;
+      }
+
+      const result = await client.query(query);
+      const data = result.rows;
+
+      const targetPath = filePath
+        ? expandTilde(filePath)
+        : getDefaultFilePath(table, format);
+
+      ensureDirectoryExists(targetPath);
+
+      let content: string;
+      if (format === 'json') {
+        content = JSON.stringify(data, null, 2);
+      } else {
+        if (data.length === 0) {
+          content = '';
+        } else {
+          const headers = Object.keys(data[0]);
+          const csvRows = [
+            headers.join(','),
+            ...data.map((row) =>
+              headers.map((h) => {
+                const val = row[h];
+                if (val === null) return '';
+                if (typeof val === 'string') return '"' + val.replace(/"/g, '""') + '"';
+                return String(val);
+              }).join(',')
+            ),
+          ];
+          content = csvRows.join('\n');
+        }
+      }
+
+      fs.writeFileSync(targetPath, content, 'utf-8');
+      const stats = fs.statSync(targetPath);
+
+      return {
+        success: true,
+        filePath: targetPath,
+        rowCount: data.length,
+        fileSize: formatFileSize(stats.size),
+      };
+    } finally {
+      client.release();
+      if (pool !== this.pool) await pool.end();
+    }
+  }
+
+  async createTable(table: string, columns: TableColumnDef[], database?: string): Promise<CreateTableResult> {
+    if (!this.pool) throw new Error('Not connected');
+
+    let pool: pg.Pool = this.pool;
+
+    if (database && database !== this.config.database) {
+      pool = new pg.Pool({
+        host: this.config.host,
+        port: this.config.port || 5432,
+        user: this.config.user,
+        password: this.config.password,
+        database,
+        max: 1,
+      });
+    }
+
+    const client = await pool.connect();
+
+    try {
+      const columnDefs = columns.map((col) => {
+        let def = `${pg.escapeIdentifier(col.name)} ${col.type}`;
+        if (col.nullable === false) def += ' NOT NULL';
+        if (col.primaryKey) def += ' PRIMARY KEY';
+        if (col.defaultValue !== undefined) {
+          def += ` DEFAULT ${this.escapeValue(col.defaultValue)}`;
+        }
+        return def;
+      });
+
+      const query = `CREATE TABLE ${pg.escapeIdentifier(table)} (${columnDefs.join(', ')})`;
+      await client.query(query);
+
+      return { success: true, tableName: table };
+    } finally {
+      client.release();
+      if (pool !== this.pool) await pool.end();
+    }
+  }
+
+  async dropTable(table: string, ifExists = false, database?: string): Promise<DropTableResult> {
+    if (!this.pool) throw new Error('Not connected');
+
+    let pool: pg.Pool = this.pool;
+
+    if (database && database !== this.config.database) {
+      pool = new pg.Pool({
+        host: this.config.host,
+        port: this.config.port || 5432,
+        user: this.config.user,
+        password: this.config.password,
+        database,
+        max: 1,
+      });
+    }
+
+    const client = await pool.connect();
+
+    try {
+      const query = ifExists
+        ? `DROP TABLE IF EXISTS ${pg.escapeIdentifier(table)}`
+        : `DROP TABLE ${pg.escapeIdentifier(table)}`;
+
+      await client.query(query);
+
+      return { success: true, tableName: table };
+    } finally {
+      client.release();
+      if (pool !== this.pool) await pool.end();
+    }
+  }
+
+  async getTableStats(table: string, database?: string): Promise<TableStatsResult> {
+    if (!this.pool) throw new Error('Not connected');
+
+    let pool: pg.Pool = this.pool;
+
+    if (database && database !== this.config.database) {
+      pool = new pg.Pool({
+        host: this.config.host,
+        port: this.config.port || 5432,
+        user: this.config.user,
+        password: this.config.password,
+        database,
+        max: 1,
+      });
+    }
+
+    const client = await pool.connect();
+
+    try {
+      // Get row count
+      const countResult = await client.query(
+        `SELECT COUNT(*) as count FROM ${pg.escapeIdentifier(table)}`
+      );
+      const rowCount = parseInt(countResult.rows[0].count as string, 10);
+
+      // Get column count
+      const colResult = await client.query(
+        `SELECT COUNT(*) as count FROM information_schema.columns WHERE table_name = $1`,
+        [table]
+      );
+      const columns = parseInt(colResult.rows[0].count as string, 10);
+
+      // Get indexes
+      const indexResult = await client.query(
+        `SELECT indexname FROM pg_indexes WHERE tablename = $1 AND schemaname = 'public'`,
+        [table]
+      );
+      const indexes = indexResult.rows.map((row) => row.indexname as string);
+
+      // Get table size
+      const sizeResult = await client.query(
+        `SELECT pg_size_pretty(pg_total_relation_size($1::regclass)) as size`,
+        [table]
+      );
+      const size = sizeResult.rows[0].size as string;
+
+      return {
+        tableName: table,
+        rowCount,
+        columns,
+        indexes,
+        size,
+      };
+    } finally {
+      client.release();
+      if (pool !== this.pool) await pool.end();
+    }
+  }
+
+  async previewData(
+    table: string,
+    page = 1,
+    pageSize = 50,
+    orderBy?: string,
+    database?: string
+  ): Promise<PreviewDataResult> {
+    if (!this.pool) throw new Error('Not connected');
+
+    let pool: pg.Pool = this.pool;
+
+    if (database && database !== this.config.database) {
+      pool = new pg.Pool({
+        host: this.config.host,
+        port: this.config.port || 5432,
+        user: this.config.user,
+        password: this.config.password,
+        database,
+        max: 1,
+      });
+    }
+
+    const client = await pool.connect();
+
+    try {
+      // Get total row count
+      const countResult = await client.query(
+        `SELECT COUNT(*) as count FROM ${pg.escapeIdentifier(table)}`
+      );
+      const totalRows = parseInt(countResult.rows[0].count as string, 10);
+
+      const totalPages = Math.ceil(totalRows / pageSize);
+      const offset = (page - 1) * pageSize;
+
+      // Build query
+      let query = `SELECT * FROM ${pg.escapeIdentifier(table)}`;
+      if (orderBy) {
+        query += ` ORDER BY ${pg.escapeIdentifier(orderBy)}`;
+      }
+      query += ` LIMIT ${pageSize} OFFSET ${offset}`;
+
+      const result = await client.query(query);
+
+      return {
+        rows: result.rows as Record<string, unknown>[],
+        currentPage: page,
+        totalPages,
+        totalRows,
+      };
+    } finally {
+      client.release();
+      if (pool !== this.pool) await pool.end();
+    }
+  }
+
+  async sampleData(table: string, count = 10, database?: string): Promise<SampleDataResult> {
+    if (!this.pool) throw new Error('Not connected');
+
+    let pool: pg.Pool = this.pool;
+
+    if (database && database !== this.config.database) {
+      pool = new pg.Pool({
+        host: this.config.host,
+        port: this.config.port || 5432,
+        user: this.config.user,
+        password: this.config.password,
+        database,
+        max: 1,
+      });
+    }
+
+    const client = await pool.connect();
+
+    try {
+      const result = await client.query(
+        `SELECT * FROM ${pg.escapeIdentifier(table)} ORDER BY RANDOM() LIMIT ${count}`
+      );
+
+      return {
+        rows: result.rows as Record<string, unknown>[],
+        sampleCount: result.rows.length,
+      };
+    } finally {
+      client.release();
+      if (pool !== this.pool) await pool.end();
+    }
+  }
 }