@vaultgradient/pipequery-cli 0.3.2 → 0.4.0

package/src/mcp/provider.ts

@@ -0,0 +1,223 @@
+import type { DataContext } from '../engine.js';
+import type { EndpointConfig, PipeQueryConfig } from '../config/schema.js';
+import type { SourceStatus } from '../server/sources/types.js';
+import { SourceManager } from '../server/sources/manager.js';
+import { query } from '../engine.js';
+
+export interface ProviderSourceInfo {
+  name: string;
+  status: SourceStatus;
+}
+
+export interface ProviderSourceDescription {
+  name: string;
+  status: SourceStatus;
+  sample: unknown[];
+  fields: string[];
+}
+
+export interface Provider {
+  listSources(): Promise<ProviderSourceInfo[]>;
+  describeSource(name: string, sampleSize?: number): Promise<ProviderSourceDescription | undefined>;
+  listEndpoints(): Promise<Array<{ path: string; config: EndpointConfig }>>;
+  callEndpoint(path: string): Promise<unknown>;
+  runQuery(expression: string): Promise<unknown>;
+  dispose(): Promise<void>;
+}
+
+// ─── LocalProvider ──────────────────────────────────────────────────────────
+// Boots its own SourceManager from a loaded config. Used when `pq mcp serve`
+// runs standalone (no attached `pq serve` daemon).
+
+export class LocalProvider implements Provider {
+  private sourceManager: SourceManager;
+  private endpoints: Map<string, EndpointConfig>;
+  private ready: Promise<void>;
+
+  constructor(private config: PipeQueryConfig, cwd: string) {
+    this.sourceManager = new SourceManager(cwd);
+    this.endpoints = new Map(Object.entries(config.endpoints));
+    this.ready = this.bootSources();
+  }
+
+  private async bootSources(): Promise<void> {
+    for (const [name, src] of Object.entries(this.config.sources)) {
+      try {
+        await this.sourceManager.addSource(name, src);
+      } catch (err) {
+        const msg = err instanceof Error ? err.message : String(err);
+        // Print to stderr; stdout is reserved for the MCP JSON-RPC stream in stdio mode.
+        process.stderr.write(`[pipequery-mcp] Source "${name}" failed to start: ${msg}\n`);
+      }
+    }
+  }
+
+  async listSources(): Promise<ProviderSourceInfo[]> {
+    await this.ready;
+    const statuses = this.sourceManager.getAllStatuses();
+    return Object.entries(statuses).map(([name, status]) => ({ name, status }));
+  }
+
+  async describeSource(name: string, sampleSize = 5): Promise<ProviderSourceDescription | undefined> {
+    await this.ready;
+    const status = this.sourceManager.getSourceStatus(name);
+    const data = this.sourceManager.getSourceData(name);
+    if (!status || !data) return undefined;
+
+    const sample = data.slice(0, sampleSize);
+    const fields = inferFields(data);
+    return { name, status, sample, fields };
+  }
+
+  async listEndpoints(): Promise<Array<{ path: string; config: EndpointConfig }>> {
+    await this.ready;
+    return [...this.endpoints].map(([path, config]) => ({ path, config }));
+  }
+
+  async callEndpoint(path: string): Promise<unknown> {
+    await this.ready;
+    const endpoint = this.endpoints.get(path);
+    if (!endpoint) {
+      throw new Error(`Endpoint "${path}" not found`);
+    }
+    const context = this.sourceManager.getContext();
+    return query(context, endpoint.query);
+  }
+
+  async runQuery(expression: string): Promise<unknown> {
+    await this.ready;
+    const context = this.sourceManager.getContext();
+    return query(context, expression);
+  }
+
+  async dispose(): Promise<void> {
+    await this.sourceManager.dispose();
+  }
+
+  // Exposed for the HTTP plugin that wants to mount the MCP server into an
+  // existing `createServer` instance; the caller passes in their own manager.
+  static withSourceManager(sourceManager: SourceManager, endpoints: Map<string, EndpointConfig>): Provider {
+    return new AttachedLocalProvider(sourceManager, endpoints);
+  }
+}
+
+class AttachedLocalProvider implements Provider {
+  constructor(private sourceManager: SourceManager, private endpoints: Map<string, EndpointConfig>) {}
+
+  async listSources(): Promise<ProviderSourceInfo[]> {
+    const statuses = this.sourceManager.getAllStatuses();
+    return Object.entries(statuses).map(([name, status]) => ({ name, status }));
+  }
+
+  async describeSource(name: string, sampleSize = 5): Promise<ProviderSourceDescription | undefined> {
+    const status = this.sourceManager.getSourceStatus(name);
+    const data = this.sourceManager.getSourceData(name);
+    if (!status || !data) return undefined;
+    return { name, status, sample: data.slice(0, sampleSize), fields: inferFields(data) };
+  }
+
+  async listEndpoints() {
+    return [...this.endpoints].map(([path, config]) => ({ path, config }));
+  }
+
+  async callEndpoint(path: string): Promise<unknown> {
+    const endpoint = this.endpoints.get(path);
+    if (!endpoint) throw new Error(`Endpoint "${path}" not found`);
+    return query(this.sourceManager.getContext(), endpoint.query);
+  }
+
+  async runQuery(expression: string): Promise<unknown> {
+    return query(this.sourceManager.getContext(), expression);
+  }
+
+  async dispose(): Promise<void> {
+    // Intentionally noop: the owning server disposes the SourceManager.
+  }
+}
+
+// ─── AttachedProvider ───────────────────────────────────────────────────────
+// Talks to a running `pq serve` via its existing /api/_control/* endpoints.
+
+export class AttachedProvider implements Provider {
+  constructor(private baseUrl: string) {
+    this.baseUrl = baseUrl.replace(/\/$/, '');
+  }
+
+  private async get<T>(path: string): Promise<T> {
+    const res = await fetch(`${this.baseUrl}${path}`);
+    if (!res.ok) throw new Error(`GET ${path} failed: HTTP ${res.status}`);
+    return res.json() as Promise<T>;
+  }
+
+  private async post<T>(path: string, body: unknown): Promise<T> {
+    const res = await fetch(`${this.baseUrl}${path}`, {
+      method: 'POST',
+      headers: { 'content-type': 'application/json' },
+      body: JSON.stringify(body),
+    });
+    if (!res.ok) throw new Error(`POST ${path} failed: HTTP ${res.status}`);
+    return res.json() as Promise<T>;
+  }
+
+  async listSources(): Promise<ProviderSourceInfo[]> {
+    const statuses = await this.get<Record<string, SourceStatus>>('/api/_control/sources');
+    return Object.entries(statuses).map(([name, status]) => ({ name, status }));
+  }
+
+  async describeSource(name: string, sampleSize = 5): Promise<ProviderSourceDescription | undefined> {
+    const res = await this.post<{
+      ok: boolean;
+      error?: string;
+      rowCount?: number;
+      sample?: unknown[];
+      status?: SourceStatus;
+    }>(`/api/_control/sources/${encodeURIComponent(name)}/test`, {});
+    if (!res.ok || !res.status || !res.sample) return undefined;
+    return {
+      name,
+      status: res.status,
+      sample: res.sample.slice(0, sampleSize),
+      fields: inferFields(res.sample),
+    };
+  }
+
+  async listEndpoints() {
+    const endpoints = await this.get<Record<string, EndpointConfig>>('/api/_control/endpoints');
+    return Object.entries(endpoints).map(([path, config]) => ({ path, config }));
+  }
+
+  async callEndpoint(path: string): Promise<unknown> {
+    // Just hit the actual public endpoint on the running server.
+    const normalized = path.startsWith('/') ? path : `/${path}`;
+    const res = await fetch(`${this.baseUrl}${normalized}`);
+    if (!res.ok) throw new Error(`GET ${normalized} failed: HTTP ${res.status}`);
+    return res.json();
+  }
+
+  async runQuery(expression: string): Promise<unknown> {
+    const res = await this.post<{ ok: boolean; result?: unknown; error?: string }>(
+      '/api/_control/query',
+      { query: expression },
+    );
+    if (!res.ok) throw new Error(res.error ?? 'Query failed');
+    return res.result;
+  }
+
+  async dispose(): Promise<void> {
+    // Stateless — nothing to clean up.
+  }
+}
+
+// ─── Helpers ────────────────────────────────────────────────────────────────
+
+function inferFields(rows: unknown[]): string[] {
+  const fields = new Set<string>();
+  for (const row of rows.slice(0, 20)) {
+    if (row && typeof row === 'object' && !Array.isArray(row)) {
+      for (const k of Object.keys(row)) fields.add(k);
+    }
+  }
+  return [...fields];
+}
+
+export type { DataContext };

