woolsocks-bigquery-mcp 2.0.0

package/proxy-client.js

@@ -0,0 +1,303 @@
+// HTTP client for woolsocks-ai-proxy BigQuery endpoints
+// Includes retry logic, circuit breaker, and rate limiting
+
+import {
+  PROXY_URL,
+  getAuthTokens,
+  REQUEST_TIMEOUT_MS,
+  RATE_LIMIT_MS,
+  RETRY_CONFIG,
+  CIRCUIT_BREAKER_CONFIG
+} from './config.js';
+
+const sleep = (ms) => new Promise(resolve => setTimeout(resolve, ms));
+
+export class ProxyApiError extends Error {
+  constructor(status, message, endpoint) {
+    super(`Proxy API error (${status}): ${message}`);
+    this.status = status;
+    this.endpoint = endpoint;
+    this.name = 'ProxyApiError';
+  }
+}
+
+class CircuitBreaker {
+  constructor(config) {
+    this.failureThreshold = config.failureThreshold;
+    this.cooldownMs = config.cooldownMs;
+    this.failures = 0;
+    this.state = 'CLOSED';  // CLOSED | OPEN | HALF_OPEN
+    this.lastFailureTime = null;
+  }
+
+  canRequest() {
+    if (this.state === 'CLOSED') return true;
+
+    if (this.state === 'OPEN') {
+      // Check if cooldown has passed
+      if (Date.now() - this.lastFailureTime >= this.cooldownMs) {
+        this.state = 'HALF_OPEN';
+        return true;
+      }
+      return false;
+    }
+
+    // HALF_OPEN - allow one request through
+    return true;
+  }
+
+  recordSuccess() {
+    this.failures = 0;
+    this.state = 'CLOSED';
+  }
+
+  recordFailure() {
+    this.failures++;
+    this.lastFailureTime = Date.now();
+
+    if (this.failures >= this.failureThreshold) {
+      this.state = 'OPEN';
+      console.error(`[CircuitBreaker] Circuit OPEN after ${this.failures} failures`);
+    }
+  }
+
+  getState() {
+    return {
+      state: this.state,
+      failures: this.failures,
+      lastFailureTime: this.lastFailureTime,
+    };
+  }
+}
+
+export class ProxyClient {
+  constructor({ identityToken, accessToken }) {
+    if (!identityToken || !accessToken) {
+      throw new Error(
+        'GCP tokens are required. Run `gcloud auth login` and `gcloud auth application-default login`.'
+      );
+    }
+    this.baseUrl = PROXY_URL;
+    this.identityToken = identityToken;  // Authenticates to proxy (Cloud Run IAM)
+    this.accessToken = accessToken;      // Passed through to BigQuery (your own IAM access)
+    this.lastRequestTime = 0;
+    this.circuitBreaker = new CircuitBreaker(CIRCUIT_BREAKER_CONFIG);
+  }
+
+  /**
+   * Create a new ProxyClient with async initialization.
+   * Fetches the current user's GCP tokens via gcloud.
+   */
+  static async create() {
+    const tokens = await getAuthTokens();
+    return new ProxyClient(tokens);
+  }
+
+  /**
+   * Rate limit requests
+   */
+  async rateLimit() {
+    const now = Date.now();
+    const timeSinceLastRequest = now - this.lastRequestTime;
+    if (timeSinceLastRequest < RATE_LIMIT_MS) {
+      await sleep(RATE_LIMIT_MS - timeSinceLastRequest);
+    }
+    this.lastRequestTime = Date.now();
+  }
+
+  /**
+   * Calculate retry delay with exponential backoff
+   */
+  getRetryDelay(attempt) {
+    const delay = Math.min(
+      RETRY_CONFIG.baseDelayMs * Math.pow(2, attempt),
+      RETRY_CONFIG.maxDelayMs
+    );
+    // Add jitter (0-25% of delay)
+    return delay + Math.random() * delay * 0.25;
+  }
+
+  /**
+   * Check if status code is retryable
+   */
+  isRetryable(status) {
+    return RETRY_CONFIG.retryableStatuses.includes(status);
+  }
+
+  /**
+   * Make HTTP request with retry and circuit breaker
+   */
+  async request(endpoint, options = {}) {
+    // Check circuit breaker
+    if (!this.circuitBreaker.canRequest()) {
+      throw new ProxyApiError(
+        503,
+        'Circuit breaker is OPEN. Proxy may be unavailable. Will retry in 30s.',
+        endpoint
+      );
+    }
+
+    let lastError;
+
+    for (let attempt = 0; attempt < RETRY_CONFIG.maxAttempts; attempt++) {
+      try {
+        await this.rateLimit();
+
+        const url = `${this.baseUrl}${endpoint}`;
+        const controller = new AbortController();
+        const timeoutId = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
+
+        const response = await fetch(url, {
+          ...options,
+          signal: controller.signal,
+          headers: {
+            'Content-Type': 'application/json',
+            // Identity token: proves who you are to the proxy (Cloud Run IAM auth)
+            'Authorization': `Bearer ${this.identityToken}`,
+            // Access token: proxy passes this to BigQuery so queries run as your GCP identity
+            'X-BQ-Access-Token': this.accessToken,
+            ...options.headers
+          }
+        });
+
+        clearTimeout(timeoutId);
+
+        if (!response.ok) {
+          const errorData = await response.json().catch(() => ({
+            error: response.statusText
+          }));
+
+          const error = new ProxyApiError(
+            response.status,
+            errorData.error || errorData.message || response.statusText,
+            endpoint
+          );
+
+          // Don't retry client errors (4xx) except rate limits
+          if (response.status < 500 && response.status !== 429) {
+            this.circuitBreaker.recordSuccess(); // Not a server issue
+            throw error;
+          }
+
+          lastError = error;
+
+          if (this.isRetryable(response.status) && attempt < RETRY_CONFIG.maxAttempts - 1) {
+            const delay = this.getRetryDelay(attempt);
+            console.error(`[ProxyClient] Retrying in ${Math.round(delay)}ms (attempt ${attempt + 1}/${RETRY_CONFIG.maxAttempts})`);
+            await sleep(delay);
+            continue;
+          }
+
+          this.circuitBreaker.recordFailure();
+          throw error;
+        }
+
+        // Success
+        this.circuitBreaker.recordSuccess();
+        return response.json();
+
+      } catch (error) {
+        if (error.name === 'AbortError') {
+          lastError = new ProxyApiError(408, 'Request timeout', endpoint);
+        } else if (error instanceof ProxyApiError) {
+          lastError = error;
+        } else {
+          lastError = new ProxyApiError(0, error.message, endpoint);
+        }
+
+        // Retry on network errors
+        if (attempt < RETRY_CONFIG.maxAttempts - 1 && !(error instanceof ProxyApiError && error.status >= 400 && error.status < 500)) {
+          const delay = this.getRetryDelay(attempt);
+          console.error(`[ProxyClient] Network error, retrying in ${Math.round(delay)}ms: ${error.message}`);
+          await sleep(delay);
+          continue;
+        }
+
+        this.circuitBreaker.recordFailure();
+        throw lastError;
+      }
+    }
+
+    throw lastError;
+  }
+
+  // =========================================================================
+  // Schema Discovery Methods
+  // =========================================================================
+
+  async listDatasets() {
+    return this.request('/v1/bigquery/datasets');
+  }
+
+  async listTables(datasetId) {
+    return this.request(`/v1/bigquery/datasets/${encodeURIComponent(datasetId)}/tables`);
+  }
+
+  async getTableSchema(datasetId, tableId) {
+    return this.request(
+      `/v1/bigquery/datasets/${encodeURIComponent(datasetId)}/tables/${encodeURIComponent(tableId)}`
+    );
+  }
+
+  async previewTable(datasetId, tableId, maxRows = 10) {
+    return this.request(
+      `/v1/bigquery/datasets/${encodeURIComponent(datasetId)}/tables/${encodeURIComponent(tableId)}/preview?maxRows=${maxRows}`
+    );
+  }
+
+  // =========================================================================
+  // Query Execution Methods
+  // =========================================================================
+
+  async executeQuery(sql) {
+    return this.request('/v1/bigquery/query', {
+      method: 'POST',
+      body: JSON.stringify({ sql })
+    });
+  }
+
+  async validateQuery(sql) {
+    return this.request('/v1/bigquery/validate', {
+      method: 'POST',
+      body: JSON.stringify({ sql })
+    });
+  }
+
+  async getStatus() {
+    return this.request('/v1/bigquery/status');
+  }
+
+  // =========================================================================
+  // Utility Methods
+  // =========================================================================
+
+  getCircuitBreakerState() {
+    return this.circuitBreaker.getState();
+  }
+}
+
+// Singleton instance
+let client = null;
+let clientPromise = null;
+
+/**
+ * Get or create the ProxyClient singleton.
+ * Initializes async on first call (fetches API key from Secret Manager).
+ */
+export async function getClient() {
+  if (client) {
+    return client;
+  }
+
+  // Prevent multiple simultaneous initializations
+  if (!clientPromise) {
+    clientPromise = ProxyClient.create().then(c => {
+      client = c;
+      return c;
+    });
+  }
+
+  return clientPromise;
+}
+
+export default ProxyClient;

