sql-database-plugin

import * as mysql from "mysql2/promise";
import * as pg from "pg";
import { DatabaseConfig } from "./config";

export type Connection =
  | mysql.Connection
  | mysql.Pool
  | pg.PoolClient
  | pg.Pool;

let connectionPool: Connection | null = null;
let currentConfig: DatabaseConfig | null = null;

export async function initializeConnection(
  config: DatabaseConfig
): Promise<void> {
  if (currentConfig && JSON.stringify(currentConfig) === JSON.stringify(config)) {
    return; // Already connected with same config
  }

  // Close existing connection
  if (connectionPool) {
    try {
      if (config.dbType === "mysql") {
        await (connectionPool as mysql.Pool).end();
      } else {
        await (connectionPool as pg.Pool).end();
      }
    } catch (error) {
      console.error("Error closing previous connection:", error);
    }
  }

  currentConfig = config;

  if (config.dbType === "mysql") {
    connectionPool = mysql.createPool({
      host: config.host,
      port: config.port,
      user: config.username,
      password: config.password,
      database: config.database,
      waitForConnections: true,
      connectionLimit: config.connectionPool || 10,
      queueLimit: 0,
    });
  } else {
    connectionPool = new pg.Pool({
      host: config.host,
      port: config.port,
      user: config.username,
      password: config.password,
      database: config.database,
      max: config.connectionPool || 10,
    });
  }

  // Test the connection
  try {
    if (config.dbType === "mysql") {
      const conn = await (connectionPool as mysql.Pool).getConnection();
      await conn.ping();
      conn.release();
    } else {
      const client = await (connectionPool as pg.Pool).connect();
      await client.query("SELECT NOW()");
      client.release();
    }
  } catch (error) {
    connectionPool = null;
    currentConfig = null;
    throw new Error(`Failed to connect to database: ${error}`);
  }
}

export async function query(
  sql: string,
  params: any[] = []
): Promise<any[]> {
  if (!connectionPool) {
    throw new Error("Database connection not initialized");
  }

  try {
    if (currentConfig?.dbType === "mysql") {
      const [rows] = await (connectionPool as mysql.Pool).execute(sql, params);
      return rows as any[];
    } else {
      const result = await (connectionPool as pg.Pool).query(sql, params);
      return result.rows;
    }
  } catch (error) {
    throw new Error(`Query execution failed: ${error}`);
  }
}

export async function execute(
  sql: string,
  params: any[] = []
): Promise<{ affectedRows: number; lastInsertId?: number }> {
  if (!connectionPool) {
    throw new Error("Database connection not initialized");
  }

  try {
    if (currentConfig?.dbType === "mysql") {
      const [result] = await (connectionPool as mysql.Pool).execute(sql, params);
      const mysqlResult = result as any;
      return {
        affectedRows: mysqlResult.affectedRows || 0,
        lastInsertId: mysqlResult.insertId,
      };
    } else {
      const result = await (connectionPool as pg.Pool).query(sql, params);
      return {
        affectedRows: result.rowCount || 0,
      };
    }
  } catch (error) {
    throw new Error(`Execute failed: ${error}`);
  }
}

export async function getSchema(): Promise<any> {
  if (!connectionPool || !currentConfig) {
    throw new Error("Database connection not initialized");
  }

  try {
    if (currentConfig.dbType === "mysql") {
      return await getMySQLSchema();
    } else {
      return await getPostgreSQLSchema();
    }
  } catch (error) {
    throw new Error(`Schema retrieval failed: ${error}`);
  }
}

async function getMySQLSchema(): Promise<any> {
  const tablesResult = await query(
    "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ?",
    [currentConfig!.database]
  );

  const schema: any = {};

  for (const tableRow of tablesResult) {
    const tableName = (tableRow as any).TABLE_NAME;
    const columnsResult = await query(
      "SELECT COLUMN_NAME, COLUMN_TYPE, IS_NULLABLE, COLUMN_KEY FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?",
      [currentConfig!.database, tableName]
    );

    schema[tableName] = columnsResult.map((col: any) => ({
      name: col.COLUMN_NAME,
      type: col.COLUMN_TYPE,
      nullable: col.IS_NULLABLE === "YES",
      key: col.COLUMN_KEY,
    }));
  }

  return schema;
}

async function getPostgreSQLSchema(): Promise<any> {
  const tablesResult = await query(
    `SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'`
  );

  const schema: any = {};

  for (const tableRow of tablesResult) {
    const tableName = (tableRow as any).table_name;
    const columnsResult = await query(
      `SELECT column_name, data_type, is_nullable FROM information_schema.columns WHERE table_schema = 'public' AND table_name = $1`,
      [tableName]
    );

    schema[tableName] = columnsResult.map((col: any) => ({
      name: col.column_name,
      type: col.data_type,
      nullable: col.is_nullable === "YES",
    }));
  }

  return schema;
}

export async function closeConnection(): Promise<void> {
  if (connectionPool) {
    try {
      if (currentConfig?.dbType === "mysql") {
        await (connectionPool as mysql.Pool).end();
      } else {
        await (connectionPool as pg.Pool).end();
      }
    } catch (error) {
      console.error("Error closing connection:", error);
    }
    connectionPool = null;
    currentConfig = null;
  }
}