@dbstudio/cli 0.1.7

package/package.json

@@ -0,0 +1,36 @@
+{
+	"name": "@dbstudio/cli",
+	"version": "0.1.7",
+	"type": "module",
+	"bin": {
+		"dbstudio": "./dist/index.js"
+	},
+	"scripts": {
+		"dev": "tsup --watch",
+		"build": "tsup",
+		"start": "bun run dist/index.js"
+	},
+	"dependencies": {
+		"@dbstudio/types": "workspace:*",
+		"@libsql/client": "^0.15.15",
+		"chalk": "^5.4.1",
+		"commander": "^13.1.0",
+		"dotenv": "^16.1.4",
+		"mysql2": "^3.14.1",
+		"ora": "^8.2.0",
+		"pg": "^8.16.0",
+		"prompts": "^2.4.2",
+		"ssh2": "^1.17.0",
+		"ws": "^8.18.0"
+	},
+	"devDependencies": {
+		"@types/better-sqlite3": "^7.6.13",
+		"@types/pg": "^8.15.2",
+		"@types/prompts": "^2.4.9",
+		"@types/ssh2": "^1.15.5",
+		"@types/ws": "^8.5.13",
+		"tsdown": "^0.11.4",
+		"tsup": "^8.5.1",
+		"typescript": "^5"
+	}
+}

package/src/agents/index.ts

