@vaultgradient/pipequery-cli 0.6.1 → 0.10.0

package/src/server/sources/snowflake.ts

@@ -0,0 +1,158 @@
+import snowflake, { type Connection } from 'snowflake-sdk';
+import type { SnowflakeSourceConfig } from '../../config/schema.js';
+import { parseDuration } from '../../utils/parseDuration.js';
+import { expandEnv } from '../../utils/expandEnv.js';
+import { parseQuery } from '../../engine.js';
+import { compileSnowflakePushdown } from './pushdown/snowflake.js';
+import type { SourceAdapter, SourceStatus } from './types.js';
+
+const DEFAULT_INTERVAL_MS = 60_000;
+const DEFAULT_MAX_ROWS = 10_000;
+
+export class SnowflakeSourceAdapter implements SourceAdapter {
+  private conn: Connection | null = null;
+  private data: unknown[] = [];
+  private timer: ReturnType<typeof setInterval> | null = null;
+  private lastFetch: Date | null = null;
+  private error?: string;
+  private listeners = new Set<(data: unknown[]) => void>();
+  private readonly intervalMs: number;
+  private readonly maxRows: number;
+
+  constructor(private config: SnowflakeSourceConfig) {
+    this.intervalMs = config.interval ? parseDuration(config.interval) : DEFAULT_INTERVAL_MS;
+    this.maxRows = config.maxRows ?? DEFAULT_MAX_ROWS;
+  }
+
+  async start(): Promise<void> {
+    this.conn = snowflake.createConnection({
+      account: expandEnv(this.config.account, 'snowflake source account'),
+      username: expandEnv(this.config.username, 'snowflake source username'),
+      password: expandEnv(this.config.password, 'snowflake source password'),
+      ...(this.config.database ? { database: expandEnv(this.config.database, 'snowflake source database') } : {}),
+      ...(this.config.schema ? { schema: expandEnv(this.config.schema, 'snowflake source schema') } : {}),
+      ...(this.config.warehouse ? { warehouse: expandEnv(this.config.warehouse, 'snowflake source warehouse') } : {}),
+      ...(this.config.role ? { role: expandEnv(this.config.role, 'snowflake source role') } : {}),
+    });
+
+    await new Promise<void>((resolve, reject) => {
+      this.conn!.connect((err) => (err ? reject(err) : resolve()));
+    });
+
+    await this.fetch();
+    this.timer = setInterval(() => {
+      void this.fetch();
+    }, this.intervalMs);
+  }
+
+  stop(): void {
+    if (this.timer) {
+      clearInterval(this.timer);
+      this.timer = null;
+    }
+    if (this.conn) {
+      this.conn.destroy(() => undefined);
+      this.conn = null;
+    }
+  }
+
+  getData(): unknown[] {
+    return this.data;
+  }
+
+  getStatus(): SourceStatus {
+    return {
+      healthy: !this.error,
+      rowCount: this.data.length,
+      lastFetch: this.lastFetch,
+      error: this.error,
+    };
+  }
+
+  onUpdate(callback: (data: unknown[]) => void): () => void {
+    this.listeners.add(callback);
+    return () => this.listeners.delete(callback);
+  }
+
+  async refresh(): Promise<void> {
+    await this.fetch();
+  }
+
+  /**
+   * Adapter-native push-down. SourceManager.runQuery dispatches here when a
+   * pipeline targets a Snowflake source; declines fall back to in-process.
+   */
+  async runPushdown(expression: string): Promise<
+    | { ok: true; rows: unknown[]; sql: string; params: unknown[] }
+    | { ok: false; reason: string }
+  > {
+    if (!this.conn) {
+      return { ok: false, reason: 'Snowflake connection not initialized — call start() first' };
+    }
+    let pipeline;
+    try {
+      pipeline = parseQuery(expression);
+    } catch (err) {
+      return { ok: false, reason: `parse error: ${err instanceof Error ? err.message : String(err)}` };
+    }
+    const compiled = compileSnowflakePushdown(pipeline, this.config.query);
+    if (!compiled.ok) {
+      return { ok: false, reason: compiled.reason };
+    }
+    try {
+      const rows = await this.executeSql(compiled.compiled.sql, compiled.compiled.params);
+      return {
+        ok: true,
+        rows,
+        sql: compiled.compiled.sql,
+        params: compiled.compiled.params,
+      };
+    } catch (err) {
+      return { ok: false, reason: `snowflake error: ${err instanceof Error ? err.message : String(err)}` };
+    }
+  }
+
+  private async fetch(): Promise<void> {
+    if (!this.conn) {
+      this.error = 'Snowflake connection not initialized';
+      return;
+    }
+    try {
+      const capQuery = `SELECT * FROM (${stripTrailingSemicolon(this.config.query)}) AS pq_src LIMIT ${this.maxRows + 1}`;
+      const rows = await this.executeSql(capQuery, []);
+      if (rows.length > this.maxRows) {
+        throw new Error(
+          `Query returned more than maxRows=${this.maxRows}. ` +
+            `Add LIMIT/WHERE to bound the result set, or raise maxRows in the source config.`,
+        );
+      }
+      this.data = rows;
+      this.lastFetch = new Date();
+      this.error = undefined;
+      this.notify();
+    } catch (err) {
+      this.error = err instanceof Error ? err.message : String(err);
+    }
+  }
+
+  private executeSql(sqlText: string, params: unknown[]): Promise<unknown[]> {
+    return new Promise((resolve, reject) => {
+      this.conn!.execute({
+        sqlText,
+        binds: params as snowflake.Binds,
+        complete: (err, _stmt, rows) => {
+          if (err) reject(err);
+          else resolve(rows ?? []);
+        },
+      });
+    });
+  }
+
+  private notify(): void {
+    for (const cb of this.listeners) cb(this.data);
+  }
+}
+
+function stripTrailingSemicolon(sql: string): string {
+  return sql.trim().replace(/;+\s*$/, '');
+}

