atoo-studio 0.0.1 → 0.0.2

package/dist/src/database/drivers/postgresql.js

@@ -0,0 +1,133 @@
+import pg from 'pg';
+export class PostgreSQLDriver {
+    pool = null;
+    async connect(params) {
+        if (params.connection_string) {
+            this.pool = new pg.Pool({ connectionString: params.connection_string, max: 5 });
+        }
+        else {
+            this.pool = new pg.Pool({
+                host: params.host || 'localhost',
+                port: params.port || 5432,
+                user: params.username || 'postgres',
+                password: params.password || '',
+                database: params.database || 'postgres',
+                max: 5,
+            });
+        }
+        // Test connection
+        const client = await this.pool.connect();
+        client.release();
+    }
+    async disconnect() {
+        if (this.pool) {
+            await this.pool.end();
+            this.pool = null;
+        }
+    }
+    isConnected() {
+        return this.pool !== null;
+    }
+    async query(sql, limit = 100) {
+        if (!this.pool)
+            throw new Error('Not connected');
+        const start = Date.now();
+        const result = await this.pool.query(sql);
+        const elapsed = Date.now() - start;
+        const rows = result.rows || [];
+        const columns = result.fields?.map(f => f.name) || (rows.length > 0 ? Object.keys(rows[0]) : []);
+        const truncated = rows.length > limit;
+        return {
+            columns,
+            rows: rows.slice(0, limit),
+            row_count: result.rowCount ?? rows.length,
+            execution_time_ms: elapsed,
+            truncated,
+        };
+    }
+    async getTables() {
+        if (!this.pool)
+            throw new Error('Not connected');
+        const result = await this.pool.query(`
+      SELECT
+        t.table_name AS name,
+        t.table_type AS type,
+        pg_total_relation_size(quote_ident(t.table_schema) || '.' || quote_ident(t.table_name))::bigint AS size_bytes,
+        s.n_live_tup AS row_count
+      FROM information_schema.tables t
+      LEFT JOIN pg_stat_user_tables s ON s.relname = t.table_name AND s.schemaname = t.table_schema
+      WHERE t.table_schema NOT IN ('information_schema', 'pg_catalog')
+      ORDER BY t.table_name
+    `);
+        return result.rows.map(r => ({
+            name: r.name,
+            type: r.type === 'BASE TABLE' ? 'table' : r.type?.toLowerCase(),
+            row_count: parseInt(r.row_count) || undefined,
+            size_bytes: parseInt(r.size_bytes) || undefined,
+        }));
+    }
+    async describeTable(table) {
+        if (!this.pool)
+            throw new Error('Not connected');
+        // Columns
+        const cols = await this.pool.query(`
+      SELECT
+        c.column_name, c.data_type, c.is_nullable, c.column_default,
+        CASE WHEN pk.column_name IS NOT NULL THEN true ELSE false END AS is_pk
+      FROM information_schema.columns c
+      LEFT JOIN (
+        SELECT ku.column_name
+        FROM information_schema.table_constraints tc
+        JOIN information_schema.key_column_usage ku ON tc.constraint_name = ku.constraint_name
+        WHERE tc.table_name = $1 AND tc.constraint_type = 'PRIMARY KEY'
+      ) pk ON pk.column_name = c.column_name
+      WHERE c.table_name = $1 AND c.table_schema NOT IN ('information_schema', 'pg_catalog')
+      ORDER BY c.ordinal_position
+    `, [table]);
+        const columns = cols.rows.map(r => ({
+            name: r.column_name,
+            type: r.data_type,
+            nullable: r.is_nullable === 'YES',
+            default_value: r.column_default || undefined,
+            primary_key: r.is_pk,
+        }));
+        // Indexes
+        const idxResult = await this.pool.query(`
+      SELECT indexname AS name, indexdef
+      FROM pg_indexes
+      WHERE tablename = $1 AND schemaname NOT IN ('information_schema', 'pg_catalog')
+    `, [table]);
+        const indexes = idxResult.rows.map(r => ({
+            name: r.name,
+            columns: (r.indexdef.match(/\(([^)]+)\)/)?.[1] || '').split(',').map((s) => s.trim()),
+            unique: /UNIQUE/i.test(r.indexdef),
+        }));
+        // Foreign keys
+        const fkResult = await this.pool.query(`
+      SELECT
+        kcu.column_name AS column,
+        ccu.table_name AS ref_table,
+        ccu.column_name AS ref_column
+      FROM information_schema.table_constraints tc
+      JOIN information_schema.key_column_usage kcu ON tc.constraint_name = kcu.constraint_name
+      JOIN information_schema.constraint_column_usage ccu ON tc.constraint_name = ccu.constraint_name
+      WHERE tc.table_name = $1 AND tc.constraint_type = 'FOREIGN KEY'
+    `, [table]);
+        return {
+            table,
+            columns,
+            indexes,
+            foreign_keys: fkResult.rows.map(r => ({
+                column: r.column,
+                ref_table: r.ref_table,
+                ref_column: r.ref_column,
+            })),
+        };
+    }
+    async getDatabases() {
+        if (!this.pool)
+            throw new Error('Not connected');
+        const result = await this.pool.query(`SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname`);
+        return result.rows.map(r => r.datname);
+    }
+}