package/secrets.js

@@ -0,0 +1,85 @@
+// GCP token auth for BigQuery MCP
+// Uses the colleague's own gcloud credentials instead of a shared API key.
+// This means BigQuery access is controlled per-person by infra via GCP IAM.
+//
+// Two tokens are fetched:
+//   identityToken  → authenticates to woolsocks-ai-proxy (proves who you are)
+//   accessToken    → passed through to BigQuery so queries run as YOUR GCP identity
+//
+// Prerequisites:
+//   gcloud auth login
+//   gcloud auth application-default login
+//   gcloud must be in PATH
+
+import { exec } from 'child_process';
+import { promisify } from 'util';
+
+const execAsync = promisify(exec);
+
+const PROXY_URL = process.env.BIGQUERY_PROXY_URL ||
+  'https://woolsocks-ai-proxy-1009713156898.europe-west1.run.app';
+
+// Cache tokens for 45 minutes (gcloud tokens expire after 60 min)
+let cachedTokens = null;
+let cacheExpiry = 0;
+const TOKEN_TTL_MS = 45 * 60 * 1000;
+
+/**
+ * Get GCP auth tokens for the current gcloud user.
+ * Returns { identityToken, accessToken }.
+ *
+ * Env var overrides for testing:
+ *   BIGQUERY_IDENTITY_TOKEN, BIGQUERY_ACCESS_TOKEN
+ */
+export async function getAuthTokens() {
+  const now = Date.now();
+
+  // Return cached tokens if still valid
+  if (cachedTokens && now < cacheExpiry) {
+    return cachedTokens;
+  }
+
+  // Env var override for local testing / CI
+  if (process.env.BIGQUERY_IDENTITY_TOKEN && process.env.BIGQUERY_ACCESS_TOKEN) {
+    console.error('[auth] Using tokens from environment variables');
+    cachedTokens = {
+      identityToken: process.env.BIGQUERY_IDENTITY_TOKEN,
+      accessToken: process.env.BIGQUERY_ACCESS_TOKEN,
+    };
+    cacheExpiry = now + TOKEN_TTL_MS;
+    return cachedTokens;
+  }
+
+  try {
+    // Identity token: authenticates to the proxy (Cloud Run IAM)
+    const { stdout: identityRaw } = await execAsync(
+      `gcloud auth print-identity-token --audiences=${PROXY_URL}`
+    );
+
+    // Access token: passed through to BigQuery so queries run as your own identity
+    const { stdout: accessRaw } = await execAsync(
+      'gcloud auth print-access-token'
+    );
+
+    cachedTokens = {
+      identityToken: identityRaw.trim(),
+      accessToken: accessRaw.trim(),
+    };
+    cacheExpiry = now + TOKEN_TTL_MS;
+
+    console.error('[auth] GCP tokens loaded from gcloud');
+    return cachedTokens;
+
+  } catch (error) {
+    console.error('[auth] Failed to get gcloud tokens:', error.message);
+    throw new Error(
+      'Could not get GCP credentials.\n' +
+      'Run the following and try again:\n' +
+      '  gcloud auth login\n' +
+      '  gcloud auth application-default login\n\n' +
+      'No BigQuery access yet? Ask infra via a Jira PERM ticket:\n' +
+      'https://woolsocks.atlassian.net/jira/software/projects/PERM/boards\n' +
+      '→ Request: "BigQuery read access for Claude Code"'
+    );
+  }
+}

