@salmanorg/db-agent/index.js

Database agent for connecting to multiple databases via SSH tunneling

require('dotenv').config({ path: require('path').resolve(process.cwd(), '.env') });

const mysql = require('mysql2/promise');
const sqlServer = require('mssql');
const { Client } = require('pg');
const { MongoClient } = require('mongodb');
const oracledb = require('oracledb');
const axios = require('axios');
const os = require('os');
const { io } = require('socket.io-client');

const config = {
  dbType: process.env.DB_TYPE || 'mysql',
  connectionUrl: process.env.URL ? decodeURIComponent(process.env.URL) : null,
  dbHost: process.env.DB_HOST || 'localhost',
  dbPort: process.env.DB_PORT,
  dbName: process.env.DB_NAME,
  dbUser: process.env.DB_USER,
  dbPassword: process.env.DB_PASSWORD,
  apiEndpoint: process.env.API_ENDPOINT || 'http://main-app:3000/api/report',
  wsEndpoint: process.env.WS_ENDPOINT || 'http://main-app:3000',
  agentId: process.env.AGENT_ID || os.hostname(),
  agentToken: process.env.AGENT_TOKEN,
  activationKey: process.env.ACTIVATION_KEY,
  debug: process.env.DEBUG === 'true',
};

let globalDbConnection = null;
let globalDbType = null;

function log(...args) {
  if (config.debug) console.log('[DEBUG]', ...args);
}

function validateEnv() {
  const required = ['AGENT_TOKEN', 'ACTIVATION_KEY'];
  const missing = required.filter((k) => !process.env[k]);
  if (missing.length) {
    console.error('❌ Missing required environment variables:', missing.join(', '));
    process.exit(1);
  }
}

async function updateAgentStatus(updatePayload) {
  const statusUrl = config.apiEndpoint.replace('/report', `/${config.agentId}/status`);
  try {
    await axios.put(statusUrl, updatePayload, {
      headers: { Authorization: `Bearer ${config.agentToken}` },
    });
    log('✅ Agent status updated:', updatePayload);
  } catch (error) {
    console.error('❌ Failed to update agent status:', error.message);
  }
}

async function activateAgent() {
  const url = config.apiEndpoint.replace('/report', '/activate');
  log('Activating agent at:', url);

  try {
    const response = await axios.post(url, { agentId: config.agentId, activationKey: config.activationKey }, { headers: { Authorization: `Bearer ${config.agentToken}` } });

    if (response.data.status !== 'activated') {
      console.error('❌ Activation failed - unexpected response:', response.data);
      process.exit(1);
    }

    console.log('✅ Agent successfully activated');
  } catch (error) {
    console.error('❌ Activation error');
    if (error.response) {
      console.error('Status:', error.response.status);
      console.error('Response:', JSON.stringify(error.response.data, null, 2));
    } else {
      console.error('Error:', error.message);
    }
    process.exit(1);
  }
}

async function updateConfiguredFlag() {
  await updateAgentStatus({ isConfigured: true });
}

async function runCustomQuery(query) {
  const { connection } = await connectToDatabase();
  const dbType = config.dbType.toLowerCase();
  let result;

  try {
    switch (dbType) {
      case 'postgres':
      case 'postgresql':
        result = await connection.query(query);
        result = result.rows;
        break;
      case 'mysql':
        [result] = await connection.query(query);
        break;
      case 'mssql':
        result = await connection.query(query);
        result = result.recordset;
        break;
      case 'oracle':
        const response = await connection.execute(query);
        result = response.rows;
        break;
      case 'mongodb':
        result = { error: 'Raw queries not supported in MongoDB this way' };
        break;
      default:
        result = { error: 'Unsupported database type' };
    }

    await sendDataToMainApp({ command: 'run', result });
    await updateAgentStatus({ status: 'ACTIVE' });
  } catch (error) {
    console.error('❌ Custom query error:', error.message);
    await updateAgentStatus({ status: 'FAILED' });
    await sendDataToMainApp({ command: 'run', error: error.message });
  }
}

function connectWebSocket() {
  const socket = io(`${config.wsEndpoint}/agent`, {
    path: '/socket.io',
    transports: ['websocket'],
    query: { token: config.agentToken },
  });

  socket.on('connect', () => {
    console.log('✅ Connected to WebSocket server');
    socket.emit('register', { agentId: config.agentId });
  });

  socket.on('connect_error', (err) => {
    console.error('❌ WebSocket connection failed:', err.message);
  });

  socket.on('command', async (message) => {
    try {
      console.log('🟡 Received command:', message.command);
      if (message.command === 'query') {
        await runAgent();
      } else if (message.command === 'stop') {
        console.log('🛑 Received stop command, exiting.');
        await updateAgentStatus({ status: 'STOPPED' });
        process.exit(0);
      } else if (message.command.startsWith('run:')) {
        const rawQuery = message.command.slice(4).trim();
        await runCustomQuery(rawQuery);
      }
    } catch (error) {
      console.error('Command error:', error.message);
      await updateAgentStatus({ status: 'FAILED' });
    }
  });

  socket.on('disconnect', () => {
    console.warn('⚠️  Socket.IO disconnected. Reconnecting...');
    updateAgentStatus({ status: 'DISCONNECTED' });
    setTimeout(connectWebSocket, 5000);
  });
}