package/dist/src/database/drivers/redis.d.ts

@@ -0,0 +1,11 @@
+import type { ConnectionParams, DatabaseDriver, QueryResult, TableInfo, SchemaInfo } from '../types.js';
+export declare class RedisDriver implements DatabaseDriver {
+    private client;
+    connect(params: ConnectionParams): Promise<void>;
+    disconnect(): Promise<void>;
+    isConnected(): boolean;
+    query(command: string, limit?: number): Promise<QueryResult>;
+    getTables(): Promise<TableInfo[]>;
+    describeTable(pattern: string): Promise<SchemaInfo>;
+    getDatabases(): Promise<string[]>;
+}

package/dist/src/database/drivers/redis.js

@@ -0,0 +1,91 @@
+import RedisLib from 'ioredis';
+const Redis = RedisLib.Redis || RedisLib;
+export class RedisDriver {
+    client = null;
+    async connect(params) {
+        if (params.connection_string) {
+            this.client = new Redis(params.connection_string, { lazyConnect: true });
+        }
+        else {
+            this.client = new Redis({
+                host: params.host || 'localhost',
+                port: params.port || 6379,
+                password: params.password || undefined,
+                db: params.database ? parseInt(params.database) : 0,
+                lazyConnect: true,
+            });
+        }
+        await this.client.connect();
+    }
+    async disconnect() {
+        if (this.client) {
+            this.client.disconnect();
+            this.client = null;
+        }
+    }
+    isConnected() {
+        return this.client !== null && this.client.status === 'ready';
+    }
+    async query(command, limit = 100) {
+        if (!this.client)
+            throw new Error('Not connected');
+        const start = Date.now();
+        const parts = command.trim().split(/\s+/);
+        const cmd = parts[0].toUpperCase();
+        const args = parts.slice(1);
+        const result = await this.client.call(cmd, ...args);
+        const elapsed = Date.now() - start;
+        // Format result based on type
+        if (Array.isArray(result)) {
+            // KEYS, SCAN results, etc
+            const rows = result.slice(0, limit).map((v, i) => ({ index: i, value: typeof v === 'object' ? JSON.stringify(v) : String(v) }));
+            return { columns: ['index', 'value'], rows, row_count: result.length, execution_time_ms: elapsed, truncated: result.length > limit };
+        }
+        if (result !== null && typeof result === 'object') {
+            const entries = Object.entries(result).slice(0, limit);
+            const rows = entries.map(([k, v]) => ({ key: k, value: String(v) }));
+            return { columns: ['key', 'value'], rows, row_count: entries.length, execution_time_ms: elapsed };
+        }
+        return {
+            columns: ['result'],
+            rows: [{ result: result === null ? '(nil)' : String(result) }],
+            row_count: 1,
+            execution_time_ms: elapsed,
+        };
+    }
+    async getTables() {
+        if (!this.client)
+            throw new Error('Not connected');
+        const dbSize = await this.client.dbsize();
+        // Get a sample of keys
+        const keys = await this.client.keys('*');
+        const sample = keys.slice(0, 100);
+        // Group by prefix (before first :)
+        const prefixes = new Map();
+        for (const key of sample) {
+            const prefix = key.includes(':') ? key.split(':')[0] + ':*' : key;
+            prefixes.set(prefix, (prefixes.get(prefix) || 0) + 1);
+        }
+        return [
+            { name: `(${dbSize} total keys)`, type: 'keyspace', row_count: dbSize },
+            ...[...prefixes.entries()].map(([name, count]) => ({ name, type: 'pattern', row_count: count })),
+        ];
+    }
+    async describeTable(pattern) {
+        if (!this.client)
+            throw new Error('Not connected');
+        // For Redis, "describe" a key pattern by sampling
+        const keys = await this.client.keys(pattern.replace('*', '') + '*');
+        const sample = keys.slice(0, 10);
+        const columns = [];
+        for (const key of sample) {
+            const type = await this.client.type(key);
+            const ttl = await this.client.ttl(key);
+            columns.push({ name: key, type: `${type} (TTL: ${ttl === -1 ? 'none' : ttl + 's'})`, nullable: false });
+        }
+        return { table: pattern, columns, indexes: [], foreign_keys: [] };
+    }
+    async getDatabases() {
+        return Array.from({ length: 16 }, (_, i) => String(i));
+    }
+}