package/src/mcp/server.ts

@@ -0,0 +1,138 @@
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import { z } from 'zod';
+import type { Provider } from './provider.js';
+
+export interface McpServerInfo {
+  name: string;
+  version: string;
+}
+
+const DEFAULT_INFO: McpServerInfo = {
+  name: 'pipequery',
+  version: '0.1.0',
+};
+
+/**
+ * Build an MCP server wired to the supplied provider.
+ *
+ * Tools exposed:
+ *   - query              run a PipeQuery expression against live sources
+ *   - list_sources       list configured sources and their health
+ *   - describe_source    return a sample + inferred field names for a source
+ *   - list_endpoints     list pre-configured endpoints and their stored queries
+ *   - call_endpoint      execute a pre-configured endpoint by path
+ */
+export function createMcpServer(provider: Provider, info: McpServerInfo = DEFAULT_INFO): McpServer {
+  const mcp = new McpServer(info);
+
+  mcp.registerTool(
+    'query',
+    {
+      title: 'Run a PipeQuery expression',
+      description:
+        'Execute a PipeQuery pipe-based expression (e.g. `crypto | where(price > 100) | sort(price desc) | first(5)`) against all currently configured sources. Returns the computed result as JSON.',
+      inputSchema: {
+        expression: z
+          .string()
+          .min(1)
+          .describe('The PipeQuery expression. Source names used in the expression must be configured; call `list_sources` to see them.'),
+      },
+    },
+    async ({ expression }) => {
+      const result = await provider.runQuery(expression);
+      return toolResult(result);
+    },
+  );
+
+  mcp.registerTool(
+    'list_sources',
+    {
+      title: 'List configured data sources',
+      description:
+        'List every data source configured in pipequery.yaml along with its current health, row count, and last-fetch timestamp.',
+    },
+    async () => {
+      const sources = await provider.listSources();
+      return toolResult({ count: sources.length, sources });
+    },
+  );
+
+  mcp.registerTool(
+    'describe_source',
+    {
+      title: 'Describe a data source',
+      description:
+        'Return a sample of rows and the inferred field names for a single source. Useful before constructing a query.',
+      inputSchema: {
+        name: z.string().min(1).describe('Source name as returned by list_sources.'),
+        sample_size: z
+          .number()
+          .int()
+          .positive()
+          .max(100)
+          .optional()
+          .describe('Number of rows to include in the sample. Defaults to 5, max 100.'),
+      },
+    },
+    async ({ name, sample_size }) => {
+      const description = await provider.describeSource(name, sample_size);
+      if (!description) {
+        return toolError(`Source "${name}" not found`);
+      }
+      return toolResult(description);
+    },
+  );
+
+  mcp.registerTool(
+    'list_endpoints',
+    {
+      title: 'List pre-configured query endpoints',
+      description:
+        'Return every endpoint declared in pipequery.yaml along with its stored query expression. These endpoints can be invoked via `call_endpoint`.',
+    },
+    async () => {
+      const endpoints = await provider.listEndpoints();
+      return toolResult({ count: endpoints.length, endpoints });
+    },
+  );
+
+  mcp.registerTool(
+    'call_endpoint',
+    {
+      title: 'Invoke a pre-configured endpoint',
+      description:
+        'Execute the stored query for an endpoint path (e.g. `/api/top-coins`) and return its JSON result. Equivalent to the HTTP endpoint on `pq serve`, without needing a live HTTP call.',
+      inputSchema: {
+        path: z.string().min(1).describe('Endpoint path, exactly as declared in pipequery.yaml (typically begins with `/`).'),
+      },
+    },
+    async ({ path }) => {
+      try {
+        const result = await provider.callEndpoint(path);
+        return toolResult(result);
+      } catch (err) {
+        return toolError(err instanceof Error ? err.message : String(err));
+      }
+    },
+  );
+
+  return mcp;
+}
+
+// ─── Helpers ────────────────────────────────────────────────────────────────
+
+function toolResult(data: unknown): { content: Array<{ type: 'text'; text: string }> } {
+  return {
+    content: [{ type: 'text', text: JSON.stringify(data, null, 2) }],
+  };
+}
+
+function toolError(message: string): {
+  content: Array<{ type: 'text'; text: string }>;
+  isError: true;
+} {
+  return {
+    content: [{ type: 'text', text: message }],
+    isError: true,
+  };
+}