package/src/server/sources/types.ts

@@ -5,6 +5,26 @@ export interface SourceStatus {
   error?: string;
 }
 
+/**
+ * Result of an adapter-native push-down attempt.
+ *
+ * `ok: true` means the adapter executed the pipe expression natively (e.g.
+ * Postgres compiled it to SQL and ran it on the database) and returned rows
+ * without materializing the source query in memory. `sql` / `params` are
+ * optional telemetry — useful for audit logs and debugging, not required
+ * for callers to act on.
+ *
+ * `ok: false` is the explicit decline channel: the adapter recognised the
+ * shape but declined (unsupported op, multi-source pipeline, etc.). The
+ * caller should fall back to in-process execution against `getContext()`.
+ *
+ * Any thrown error means *neither* path succeeded — let it propagate; the
+ * caller has no useful fallback.
+ */
+export type PushdownResult =
+  | { ok: true; rows: unknown[]; sql?: string; params?: unknown[] }
+  | { ok: false; reason: string };
+
 export interface SourceAdapter {
   start(): Promise<void>;
   stop(): void;
@@ -12,4 +32,11 @@ export interface SourceAdapter {
   getStatus(): SourceStatus;
   onUpdate(callback: (data: unknown[]) => void): () => void;
   refresh?(): Promise<void>;
+  /**
+   * Optional: compile + execute the pipe expression natively against the
+   * underlying engine instead of materializing into memory. When absent
+   * (or when the adapter declines), SourceManager falls back to in-process
+   * execution. Today only PostgresSourceAdapter implements this.
+   */
+  runPushdown?(expression: string): Promise<PushdownResult>;
 }

package/src/server/sources/websocket.ts

@@ -1,6 +1,7 @@
 import WebSocket from 'ws';
 import type { WebSocketSourceConfig } from '../../config/schema.js';
 import type { SourceAdapter, SourceStatus } from './types.js';
+import { parseDuration } from '../../utils/parseDuration.js';
 
 export class WebSocketSourceAdapter implements SourceAdapter {
   private data: unknown[] = [];
@@ -9,21 +10,44 @@ export class WebSocketSourceAdapter implements SourceAdapter {
   private error?: string;
   private listeners = new Set<(data: unknown[]) => void>();
   private maxBuffer: number;
+  private heartbeatTimer: ReturnType<typeof setInterval> | null = null;
+  /** Tracks whether stop() was called so reconnect respects intent. */
+  private stopped = false;
 
   constructor(private config: WebSocketSourceConfig) {
     this.maxBuffer = config.maxBuffer ?? 1000;
   }
 
   async start(): Promise<void> {
+    this.stopped = false;
     return new Promise((resolve, reject) => {
-      this.ws = new WebSocket(this.config.url);
+      const ws = new WebSocket(this.config.url);
+      this.ws = ws;
+      let settled = false;
 
-      this.ws.on('open', () => {
+      ws.on('open', () => {
         this.error = undefined;
-        resolve();
+        try {
+          this.sendSubscribePayloads(ws);
+          this.startHeartbeat(ws);
+        } catch (err) {
+          // A failure to send the subscribe frame leaves the socket open
+          // but useless; record it and reject the start() promise.
+          this.error = err instanceof Error ? err.message : String(err);
+          ws.close();
+          if (!settled) {
+            settled = true;
+            reject(err instanceof Error ? err : new Error(this.error));
+          }
+          return;
+        }
+        if (!settled) {
+          settled = true;
+          resolve();
+        }
       });
 
-      this.ws.on('message', (raw) => {
+      ws.on('message', (raw) => {
         try {
           const parsed = JSON.parse(raw.toString());
           if (Array.isArray(parsed)) {
@@ -42,21 +66,30 @@ export class WebSocketSourceAdapter implements SourceAdapter {
         }
       });
 
-      this.ws.on('error', (err) => {
+      ws.on('error', (err) => {
         this.error = err.message;
-        reject(err);
+        if (!settled) {
+          settled = true;
+          reject(err);
+        }
       });
 
-      this.ws.on('close', () => {
-        // Auto-reconnect after 5s
+      ws.on('close', () => {
+        this.stopHeartbeat();
+        // Auto-reconnect after 5s unless stop() was explicitly called.
+        // Re-running start() also re-sends the subscribe payloads, which
+        // is exactly what every exchange expects after a connection drop.
+        if (this.stopped) return;
         setTimeout(() => {
-          if (this.ws) this.start().catch(() => {});
+          if (!this.stopped) this.start().catch(() => {});
         }, 5000);
       });
     });
   }
 
   stop(): void {
+    this.stopped = true;
+    this.stopHeartbeat();
     const ws = this.ws;
     this.ws = null;
     ws?.close();
@@ -83,4 +116,35 @@ export class WebSocketSourceAdapter implements SourceAdapter {
   private notify(): void {
     for (const cb of this.listeners) cb(this.data);
   }
+
+  private sendSubscribePayloads(ws: WebSocket): void {
+    const subscribe = this.config.subscribe;
+    if (subscribe === undefined || subscribe === null) return;
+    const list = Array.isArray(subscribe) ? subscribe : [subscribe];
+    for (const payload of list) {
+      ws.send(JSON.stringify(payload));
+    }
+  }
+
+  private startHeartbeat(ws: WebSocket): void {
+    this.stopHeartbeat();
+    const hb = this.config.heartbeat;
+    if (!hb) return;
+    const intervalMs = parseDuration(hb.interval);
+    this.heartbeatTimer = setInterval(() => {
+      if (ws.readyState !== WebSocket.OPEN) return;
+      try {
+        ws.send(JSON.stringify(hb.payload));
+      } catch {
+        // Ignore — next 'close' / 'error' will surface anything fatal.
+      }
+    }, intervalMs);
+  }
+
+  private stopHeartbeat(): void {
+    if (this.heartbeatTimer) {
+      clearInterval(this.heartbeatTimer);
+      this.heartbeatTimer = null;
+    }
+  }
 }

package/src/telegram/bot.ts

@@ -10,6 +10,7 @@ import { Bot, type Context, GrammyError, HttpError } from 'grammy';
 import type { Provider } from '../mcp/provider.js';
 import { formatError, formatResult, escapeHtml } from './format.js';
 import type { NLTranslator } from './nl.js';
+import type { EventLogger, UserRef } from './event-log.js';
 
 export interface BotOptions {
   /** Whitelist of Telegram user IDs (numeric) or @usernames allowed to query.
@@ -18,12 +19,16 @@ export interface BotOptions {
   /** Optional natural-language translator. If present, plain-text messages
    *  (no leading slash) are routed through it. */
   nl?: NLTranslator;
+  /** Optional structured event logger. When present, every command,
+   *  NL query, and unauthorized attempt is recorded. */
+  logger?: EventLogger;
 }
 
 export function buildBot(token: string, provider: Provider, opts: BotOptions = {}): Bot {
   const bot = new Bot(token);
   const allowList = normalizeAllowList(opts.allowUsers ?? []);
   const nl = opts.nl;
+  const logger = opts.logger;
 
   // Auth middleware — runs before any handler.
   bot.use(async (ctx, next) => {
@@ -36,6 +41,7 @@ export function buildBot(token: string, provider: Provider, opts: BotOptions = {
       return;
     }
     if (!isAllowed(ctx, allowList)) {
+      logger?.log({ kind: 'unauthorized', user: userOf(ctx) });
       await ctx.reply('🔒 You are not authorized to use this bot.');
       return;
     }
@@ -43,6 +49,7 @@ export function buildBot(token: string, provider: Provider, opts: BotOptions = {
   });
 
   bot.command('start', async (ctx) => {
+    const start = Date.now();
     await ctx.reply(
       [
         '👋 <b>pipequery bot</b>',
@@ -59,9 +66,11 @@ export function buildBot(token: string, provider: Provider, opts: BotOptions = {
       ].join('\n'),
       { parse_mode: 'HTML' },
     );
+    logger?.log({ kind: 'command', user: userOf(ctx), cmd: '/start', outcome: 'success', latencyMs: Date.now() - start });
   });
 
   bot.command('help', async (ctx) => {
+    const start = Date.now();
     const lines = [
       '<b>Examples</b>',
       '<code>/sources</code>',
@@ -86,37 +95,73 @@ export function buildBot(token: string, provider: Provider, opts: BotOptions = {
       '• Results &gt; 30 rows are truncated with a footer; tighten your query to see all of them.',
     );
     await ctx.reply(lines.join('\n'), { parse_mode: 'HTML' });
+    logger?.log({ kind: 'command', user: userOf(ctx), cmd: '/help', outcome: 'success', latencyMs: Date.now() - start });
   });
 
   bot.command('sources', async (ctx) => {
+    const start = Date.now();
     try {
       const sources = await provider.listSources();
       if (sources.length === 0) {
         await ctx.reply('<i>No sources configured.</i>', { parse_mode: 'HTML' });
-        return;
+      } else {
+        const lines = sources.map((s) => {
+          const dot = s.status.healthy ? '🟢' : '🔴';
+          const rowCount = s.status.rowCount.toLocaleString();
+          const err = s.status.error ? ` <i>(${escapeHtml(s.status.error)})</i>` : '';
+          return `${dot} <b>${escapeHtml(s.name)}</b> — ${rowCount} rows${err}`;
+        });
+        await ctx.reply(lines.join('\n'), { parse_mode: 'HTML' });
       }
-      const lines = sources.map((s) => {
-        const dot = s.status.healthy ? '🟢' : '🔴';
-        const rowCount = s.status.rowCount.toLocaleString();
-        const err = s.status.error ? ` <i>(${escapeHtml(s.status.error)})</i>` : '';
-        return `${dot} <b>${escapeHtml(s.name)}</b> — ${rowCount} rows${err}`;
+      logger?.log({
+        kind: 'command',
+        user: userOf(ctx),
+        cmd: '/sources',
+        outcome: 'success',
+        rowCount: sources.length,
+        latencyMs: Date.now() - start,
       });
-      await ctx.reply(lines.join('\n'), { parse_mode: 'HTML' });
     } catch (err) {
       await ctx.reply(formatError(err), { parse_mode: 'HTML' });
+      logger?.log({
+        kind: 'command',
+        user: userOf(ctx),
+        cmd: '/sources',
+        outcome: 'error',
+        latencyMs: Date.now() - start,
+        error: errMessage(err),
+      });
     }
   });
 
   bot.command('describe', async (ctx) => {
+    const start = Date.now();
     const name = ctx.match.trim();
     if (!name) {
       await ctx.reply('Usage: <code>/describe &lt;source-name&gt;</code>', { parse_mode: 'HTML' });
+      logger?.log({
+        kind: 'command',
+        user: userOf(ctx),
+        cmd: '/describe',
+        outcome: 'error',
+        latencyMs: Date.now() - start,
+        error: 'missing source name',
+      });
       return;
     }
     try {
       const desc = await provider.describeSource(name, 5);
       if (!desc) {
         await ctx.reply(`Source <code>${escapeHtml(name)}</code> not found.`, { parse_mode: 'HTML' });
+        logger?.log({
+          kind: 'command',
+          user: userOf(ctx),
+          cmd: '/describe',
+          args: name,
+          outcome: 'error',
+          latencyMs: Date.now() - start,
+          error: 'source not found',
+        });
         return;
       }
       const fields = desc.fields.length ? desc.fields.map((f) => `<code>${escapeHtml(f)}</code>`).join(', ') : '<i>(none inferred)</i>';
@@ -125,52 +170,141 @@ export function buildBot(token: string, provider: Provider, opts: BotOptions = {
         `<b>${escapeHtml(name)}</b> — ${desc.status.rowCount.toLocaleString()} rows\n<b>Fields:</b> ${fields}\n\n<b>Sample:</b>\n${sample}`,
         { parse_mode: 'HTML' },
       );
+      logger?.log({
+        kind: 'command',
+        user: userOf(ctx),
+        cmd: '/describe',
+        args: name,
+        outcome: 'success',
+        rowCount: desc.sample.length,
+        latencyMs: Date.now() - start,
+      });
     } catch (err) {
       await ctx.reply(formatError(err), { parse_mode: 'HTML' });
+      logger?.log({
+        kind: 'command',
+        user: userOf(ctx),
+        cmd: '/describe',
+        args: name,
+        outcome: 'error',
+        latencyMs: Date.now() - start,
+        error: errMessage(err),
+      });
     }
   });
 
   bot.command('endpoints', async (ctx) => {
+    const start = Date.now();
     try {
       const endpoints = await provider.listEndpoints();
       if (endpoints.length === 0) {
         await ctx.reply('<i>No endpoints configured.</i>', { parse_mode: 'HTML' });
-        return;
+      } else {
+        const lines = endpoints.map(
+          (e) => `<b>${escapeHtml(e.path)}</b>\n<code>${escapeHtml(e.config.query)}</code>`,
+        );
+        await ctx.reply(lines.join('\n\n'), { parse_mode: 'HTML' });
       }
-      const lines = endpoints.map(
-        (e) => `<b>${escapeHtml(e.path)}</b>\n<code>${escapeHtml(e.config.query)}</code>`,
-      );
-      await ctx.reply(lines.join('\n\n'), { parse_mode: 'HTML' });
+      logger?.log({
+        kind: 'command',
+        user: userOf(ctx),
+        cmd: '/endpoints',
+        outcome: 'success',
+        rowCount: endpoints.length,
+        latencyMs: Date.now() - start,
+      });
     } catch (err) {
       await ctx.reply(formatError(err), { parse_mode: 'HTML' });
+      logger?.log({
+        kind: 'command',
+        user: userOf(ctx),
+        cmd: '/endpoints',
+        outcome: 'error',
+        latencyMs: Date.now() - start,
+        error: errMessage(err),
+      });
     }
   });
 
   bot.command('call', async (ctx) => {
+    const start = Date.now();
     const path = ctx.match.trim();
     if (!path) {
       await ctx.reply('Usage: <code>/call &lt;endpoint-path&gt;</code>', { parse_mode: 'HTML' });
+      logger?.log({
+        kind: 'command',
+        user: userOf(ctx),
+        cmd: '/call',
+        outcome: 'error',
+        latencyMs: Date.now() - start,
+        error: 'missing path',
+      });
       return;
     }
     try {
       const result = await provider.callEndpoint(path);
       await ctx.reply(formatResult(result), { parse_mode: 'HTML' });
+      logger?.log({
+        kind: 'command',
+        user: userOf(ctx),
+        cmd: '/call',
+        args: path,
+        outcome: 'success',
+        rowCount: rowCountOf(result),
+        latencyMs: Date.now() - start,
+      });
     } catch (err) {
       await ctx.reply(formatError(err), { parse_mode: 'HTML' });
+      logger?.log({
+        kind: 'command',
+        user: userOf(ctx),
+        cmd: '/call',
+        args: path,
+        outcome: 'error',
+        latencyMs: Date.now() - start,
+        error: errMessage(err),
+      });
     }
   });
 
   bot.command('query', async (ctx) => {
+    const start = Date.now();
     const expr = ctx.match.trim();
     if (!expr) {
       await ctx.reply('Usage: <code>/query &lt;pipequery expression&gt;</code>', { parse_mode: 'HTML' });
+      logger?.log({
+        kind: 'command',
+        user: userOf(ctx),
+        cmd: '/query',
+        outcome: 'error',
+        latencyMs: Date.now() - start,
+        error: 'missing expression',
+      });
       return;
     }
     try {
       const result = await provider.runQuery(expr);
       await ctx.reply(formatResult(result), { parse_mode: 'HTML' });
+      logger?.log({
+        kind: 'command',
+        user: userOf(ctx),
+        cmd: '/query',
+        args: expr,
+        outcome: 'success',
+        rowCount: rowCountOf(result),
+        latencyMs: Date.now() - start,
+      });
     } catch (err) {
       await ctx.reply(formatError(err), { parse_mode: 'HTML' });
+      logger?.log({
+        kind: 'command',
+        user: userOf(ctx),
+        cmd: '/query',
+        args: expr,
+        outcome: 'error',
+        latencyMs: Date.now() - start,
+        error: errMessage(err),
+      });
     }
   });
 
@@ -180,12 +314,21 @@ export function buildBot(token: string, provider: Provider, opts: BotOptions = {
   bot.on('message:text', async (ctx) => {
     const text = ctx.message.text.trim();
     if (!text || text.startsWith('/')) return;
+    const start = Date.now();
     if (!nl) {
       await ctx.reply(
         'I don\'t recognise that. Use <code>/help</code> to see the available commands. ' +
           'Natural-language queries are disabled — start the bot with <code>--anthropic-key</code> to enable them.',
         { parse_mode: 'HTML' },
       );
+      logger?.log({
+        kind: 'nl',
+        user: userOf(ctx),
+        text,
+        outcome: 'error',
+        latencyMs: Date.now() - start,
+        error: 'NL translation disabled',
+      });
       return;
     }
     try {
@@ -194,13 +337,39 @@ export function buildBot(token: string, provider: Provider, opts: BotOptions = {
       if (!expression) {
         const reason = explanation || 'I couldn\'t translate that into a pipequery expression.';
         await ctx.reply(`🤔 ${escapeHtml(reason)}`, { parse_mode: 'HTML' });
+        logger?.log({
+          kind: 'nl',
+          user: userOf(ctx),
+          text,
+          explanation,
+          outcome: 'not_answerable',
+          latencyMs: Date.now() - start,
+        });
         return;
       }
       const result = await provider.runQuery(expression);
       const header = `<i>${escapeHtml(explanation || 'Translated query')}</i>\n<code>${escapeHtml(expression)}</code>\n\n`;
       await ctx.reply(header + formatResult(result), { parse_mode: 'HTML' });
+      logger?.log({
+        kind: 'nl',
+        user: userOf(ctx),
+        text,
+        expression,
+        explanation,
+        outcome: 'success',
+        rowCount: rowCountOf(result),
+        latencyMs: Date.now() - start,
+      });
     } catch (err) {
       await ctx.reply(formatError(err), { parse_mode: 'HTML' });
+      logger?.log({
+        kind: 'nl',
+        user: userOf(ctx),
+        text,
+        outcome: 'error',
+        latencyMs: Date.now() - start,
+        error: errMessage(err),
+      });
     }
   });
 
@@ -245,3 +414,20 @@ function trackUnauthSeen(id: number, username?: string): void {
       `Restart with --allow-user @yourname to lock it down.\n`,
   );
 }
+
+// ─── Logging helpers ────────────────────────────────────────────────────────
+
+function userOf(ctx: Context): UserRef {
+  return {
+    id: ctx.from?.id,
+    username: ctx.from?.username,
+  };
+}
+
+function rowCountOf(value: unknown): number | undefined {
+  return Array.isArray(value) ? value.length : undefined;
+}
+
+function errMessage(err: unknown): string {
+  return err instanceof Error ? err.message : String(err);
+}