async function connectToDatabase() {
  if (globalDbConnection) {
    return { connection: globalDbConnection };
  }

  const dbType = config.dbType.toLowerCase();
  globalDbType = dbType;
  log('Connecting to DB:', dbType);

  try {
    let connection;

    if (config.connectionUrl) {
      switch (dbType) {
        case 'mysql':
          connection = await mysql.createConnection(config.connectionUrl);
          break;
        case 'postgres':
        case 'postgresql':
          connection = new Client({ connectionString: config.connectionUrl });
          await connection.connect();
          break;
        case 'mongodb':
          connection = await MongoClient.connect(config.connectionUrl, {
            useNewUrlParser: true,
            useUnifiedTopology: true,
          });
          break;
        case 'mssql':
          connection = await sqlServer.connect(config.connectionUrl);
          break;
        case 'oracle':
          connection = await oracledb.getConnection({ connectString: config.connectionUrl });
          break;
        default:
          throw new Error('Unsupported database type');
      }
    } else {
      switch (dbType) {
        case 'mysql':
          connection = await mysql.createConnection({
            host: config.dbHost,
            port: config.dbPort || 3306,
            user: config.dbUser,
            password: config.dbPassword,
            database: config.dbName,
          });
          break;
        case 'postgres':
        case 'postgresql':
          connection = new Client({
            host: config.dbHost,
            port: config.dbPort || 5432,
            user: config.dbUser,
            password: config.dbPassword,
            database: config.dbName,
          });
          await connection.connect();
          break;
        case 'mongodb':
          connection = await MongoClient.connect(`mongodb://${config.dbHost}:${config.dbPort || 27017}/${config.dbName}`, {
            useNewUrlParser: true,
            useUnifiedTopology: true,
          });
          break;
        case 'mssql':
          connection = await sqlServer.connect({
            server: config.dbHost,
            port: parseInt(config.dbPort) || 1433,
            user: config.dbUser,
            password: config.dbPassword,
            database: config.dbName,
            options: { encrypt: true },
          });
          break;
        case 'oracle':
          connection = await oracledb.getConnection({
            user: config.dbUser,
            password: config.dbPassword,
            connectString: `${config.dbHost}:${config.dbPort || 1521}/${config.dbName}`,
          });
          break;
        default:
          throw new Error('Unsupported database type');
      }
    }

    globalDbConnection = connection;
    await updateConfiguredFlag();
    return { connection };
  } catch (error) {
    console.error('❌ Database connection error:', error.message);
    await updateAgentStatus({ status: 'FAILED' });
    throw error;
  }
}

async function sendDataToMainApp(data) {
  try {
    await axios.post(
      config.apiEndpoint,
      { agentId: config.agentId, data },
      {
        headers: { Authorization: `Bearer ${config.agentToken}` },
      },
    );
    console.log('✅ Data sent to main application');
  } catch (error) {
    console.error('❌ Error sending data:');
    if (error.response) {
      console.error('Status:', error.response.status);
      console.error('Response:', error.response.data);
    } else {
      console.error('Message:', error.message);
    }
  }
}

async function runAgent() {
  let connection;
  const dbType = config.dbType.toLowerCase();
  try {
    const result = await connectToDatabase();
    connection = result.connection;

    let resultData;
    switch (dbType) {
      case 'mysql':
        [resultData] = await connection.query('SELECT 1 AS test');
        break;
      case 'postgres':
      case 'postgresql':
        resultData = await connection.query('SELECT 1 AS test');
        resultData = resultData.rows;
        break;
      case 'mongodb':
        resultData = await connection.db(config.dbName).listCollections().toArray();
        break;
      case 'mssql':
        resultData = await connection.query('SELECT 1 AS test');
        resultData = resultData.recordset;
        break;
      case 'oracle':
        resultData = await connection.execute('SELECT 1 AS test FROM dual');
        resultData = resultData.rows;
        break;
      default:
        throw new Error('Unsupported DB type for query');
    }

    await sendDataToMainApp(resultData);
    await updateAgentStatus({ status: 'ACTIVE' });
  } catch (error) {
    console.error('❌ Agent runtime error:', error.message);
    await updateAgentStatus({ status: 'FAILED' });
  }
}

process.on('SIGINT', async () => {
  console.log('Cleaning up before exit...');
  if (globalDbConnection) {
    try {
      if (['mysql', 'postgresql', 'postgres'].includes(globalDbType)) await globalDbConnection.end();
      if (['mssql', 'oracle'].includes(globalDbType)) await globalDbConnection.close();
    } catch (err) {
      console.warn('⚠️ Error closing DB connection:', err.message);
    }
  }
  await updateAgentStatus({ status: 'STOPPED' });
  process.exit(0);
});

async function main() {
  validateEnv();
  await activateAgent();
  connectWebSocket();
  setInterval(runAgent, 60000);
  await runAgent();
}

main().catch((error) => {
  console.error('❌ Agent failed to start:', error.message);
  updateAgentStatus({ status: 'FAILED' });
  process.exit(1);
});