package/src/mcp/stdio.ts

@@ -0,0 +1,26 @@
+import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
+import type { Provider } from './provider.js';
+import { createMcpServer, type McpServerInfo } from './server.js';
+
+export async function runStdio(provider: Provider, info?: McpServerInfo): Promise<void> {
+  const mcp = createMcpServer(provider, info);
+  const transport = new StdioServerTransport();
+  await mcp.connect(transport);
+
+  // Clean up on stdio close — Claude Desktop etc. close the child process
+  // by closing its stdin.
+  const cleanup = async () => {
+    try {
+      await mcp.close();
+    } finally {
+      await provider.dispose();
+    }
+  };
+
+  process.stdin.on('close', () => {
+    void cleanup().then(() => process.exit(0));
+  });
+
+  process.on('SIGINT', () => void cleanup().then(() => process.exit(0)));
+  process.on('SIGTERM', () => void cleanup().then(() => process.exit(0)));
+}

package/src/server/sources/kafka.ts

@@ -0,0 +1,166 @@
+import { randomBytes } from 'node:crypto';
+import { Kafka, logLevel, type Consumer, type SASLOptions } from 'kafkajs';
+import type { KafkaSourceConfig } from '../../config/schema.js';
+import type { SourceAdapter, SourceStatus } from './types.js';
+
+const DEFAULT_MAX_BUFFER = 1000;
+
+function expandEnv(input: string): string {
+  return input.replace(/\$\{([A-Z0-9_]+)\}/gi, (_, name: string) => {
+    const value = process.env[name];
+    if (value === undefined) {
+      process.stderr.write(
+        `[pipequery] env var "${name}" referenced in kafka source is not set\n`,
+      );
+      return '';
+    }
+    return value;
+  });
+}
+
+function normalizeBrokers(brokers: string | string[]): string[] {
+  const list = Array.isArray(brokers) ? brokers : brokers.split(',');
+  return list.map((b) => expandEnv(b.trim())).filter(Boolean);
+}
+
+export class KafkaSourceAdapter implements SourceAdapter {
+  private data: unknown[] = [];
+  private kafka: Kafka | null = null;
+  private consumer: Consumer | null = null;
+  private lastFetch: Date | null = null;
+  private error?: string;
+  private connected = false;
+  private listeners = new Set<(data: unknown[]) => void>();
+  private readonly maxBuffer: number;
+  private readonly groupId: string;
+
+  constructor(private config: KafkaSourceConfig) {
+    this.maxBuffer = config.maxBuffer ?? DEFAULT_MAX_BUFFER;
+    // Auto-generate a per-process consumer group by default so each pipequery
+    // instance sees the full firehose. Users who want load-balanced consumption
+    // across replicas can set groupId explicitly in the yaml.
+    this.groupId = config.groupId ?? `pipequery-${randomBytes(6).toString('hex')}`;
+  }
+
+  async start(): Promise<void> {
+    const brokers = normalizeBrokers(this.config.brokers);
+    if (brokers.length === 0) {
+      throw new Error('Kafka source has no brokers configured');
+    }
+
+    this.kafka = new Kafka({
+      clientId: `pipequery-${this.config.topic}`,
+      brokers,
+      ssl: this.config.ssl ?? false,
+      sasl: this.config.sasl
+        ? ({
+            mechanism: this.config.sasl.mechanism,
+            username: expandEnv(this.config.sasl.username),
+            password: expandEnv(this.config.sasl.password),
+          } as SASLOptions)
+        : undefined,
+      // Silence kafkajs's own logger — we surface errors via SourceStatus.
+      logLevel: logLevel.ERROR,
+    });
+
+    this.consumer = this.kafka.consumer({ groupId: this.groupId });
+    await this.consumer.connect();
+    await this.consumer.subscribe({
+      topic: this.config.topic,
+      fromBeginning: this.config.fromBeginning ?? false,
+    });
+    this.connected = true;
+    this.error = undefined;
+
+    // Fire-and-forget; kafkajs's run() resolves immediately and keeps running
+    // in the background until disconnect().
+    void this.consumer.run({
+      eachMessage: async ({ message, partition }) => {
+        try {
+          const row = this.decode(message.value, message.key, message.timestamp, partition, message.offset);
+          this.data.push(row);
+          if (this.data.length > this.maxBuffer) {
+            this.data = this.data.slice(-this.maxBuffer);
+          }
+          this.lastFetch = new Date();
+          this.error = undefined;
+          this.notify();
+        } catch {
+          // Skip malformed messages rather than crashing the consumer.
+        }
+      },
+    });
+  }
+
+  stop(): void {
+    this.connected = false;
+    if (this.consumer) {
+      const consumer = this.consumer;
+      this.consumer = null;
+      // Fire-and-forget disconnect; SourceAdapter.stop() is sync, matching
+      // the rest of the adapters.
+      void consumer.disconnect().catch(() => {
+        // Swallow — we're shutting down anyway.
+      });
+    }
+    this.kafka = null;
+  }
+
+  getData(): unknown[] {
+    return this.data;
+  }
+
+  getStatus(): SourceStatus {
+    return {
+      healthy: this.connected && !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);
+  }
+
+  private decode(
+    value: Buffer | null,
+    key: Buffer | null,
+    timestamp: string,
+    partition: number,
+    offset: string,
+  ): Record<string, unknown> {
+    const base: Record<string, unknown> = {
+      _kafka_topic: this.config.topic,
+      _kafka_partition: partition,
+      _kafka_offset: offset,
+      _kafka_timestamp: Number.parseInt(timestamp, 10),
+      _kafka_key: key ? key.toString('utf8') : null,
+    };
+
+    if (!value) {
+      return { ...base, value: null };
+    }
+
+    const format = this.config.valueFormat ?? 'json';
+    if (format === 'string') {
+      return { ...base, value: value.toString('utf8') };
+    }
+    if (format === 'raw') {
+      return { ...base, value: value.toString('base64') };
+    }
+    // json
+    const decoded = JSON.parse(value.toString('utf8')) as unknown;
+    if (decoded !== null && typeof decoded === 'object' && !Array.isArray(decoded)) {
+      // Spread the message body into the top-level row so queries can reference
+      // fields directly (e.g. `events | where(event_type == 'order.paid')`).
+      return { ...base, ...(decoded as Record<string, unknown>) };
+    }
+    return { ...base, value: decoded };
+  }
+
+  private notify(): void {
+    for (const cb of this.listeners) cb(this.data);
+  }
+}

package/src/server/sources/manager.ts

@@ -4,6 +4,10 @@ import { RestSourceAdapter } from './rest.js';
 import { WebSocketSourceAdapter } from './websocket.js';
 import { FileSourceAdapter } from './file.js';
 import { StaticSourceAdapter } from './static.js';
+import { PostgresSourceAdapter } from './postgres.js';
+import { MysqlSourceAdapter } from './mysql.js';
+import { SqliteSourceAdapter } from './sqlite.js';
+import { KafkaSourceAdapter } from './kafka.js';
 
 export type DataContext = Record<string, unknown[]>;
 
@@ -113,6 +117,14 @@ export class SourceManager {
         return new FileSourceAdapter(config, this.cwd);
       case 'static':
         return new StaticSourceAdapter(config);
+      case 'postgres':
+        return new PostgresSourceAdapter(config);
+      case 'mysql':
+        return new MysqlSourceAdapter(config);
+      case 'sqlite':
+        return new SqliteSourceAdapter(config, this.cwd);
+      case 'kafka':
+        return new KafkaSourceAdapter(config);
       default:
         throw new Error(`Unknown source type: ${(config as { type: string }).type}`);
     }