@@ -0,0 +1,458 @@
+import fs from "node:fs";
+import { createServer, type Server } from "node:net";
+import os from "node:os";
+import path from "node:path";
+import type {
+	AuthMessage,
+	DatabaseConfig,
+	GetTableSchemaMessage,
+	GetTablesMessage,
+	InsertRowMessage,
+	QueryMessage,
+	UpdateRowMessage,
+	WSMessage,
+} from "@dbstudio/types";
+import { Client } from "ssh2";
+import WebSocket from "ws";
+import { createDriver, type DatabaseDriver } from "../drivers";
+
+const CONFIG_DIR = path.join(os.homedir(), ".dbstudio");
+const STATUS_FILE = path.join(CONFIG_DIR, "status.json");
+
+interface AgentConfig {
+	serverUrl: string;
+	token: string;
+	dbConfig: DatabaseConfig;
+}
+
+export class Agent {
+	private ws: WebSocket | null = null;
+	private driver: DatabaseDriver | null = null;
+	private config: AgentConfig;
+	private reconnectAttempts = 0;
+	private maxReconnectAttempts = 5;
+	private pingInterval: ReturnType<typeof setInterval> | null = null;
+	private sshClient: Client | null = null;
+	private tunnelServer: Server | null = null;
+	private isReconnecting = false;
+	private reconnectTimeout: ReturnType<typeof setTimeout> | null = null;
+
+	constructor(config: AgentConfig) {
+		this.config = config;
+		this.ensureConfigDir();
+	}
+
+	private ensureConfigDir(): void {
+		if (!fs.existsSync(CONFIG_DIR)) {
+			fs.mkdirSync(CONFIG_DIR, { recursive: true });
+		}
+	}
+
+	async connect(): Promise<void> {
+		// Set up SSH tunnel if configured
+		if (this.config.dbConfig.ssh) {
+			await this.setupSshTunnel();
+		}
+
+		// Connect to database first
+		this.driver = createDriver(this.config.dbConfig);
+		await this.driver.connect();
+
+		// Connect to WebSocket server
+		await this.connectWebSocket();
+
+		// Save status
+		this.saveStatus("connected");
+	}
+
+	private async setupSshTunnel(): Promise<void> {
+		const { ssh, host, port, type } = this.config.dbConfig;
+		if (!ssh) return;
+
+		return new Promise((resolve, reject) => {
+			this.sshClient = new Client();
+
+			this.sshClient
+				.on("ready", () => {
+					// Create local server to forward traffic
+					this.tunnelServer = createServer((sock) => {
+						this.sshClient?.forwardOut(
+							"127.0.0.1",
+							sock.remotePort || 0,
+							host || "localhost",
+							port || (type === "postgresql" ? 5432 : 3306),
+							(err, stream) => {
+								if (err) {
+									sock.end();
+									return;
+								}
+								sock.pipe(stream).pipe(sock);
+							},
+						);
+					});
+
+					this.tunnelServer.listen(0, "127.0.0.1", () => {
+						const addr = this.tunnelServer?.address();
+						if (addr && typeof addr !== "string") {
+							// Update config to point to local tunnel
+							this.config.dbConfig.host = "127.0.0.1";
+							this.config.dbConfig.port = addr.port;
+							console.log(
+								`SSH Tunnel established: 127.0.0.1:${addr.port} -> ${host}:${port || (type === "postgresql" ? 5432 : 3306)}`,
+							);
+							resolve();
+						} else {
+							reject(new Error("Failed to get tunnel address"));
+						}
+					});
+
+					this.tunnelServer.on("error", reject);
+				})
+				.on("error", reject)
+				.connect({
+					host: ssh.host,
+					port: ssh.port || 22,
+					username: ssh.username,
+					password: ssh.password,
+					privateKey: ssh.privateKey,
+					passphrase: ssh.passphrase,
+				});
+		});
+	}
+
+	private async connectWebSocket(): Promise<void> {
+		return new Promise((resolve, reject) => {
+			this.ws = new WebSocket(this.config.serverUrl);
+
+			this.ws.on("open", () => {
+				console.log("WebSocket connected, authenticating...");
+				this.authenticate();
+				this.startPingInterval();
+				resolve();
+			});
+
+			this.ws.on("message", (data) => {
+				this.handleMessage(data.toString());
+			});
+
+			this.ws.on("close", () => {
+				console.log("WebSocket closed");
+				this.stopPingInterval();
+				this.handleReconnect();
+			});
+
+			this.ws.on("error", (error) => {
+				console.error("WebSocket error:", error.message);
+				reject(error);
+			});
+		});
+	}
+
+	private authenticate(): void {
+		const authMessage: AuthMessage = {
+			id: crypto.randomUUID(),
+			type: "auth",
+			timestamp: Date.now(),
+			payload: {
+				token: this.config.token,
+				connectionId: crypto.randomUUID(),
+				dbConfig: this.config.dbConfig,
+			},
+		};
+		this.send(authMessage);
+	}
+
+	private async handleMessage(data: string): Promise<void> {
+		try {
+			const message = JSON.parse(data) as WSMessage;
+
+			switch (message.type) {
+				case "auth_success":
+					console.log("Authenticated successfully");
+					break;
+
+				case "auth_error":
+					console.error("Authentication failed:", message.payload.error);
+					await this.disconnect();
+					break;
+
+				case "ping":
+					this.send({ id: message.id, type: "pong", timestamp: Date.now() });
+					break;
+
+				case "pong":
+					// Received pong from server, nothing to do
+					break;
+
+				case "query":
+					await this.handleQuery(message as QueryMessage);
+					break;
+
+				case "get_tables":
+					await this.handleGetTables(message as GetTablesMessage);
+					break;
+
+				case "get_table_schema":
+					await this.handleGetTableSchema(message as GetTableSchemaMessage);
+					break;
+
+				case "insert_row":
+					await this.handleInsertRow(message as InsertRowMessage);
+					break;
+
+				case "update_row":
+					await this.handleUpdateRow(message as UpdateRowMessage);
+					break;
+
+				default:
+					console.log("Unknown message type:", message.type);
+			}
+		} catch (error) {
+			console.error("Error handling message:", error);
+		}
+	}
+
+	private async handleQuery(message: QueryMessage): Promise<void> {
+		if (!this.driver) return;
+
+		try {
+			const result = await this.driver.query(
+				message.payload.sql,
+				message.payload.params,
+			);
+			this.send({
+				id: message.id,
+				type: "query_result",
+				timestamp: Date.now(),
+				payload: result,
+			});
+		} catch (error) {
+			this.send({
+				id: message.id,
+				type: "query_error",
+				timestamp: Date.now(),
+				payload: {
+					message: error instanceof Error ? error.message : "Query failed",
+					code: (error as any)?.code,
+				},
+			});
+		}
+	}
+
+	private async handleGetTables(message: GetTablesMessage): Promise<void> {
+		if (!this.driver) return;
+
+		try {
+			const tables = await this.driver.getTables(message.payload.schema);
+			this.send({
+				id: message.id,
+				type: "tables_result",
+				timestamp: Date.now(),
+				payload: { tables },
+			});
+		} catch (error) {
+			console.error("Error in handleGetTables:", error);
+			this.send({
+				id: message.id,
+				type: "query_error",
+				timestamp: Date.now(),
+				payload: {
+					message:
+						error instanceof Error ? error.message : "Failed to get tables",
+				},
+			});
+		}
+	}
+
+	private async handleGetTableSchema(
+		message: GetTableSchemaMessage,
+	): Promise<void> {
+		if (!this.driver) return;
+
+		try {
+			const schema = await this.driver.getTableSchema(
+				message.payload.table,
+				message.payload.schema,
+			);
+			this.send({
+				id: message.id,
+				type: "table_schema_result",
+				timestamp: Date.now(),
+				payload: schema,
+			});
+		} catch (error) {
+			this.send({
+				id: message.id,
+				type: "query_error",
+				timestamp: Date.now(),
+				payload: {
+					message:
+						error instanceof Error ? error.message : "Failed to get schema",
+				},
+			});
+		}
+	}
+
+	private async handleInsertRow(message: InsertRowMessage): Promise<void> {
+		if (!this.driver) return;
+
+		try {
+			const result = await this.driver.insertRow(
+				message.payload.table,
+				message.payload.data,
+			);
+			this.send({
+				id: message.id,
+				type: "query_result",
+				timestamp: Date.now(),
+				payload: result,
+			});
+		} catch (error) {
+			this.send({
+				id: message.id,
+				type: "query_error",
+				timestamp: Date.now(),
+				payload: {
+					message:
+						error instanceof Error ? error.message : "Failed to insert row",
+				},
+			});
+		}
+	}
+
+	private async handleUpdateRow(message: UpdateRowMessage): Promise<void> {
+		if (!this.driver) return;
+
+		try {
+			const result = await this.driver.updateRow(
+				message.payload.table,
+				message.payload.data,
+				message.payload.where,
+			);
+			this.send({
+				id: message.id,
+				type: "query_result",
+				timestamp: Date.now(),
+				payload: result,
+			});
+		} catch (error) {
+			this.send({
+				id: message.id,
+				type: "query_error",
+				timestamp: Date.now(),
+				payload: {
+					message:
+						error instanceof Error ? error.message : "Failed to update row",
+				},
+			});
+		}
+	}
+
+	private send(message: WSMessage | Record<string, unknown>): void {
+		if (this.ws?.readyState === WebSocket.OPEN) {
+			this.ws.send(JSON.stringify(message));
+		}
+	}
+
+	private startPingInterval(): void {
+		this.pingInterval = setInterval(() => {
+			this.send({
+				id: crypto.randomUUID(),
+				type: "ping",
+				timestamp: Date.now(),
+			});
+		}, 30000);
+	}
+
+	private stopPingInterval(): void {
+		if (this.pingInterval) {
+			clearInterval(this.pingInterval);
+			this.pingInterval = null;
+		}
+	}
+
+	private clearReconnect(): void {
+		if (this.reconnectTimeout) {
+			clearTimeout(this.reconnectTimeout);
+			this.reconnectTimeout = null;
+		}
+		this.isReconnecting = false;
+	}
+
+	private async handleReconnect(): Promise<void> {
+		if (this.isReconnecting) return;
+		this.isReconnecting = true;
+
+		if (this.reconnectAttempts >= this.maxReconnectAttempts) {
+			console.error("Max reconnect attempts reached");
+			this.saveStatus("error");
+			this.isReconnecting = false;
+			return;
+		}
+
+		this.reconnectAttempts++;
+		const delay = Math.min(1000 * 2 ** this.reconnectAttempts, 30000);
+		console.log(`Reconnecting in ${delay / 1000}s...`);
+
+		this.saveStatus("connecting");
+
+		if (this.reconnectTimeout) clearTimeout(this.reconnectTimeout);
+		this.reconnectTimeout = setTimeout(async () => {
+			try {
+				await this.connectWebSocket();
+				this.reconnectAttempts = 0;
+				this.saveStatus("connected");
+			} catch (_error) {
+				// Failed again, allow next reconnect trigger
+			} finally {
+				this.isReconnecting = false;
+				this.reconnectTimeout = null;
+			}
+		}, delay);
+	}
+
+	async disconnect(): Promise<void> {
+		this.stopPingInterval();
+		this.clearReconnect();
+
+		if (this.ws) {
+			this.ws.close();
+			this.ws = null;
+		}
+
+		if (this.driver) {
+			await this.driver.disconnect();
+			this.driver = null;
+		}
+
+		if (this.tunnelServer) {
+			this.tunnelServer.close();
+			this.tunnelServer = null;
+		}
+
+		if (this.sshClient) {
+			this.sshClient.end();
+			this.sshClient = null;
+		}
+
+		this.saveStatus("disconnected");
+		this.removeStatus();
+	}
+
+	private saveStatus(status: string): void {
+		const statusData = {
+			status,
+			database: this.config.dbConfig.database || this.config.dbConfig.filepath,
+			type: this.config.dbConfig.type,
+			connectedAt: Date.now(),
+			pid: process.pid,
+		};
+		fs.writeFileSync(STATUS_FILE, JSON.stringify(statusData, null, 2));
+	}
+
+	private removeStatus(): void {
+		if (fs.existsSync(STATUS_FILE)) {
+			fs.unlinkSync(STATUS_FILE);
+		}
+	}
+}