package/tools/index.js

@@ -0,0 +1,15 @@
+// Tool registry - combines all tools from all modules
+
+import schemaTools from './schema/index.js';
+import queryTools from './query/index.js';
+
+// Combine all tools into a single array
+export const TOOLS = [
+  ...schemaTools,
+  ...queryTools
+];
+
+// Create a map for quick lookup by name
+export const TOOL_MAP = new Map(TOOLS.map(tool => [tool.name, tool]));
+
+export default TOOLS;

package/tools/query/execute-query.js

@@ -0,0 +1,115 @@
+// Execute a BigQuery query
+
+import { getClient } from '../../proxy-client.js';
+import { MAX_QUERY_ROWS } from '../../config.js';
+
+export default {
+  name: 'bigquery__execute-query',
+  description: `Execute a BigQuery SQL query. IMPORTANT RULES — follow these every time:
+1. ALWAYS call bigquery__validate-query FIRST to check cost before executing.
+2. ALWAYS include a LIMIT clause (max 1000 rows for exploration).
+3. ALWAYS add a date/time filter on any table larger than 1 GB (use WHERE created_at, _PARTITIONTIME, or similar).
+4. NEVER use SELECT * on large tables — select only the columns you need.
+5. If you get a 401/403 error, the user does not have BigQuery access — tell them to request it via a Jira ticket in the PERM project (see error hint).
+Guardrails enforced server-side: 10 GB per-query scan limit, 50 GB daily limit, EU region only. Max ${MAX_QUERY_ROWS} rows returned.`,
+  inputSchema: {
+    type: 'object',
+    properties: {
+      sql: {
+        type: 'string',
+        description: 'The SQL query to execute. Use fully qualified table names (dataset.table). Always add date/partition filters.'
+      }
+    },
+    required: ['sql']
+  },
+
+  async handler(args) {
+    const { sql } = args;
+
+    if (!sql || typeof sql !== 'string') {
+      return {
+        content: [{
+          type: 'text',
+          text: JSON.stringify({
+            error: 'sql parameter is required and must be a string',
+            hint: 'Use bigquery__get-table-schema to understand table structure before querying'
+          }, null, 2)
+        }],
+        isError: true
+      };
+    }
+
+    const client = await getClient();
+    const result = await client.executeQuery(sql);
+
+    // Handle blocked queries
+    if (result.blocked) {
+      return {
+        content: [{
+          type: 'text',
+          text: JSON.stringify({
+            blocked: true,
+            reason: result.reason,
+            message: result.message,
+            details: {
+              bytesScanned: result.bytesScanned,
+              estimatedGB: result.estimatedGB,
+              scanLimitGB: result.scanLimitGB,
+              dailyUsedGB: result.dailyUsedGB,
+              dailyLimitGB: result.dailyLimitGB,
+            },
+            hint: 'Reduce query cost by: (1) adding a date filter e.g. WHERE created_at >= DATE_SUB(CURRENT_DATE(), INTERVAL 30 DAY), (2) selecting specific columns instead of *, (3) adding LIMIT 1000.'
+          }, null, 2)
+        }],
+        isError: true
+      };
+    }
+
+    // Handle execution failures
+    if (!result.success) {
+      const isAuthError = result.error?.includes('401') || result.error?.includes('403') ||
+        result.error?.includes('permission') || result.error?.includes('Access Denied');
+      return {
+        content: [{
+          type: 'text',
+          text: JSON.stringify({
+            success: false,
+            error: result.error,
+            message: result.message,
+            hint: isAuthError
+              ? 'You do not have BigQuery access. Ask your infra manager to grant access by creating a ticket in the Jira PERM project: https://woolsocks.atlassian.net/jira/software/projects/PERM/boards — request "BigQuery read access for Claude Code".'
+              : 'Check SQL syntax and table names. Use bigquery__get-table-schema to verify column names.'
+          }, null, 2)
+        }],
+        isError: true
+      };
+    }
+
+    // Truncate results if too many rows
+    let rows = result.rows;
+    let truncated = false;
+    if (rows.length > MAX_QUERY_ROWS) {
+      rows = rows.slice(0, MAX_QUERY_ROWS);
+      truncated = true;
+    }
+
+    return {
+      content: [{
+        type: 'text',
+        text: JSON.stringify({
+          success: true,
+          rowCount: result.rowCount,
+          returnedRows: rows.length,
+          truncated,
+          bytesScanned: result.bytesScanned,
+          estimatedGB: result.estimatedGB,
+          dailyUsedGB: result.dailyUsedGB,
+          dailyLimitGB: result.dailyLimitGB,
+          costNote: `Scanned ${result.estimatedGB}GB. Daily usage: ${result.dailyUsedGB}GB / ${result.dailyLimitGB}GB`,
+          rows,
+          ...(truncated && { warning: `Results truncated to ${MAX_QUERY_ROWS} rows. Add LIMIT to your query to control output size.` })
+        }, null, 2)
+      }]
+    };
+  }
+};