package/dist/src/database/drivers/sqlite.d.ts

@@ -0,0 +1,10 @@
+import type { ConnectionParams, DatabaseDriver, QueryResult, TableInfo, SchemaInfo } from '../types.js';
+export declare class SQLiteDriver implements DatabaseDriver {
+    private db;
+    connect(params: ConnectionParams): Promise<void>;
+    disconnect(): Promise<void>;
+    isConnected(): boolean;
+    query(sql: string, limit?: number): Promise<QueryResult>;
+    getTables(): Promise<TableInfo[]>;
+    describeTable(table: string): Promise<SchemaInfo>;
+}

package/dist/src/database/drivers/sqlite.js

@@ -0,0 +1,100 @@
+import Database from 'better-sqlite3';
+export class SQLiteDriver {
+    db = null;
+    async connect(params) {
+        const file = params.filename || params.database;
+        if (!file)
+            throw new Error('filename is required for SQLite');
+        const isReadonly = params._readonly === true;
+        this.db = new Database(file, { readonly: isReadonly });
+        if (!isReadonly) {
+            this.db.pragma('journal_mode = WAL');
+        }
+    }
+    async disconnect() {
+        if (this.db) {
+            this.db.close();
+            this.db = null;
+        }
+    }
+    isConnected() {
+        return this.db !== null;
+    }
+    async query(sql, limit = 100) {
+        if (!this.db)
+            throw new Error('Not connected');
+        const start = Date.now();
+        const trimmed = sql.trim();
+        const isSelect = /^(SELECT|PRAGMA|EXPLAIN|WITH)\b/i.test(trimmed);
+        if (isSelect) {
+            const stmt = this.db.prepare(sql);
+            const rows = stmt.all();
+            const elapsed = Date.now() - start;
+            const columns = rows.length > 0 ? Object.keys(rows[0]) : [];
+            const truncated = rows.length > limit;
+            return {
+                columns,
+                rows: rows.slice(0, limit),
+                row_count: rows.length,
+                execution_time_ms: elapsed,
+                truncated,
+            };
+        }
+        else {
+            const result = this.db.prepare(sql).run();
+            const elapsed = Date.now() - start;
+            return {
+                columns: ['changes', 'lastInsertRowid'],
+                rows: [{ changes: result.changes, lastInsertRowid: Number(result.lastInsertRowid) }],
+                row_count: result.changes,
+                execution_time_ms: elapsed,
+            };
+        }
+    }
+    async getTables() {
+        if (!this.db)
+            throw new Error('Not connected');
+        const tables = this.db.prepare(`
+      SELECT name, type FROM sqlite_master
+      WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%'
+      ORDER BY name
+    `).all();
+        return tables.map(t => {
+            let row_count;
+            try {
+                const cnt = this.db.prepare(`SELECT COUNT(*) as c FROM "${t.name}"`).get();
+                row_count = cnt?.c;
+            }
+            catch { }
+            return { name: t.name, type: t.type, row_count };
+        });
+    }
+    async describeTable(table) {
+        if (!this.db)
+            throw new Error('Not connected');
+        const cols = this.db.prepare(`PRAGMA table_info("${table}")`).all();
+        const columns = cols.map(c => ({
+            name: c.name,
+            type: c.type,
+            nullable: !c.notnull,
+            default_value: c.dflt_value || undefined,
+            primary_key: !!c.pk,
+        }));
+        const idxList = this.db.prepare(`PRAGMA index_list("${table}")`).all();
+        const indexes = idxList.map(idx => {
+            const idxInfo = this.db.prepare(`PRAGMA index_info("${idx.name}")`).all();
+            return {
+                name: idx.name,
+                columns: idxInfo.map(i => i.name),
+                unique: !!idx.unique,
+            };
+        });
+        const fks = this.db.prepare(`PRAGMA foreign_key_list("${table}")`).all();
+        const foreign_keys = fks.map(fk => ({
+            column: fk.from,
+            ref_table: fk.table,
+            ref_column: fk.to,
+        }));
+        return { table, columns, indexes, foreign_keys };
+    }
+}