package/tools/query/get-status.js

@@ -0,0 +1,71 @@
+// Get BigQuery usage status
+
+import { getClient } from '../../proxy-client.js';
+import { schemaCache } from '../../cache.js';
+
+export default {
+  name: 'bigquery__get-status',
+  description: 'Get current BigQuery usage status including daily/monthly limits, scan counts, and remaining budget. Also shows circuit breaker state and cache stats. Check this before running large queries.',
+  inputSchema: {
+    type: 'object',
+    properties: {},
+    required: []
+  },
+
+  async handler(args) {
+    const client = await getClient();
+    const status = await client.getStatus();
+
+    // Get client state
+    const circuitState = client.getCircuitBreakerState();
+    const cacheStats = schemaCache.getStats();
+
+    // Build status message
+    const lines = [
+      'BigQuery Status:',
+      `  Ready: ${status.ready ? 'Yes' : 'No'}`,
+      '',
+      'Daily Usage:',
+      `  Scanned: ${status.daily.usedGB} GB / ${status.daily.limitGB} GB (${status.daily.percentUsed}%)`,
+      `  Queries: ${status.daily.queryCount}`,
+      '',
+      'Monthly Usage:',
+      `  Cost: EUR ${status.monthly.spentEUR} / EUR ${status.monthly.limitEUR} (${status.monthly.percentUsed}%)`,
+      `  Queries: ${status.monthly.queryCount}`,
+      '',
+      'Limits:',
+      `  Per query: ${status.limits.perQueryGB} GB`,
+      `  Daily: ${status.limits.dailyGB} GB`,
+      `  Monthly: EUR ${status.limits.monthlyEUR}`,
+    ];
+
+    // Add warnings if close to limits
+    let warning = null;
+    if (status.daily.percentUsed >= 90) {
+      warning = 'CRITICAL: Daily limit nearly exhausted! Wait until tomorrow for large queries.';
+    } else if (status.daily.percentUsed >= 80) {
+      warning = 'WARNING: Daily limit at 80%. Use queries sparingly.';
+    } else if (status.monthly.percentUsed >= 90) {
+      warning = 'CRITICAL: Monthly budget nearly exhausted!';
+    } else if (status.monthly.percentUsed >= 80) {
+      warning = 'WARNING: Monthly budget at 80%.';
+    }
+
+    return {
+      content: [{
+        type: 'text',
+        text: JSON.stringify({
+          ready: status.ready,
+          error: status.error,
+          limits: status.limits,
+          daily: status.daily,
+          monthly: status.monthly,
+          warning,
+          statusMessage: lines.join('\n'),
+          circuitBreaker: circuitState,
+          cache: cacheStats,
+        }, null, 2)
+      }]
+    };
+  }
+};

package/tools/query/index.js

@@ -0,0 +1,11 @@
+// Query execution tools
+
+import executeQuery from './execute-query.js';
+import validateQuery from './validate-query.js';
+import getStatus from './get-status.js';
+
+export default [
+  executeQuery,
+  validateQuery,
+  getStatus
+];

package/tools/query/validate-query.js

@@ -0,0 +1,83 @@
+// Validate a query without executing (dry-run)
+
+import { getClient } from '../../proxy-client.js';
+
+export default {
+  name: 'bigquery__validate-query',
+  description: 'Validate a BigQuery query without executing it. Returns estimated bytes to scan and checks against cost limits. ALWAYS call this before bigquery__execute-query — it is free (no data scanned) and prevents surprise cost overruns. If validation fails, revise the query (add date filter, LIMIT, or fewer columns) before executing.',
+  inputSchema: {
+    type: 'object',
+    properties: {
+      sql: {
+        type: 'string',
+        description: 'The SQL query to validate'
+      }
+    },
+    required: ['sql']
+  },
+
+  async handler(args) {
+    const { sql } = args;
+
+    if (!sql || typeof sql !== 'string') {
+      return {
+        content: [{
+          type: 'text',
+          text: JSON.stringify({
+            error: 'sql parameter is required and must be a string',
+            hint: 'Provide a valid SQL query to validate'
+          }, null, 2)
+        }],
+        isError: true
+      };
+    }
+
+    const client = await getClient();
+    const validation = await client.validateQuery(sql);
+
+    if (!validation.approved) {
+      const isAuthError = validation.reason === 'unauthorized' || validation.reason === 'forbidden' ||
+        validation.message?.toLowerCase().includes('permission') ||
+        validation.message?.toLowerCase().includes('access denied');
+      return {
+        content: [{
+          type: 'text',
+          text: JSON.stringify({
+            approved: false,
+            reason: validation.reason,
+            message: validation.message,
+            details: {
+              bytesScanned: validation.bytesScanned,
+              estimatedGB: validation.estimatedGB,
+              scanLimitGB: validation.scanLimitGB,
+              dailyUsedGB: validation.dailyUsedGB,
+              dailyLimitGB: validation.dailyLimitGB,
+              tables: validation.tables,
+            },
+            hint: isAuthError
+              ? 'You do not have BigQuery access. Ask your infra manager to grant access by creating a ticket in the Jira PERM project: https://woolsocks.atlassian.net/jira/software/projects/PERM/boards — request "BigQuery read access for Claude Code".'
+              : 'Reduce query cost: (1) add a date filter e.g. WHERE created_at >= DATE_SUB(CURRENT_DATE(), INTERVAL 30 DAY), (2) select specific columns instead of *, (3) add LIMIT 1000.'
+          }, null, 2)
+        }],
+        isError: true
+      };
+    }
+
+    return {
+      content: [{
+        type: 'text',
+        text: JSON.stringify({
+          approved: true,
+          estimatedGB: validation.estimatedGB,
+          scanLimitGB: validation.scanLimitGB,
+          dailyUsedGB: validation.dailyUsedGB,
+          dailyLimitGB: validation.dailyLimitGB,
+          remainingDailyGB: validation.dailyLimitGB - validation.dailyUsedGB,
+          tables: validation.tables,
+          message: `Query approved. Will scan ~${validation.estimatedGB}GB (limit: ${validation.scanLimitGB}GB per query).`,
+          hint: 'Safe to execute with bigquery__execute-query'
+        }, null, 2)
+      }]
+    };
+  }
+};