package/dist/src/database/query-stream.d.ts

@@ -0,0 +1,5 @@
+import type { WebSocketServer } from 'ws';
+import type { IncomingMessage } from 'http';
+import type { Duplex } from 'stream';
+export declare function isDatabaseWsUpgrade(url: string): boolean;
+export declare function handleDatabaseWsUpgrade(wss: WebSocketServer, req: IncomingMessage, socket: Duplex, head: Buffer): void;

package/dist/src/database/query-stream.js

@@ -0,0 +1,75 @@
+import { connectionManager } from './connection-manager.js';
+const WS_PATH_RE = /^\/ws\/database-query\/([^/?]+)/;
+const BATCH_SIZE = 100;
+const KEEPALIVE_INTERVAL_MS = 30_000;
+export function isDatabaseWsUpgrade(url) {
+    return WS_PATH_RE.test(url);
+}
+export function handleDatabaseWsUpgrade(wss, req, socket, head) {
+    const match = (req.url ?? '').match(WS_PATH_RE);
+    if (!match) {
+        socket.destroy();
+        return;
+    }
+    const connectionId = decodeURIComponent(match[1]);
+    wss.handleUpgrade(req, socket, head, (ws) => {
+        wss.emit('connection', ws, req);
+        setupQuerySocket(ws, connectionId);
+    });
+}
+function send(ws, data) {
+    if (ws.readyState === ws.OPEN) {
+        ws.send(JSON.stringify(data));
+    }
+}
+function setupQuerySocket(ws, connectionId) {
+    const keepAlive = setInterval(() => {
+        if (ws.readyState === ws.OPEN) {
+            ws.ping();
+        }
+    }, KEEPALIVE_INTERVAL_MS);
+    ws.on('close', () => {
+        clearInterval(keepAlive);
+    });
+    ws.on('message', async (raw) => {
+        let msg;
+        try {
+            msg = JSON.parse(typeof raw === 'string' ? raw : raw.toString('utf-8'));
+        }
+        catch {
+            send(ws, { type: 'error', message: 'Invalid JSON' });
+            return;
+        }
+        if (msg.type !== 'query' || typeof msg.query !== 'string') {
+            send(ws, { type: 'error', message: 'Expected message of type "query" with a "query" string field' });
+            return;
+        }
+        const limit = typeof msg.limit === 'number' && msg.limit > 0 ? msg.limit : undefined;
+        try {
+            const startTime = performance.now();
+            const result = await connectionManager.query(connectionId, msg.query, limit);
+            const executionTimeMs = Math.round(performance.now() - startTime);
+            // Send column names first
+            send(ws, { type: 'columns', columns: result.columns });
+            // Send rows in batches
+            const rows = result.rows;
+            let batchNumber = 0;
+            for (let offset = 0; offset < rows.length; offset += BATCH_SIZE) {
+                const batch = rows.slice(offset, offset + BATCH_SIZE);
+                send(ws, { type: 'rows', rows: batch, batch: batchNumber });
+                batchNumber++;
+            }
+            // Send completion message
+            send(ws, {
+                type: 'complete',
+                row_count: result.row_count,
+                execution_time_ms: executionTimeMs,
+                truncated: result.truncated ?? false,
+            });
+        }
+        catch (err) {
+            const message = err instanceof Error ? err.message : String(err);
+            send(ws, { type: 'error', message });
+        }
+    });
+}

package/dist/src/database/types.d.ts

@@ -0,0 +1,71 @@
+export type DbType = 'postgresql' | 'mysql' | 'mariadb' | 'sqlite' | 'redis' | 'mongodb' | 'elasticsearch' | 'opensearch' | 'clickhouse' | 'cockroachdb' | 'cassandra' | 'scylladb' | 'neo4j' | 'influxdb' | 'memcached';
+export type DiscoverySource = 'local' | 'container' | 'manual' | 'ssh';
+export interface ConnectionParams {
+    host?: string;
+    port?: number;
+    username?: string;
+    password?: string;
+    database?: string;
+    filename?: string;
+    connection_string?: string;
+    /** SSH tunnel: Atoo Studio SSH connection ID */
+    ssh_connection_id?: string;
+    /** SSH tunnel: remote host (default 127.0.0.1) */
+    ssh_remote_host?: string;
+    /** SSH tunnel: remote port (defaults to the DB's default port) */
+    ssh_remote_port?: number;
+}
+export interface DiscoveredDatabase {
+    id: string;
+    name: string;
+    db_type: DbType;
+    source: DiscoverySource;
+    params: ConnectionParams;
+    /** e.g. container name, file path */
+    source_detail?: string;
+    connected?: boolean;
+}
+export interface QueryResult {
+    columns: string[];
+    rows: Record<string, any>[];
+    row_count: number;
+    execution_time_ms: number;
+    truncated?: boolean;
+}
+export interface TableInfo {
+    name: string;
+    type?: string;
+    row_count?: number;
+    size_bytes?: number;
+}
+export interface ColumnInfo {
+    name: string;
+    type: string;
+    nullable: boolean;
+    default_value?: string;
+    primary_key?: boolean;
+    comment?: string;
+}
+export interface SchemaInfo {
+    table: string;
+    columns: ColumnInfo[];
+    indexes?: {
+        name: string;
+        columns: string[];
+        unique: boolean;
+    }[];
+    foreign_keys?: {
+        column: string;
+        ref_table: string;
+        ref_column: string;
+    }[];
+}
+export interface DatabaseDriver {
+    connect(params: ConnectionParams): Promise<void>;
+    disconnect(): Promise<void>;
+    isConnected(): boolean;
+    query(sql: string, limit?: number): Promise<QueryResult>;
+    getTables(): Promise<TableInfo[]>;
+    describeTable(table: string): Promise<SchemaInfo>;
+    getDatabases?(): Promise<string[]>;
+}

package/dist/src/database/types.js

@@ -0,0 +1 @@
+export {};

package/dist/src/events/index.d.ts

@@ -0,0 +1,3 @@
+export * from './types.js';
+export { toWireMessages, tryParseContextUsageWire } from './wire.js';
+export type { WireMessage, WireUserMessage, WireAssistantMessage, WireThinking, WireToolUse, WireToolRequest, WireQuestion, WireQuestionItem, WirePlanApproval, WireStatusUpdate, WireContextUsage, WireSystemMessage, WireResult, WireAttachment, } from './wire.js';

package/dist/src/events/index.js

@@ -0,0 +1,3 @@
+// Barrel export for events module
+export * from './types.js';
+export { toWireMessages, tryParseContextUsageWire } from './wire.js';