s3-querier 1.0.3 → 1.1.0

package/README.md

@@ -159,6 +159,171 @@ const results = await s3Querier({
 });
 ```
 
+## MCP Server
+
+s3-querier ships a [Model Context Protocol](https://modelcontextprotocol.io/) server that exposes two tools to any MCP-compatible client (Claude Desktop, Claude Code, etc.):
+
+- **`query`** — runs a DuckDB SQL query against your S3 data
+- **`list_files`** — lists objects under a prefix so an LLM can discover available data
+
+### Environment variables
+
+| Variable | Required | Description |
+| --- | --- | --- |
+| `S3_ENDPOINT` | ✓ | S3 endpoint URL |
+| `S3_BUCKET` | ✓ | Default bucket |
+| `S3_ACCESS_KEY_ID` | ✓ * | HMAC access key |
+| `S3_SECRET_ACCESS_KEY` | ✓ * | HMAC secret key |
+| `S3_API_KEY` | ✓ * | IBM IAM API key (alternative to HMAC) |
+| `S3_BUCKETS_DIR` | | Local cache directory (default `/tmp/s3-querier`) |
+
+\* Either HMAC pair **or** `S3_API_KEY` is required.
+
+### Basic server
+
+The built-in server entry point requires no configuration beyond environment variables.
+
+**Claude Code / Claude Desktop**
+
+```bash
+claude mcp add s3-querier \
+  -e S3_ENDPOINT=https://s3.amazonaws.com \
+  -e S3_BUCKET=my-bucket \
+  -e S3_ACCESS_KEY_ID=key \
+  -e S3_SECRET_ACCESS_KEY=secret \
+  -- node node_modules/s3-querier/src/mcp/server.js
+```
+
+**IBM Bob**
+
+Add to `mcp_settings.json` (global, applies across all workspaces) or `.Bob/mcp.json` (project-level, committed with your repo):
+
+```json
+{
+  "mcpServers": {
+    "s3-querier": {
+      "command": "node",
+      "args": ["/absolute/path/to/node_modules/s3-querier/src/mcp/server.js"],
+      "env": {
+        "S3_ENDPOINT": "https://s3.amazonaws.com",
+        "S3_BUCKET": "my-bucket",
+        "S3_ACCESS_KEY_ID": "key",
+        "S3_SECRET_ACCESS_KEY": "secret"
+      },
+      "disabled": false,
+      "alwaysAllow": [],
+      "disabledTools": []
+    }
+  }
+}
+```
+
+### Extending with `S3QuerierMCP`
+
+For richer LLM context, create a custom server using `S3QuerierMCP` and pass a `datasets` array. Each entry describes a dataset so the model knows what data is available and how it's structured — without having to explore the bucket first.
+
+```js
+// my-server.js
+import { S3QuerierMCP } from 's3-querier/src/mcp/s3querier-mcp.js';
+
+new S3QuerierMCP({
+  datasets: [
+    {
+      name: 'sales',
+      description: 'Monthly sales transactions partitioned by year and month.',
+      prefix: 'sales/',
+      partitioning: 'year/month',
+      files: {
+        data: {
+          description: 'Sales records — id (int), date, product, amount (float), region',
+        },
+      },
+    },
+    {
+      name: 'products',
+      description: 'Product catalog — static reference data, no partitioning.',
+      prefix: 'products/',
+      files: {
+        catalog: {
+          description: 'Products — name, category, price (float)',
+        },
+      },
+    },
+  ],
+}).start();
+```
+
+**Claude Code / Claude Desktop**
+
+```bash
+claude mcp add my-datalake \
+  -e S3_ENDPOINT=https://s3.amazonaws.com \
+  -e S3_BUCKET=my-bucket \
+  -e S3_ACCESS_KEY_ID=key \
+  -e S3_SECRET_ACCESS_KEY=secret \
+  -- node my-server.js
+```
+
+**IBM Bob**
+
+```json
+{
+  "mcpServers": {
+    "my-datalake": {
+      "command": "node",
+      "args": ["/absolute/path/to/my-server.js"],
+      "env": {
+        "S3_ENDPOINT": "https://s3.amazonaws.com",
+        "S3_BUCKET": "my-bucket",
+        "S3_ACCESS_KEY_ID": "key",
+        "S3_SECRET_ACCESS_KEY": "secret"
+      },
+      "disabled": false,
+      "alwaysAllow": [],
+      "disabledTools": []
+    }
+  }
+}
+```
+
+#### Dataset options
+
+| Field | Description |
+| --- | --- |
+| `name` | Dataset identifier |
+| `description` | Narrative description injected into the tool prompt |
+| `prefix` | S3 path prefix (e.g. `"sales/"`) |
+| `partitioning` | Partitioning scheme hint (e.g. `"year/month"`) |
+| `bucket` | Overrides `S3_BUCKET` for this dataset |
+| `endpoint` | Overrides `S3_ENDPOINT` for this dataset |
+| `files` | Map of logical file names to `{ description }` |
+
+### Adding custom tools
+
+Pass a `tools` array to register additional MCP tools alongside the built-in ones:
+
+```js
+import { z } from 'zod';
+import { S3QuerierMCP } from 's3-querier/src/mcp/s3querier-mcp.js';
+
+new S3QuerierMCP({
+  datasets: [ /* ... */ ],
+  tools: [
+    {
+      name: 'get_report',
+      description: 'Returns the latest weekly summary report.',
+      inputSchema: {
+        week: z.string().describe('ISO week string, e.g. "2025-W03"'),
+      },
+      handler: async ({ week }) => {
+        // your logic here
+        return { content: [{ type: 'text', text: `Report for ${week}` }] };
+      },
+    },
+  ],
+}).start();
+```
+
 ## Examples
 
 The `examples/` directory contains a local interactive demo and standalone scripts. All examples target a local MinIO instance — you'll need [Docker](https://docs.docker.com/get-docker/) and [Docker Compose](https://docs.docker.com/compose/install/) installed. Both are bundled with Docker Desktop on Mac and Windows; on Linux, install the Compose plugin separately.

package/package.json

@@ -1,11 +1,15 @@
 {
   "name": "s3-querier",
-  "version": "1.0.3",
+  "version": "1.1.0",
   "description": "Query S3-compatible storage with DuckDB and SQL",
   "type": "module",
   "main": "src/s3-querier.js",
+  "bin": {
+    "s3-querier-mcp": "src/mcp/server.js"
+  },
   "exports": {
-    ".": "./src/s3-querier.js"
+    ".": "./src/s3-querier.js",
+    "./mcp": "./src/mcp/datalake-mcp.js"
   },
   "files": [
     "src/**/*.js",
@@ -44,8 +48,10 @@
     "@aws-sdk/client-s3": "^3.0.0",
     "@derekstride/tree-sitter-sql": "^0.3.11",
     "@duckdb/node-api": "^1.5.3-r.3",
+    "@modelcontextprotocol/sdk": "^1.29.0",
     "avsc": "^5.7.7",
     "date-fns": "^4.0.0",
+    "hyparquet": "^1.26.0",
     "lru-cache": "^11.0.0",
     "peggy": "^5.1.0",
     "pino": "^10.3.1",

package/src/mcp/handlers/list-files.js

@@ -0,0 +1,63 @@
+import { ListObjectsV2Command } from '@aws-sdk/client-s3';
+import { bigintReplacer } from '../../s3-querier.js';
+import { buildS3Client } from '../../s3/s3.js';
+import { readParquetColumns } from '../../utils/parquet-schema-reader.js';
+
+const { S3_ACCESS_KEY_ID, S3_SECRET_ACCESS_KEY, S3_API_KEY, S3_ENDPOINT, S3_BUCKET } = process.env;
+
+export async function handleListFiles({ prefix = '', maxResults = 100, endpoint, bucket }) {
+  const resolvedEndpoint = endpoint || S3_ENDPOINT;
+  const resolvedBucket = bucket || S3_BUCKET;
+  const clientConfig = {
+    apiKey: S3_API_KEY,
+    accessKeyId: S3_ACCESS_KEY_ID,
+    secretAccessKey: S3_SECRET_ACCESS_KEY,
+    endpoint: resolvedEndpoint,
+  };
+  const s3Client = buildS3Client(clientConfig);
+  const listCommand = new ListObjectsV2Command({
+    Bucket: resolvedBucket,
+    Prefix: prefix,
+    MaxKeys: maxResults,
+    Delimiter: '/',
+  });
+  const response = await s3Client.send(listCommand);
+  const directories = (response.CommonPrefixes ?? []).map(({ Prefix }) => Prefix);
+  const files = (response.Contents ?? []).map(({ Key, Size }) => ({ file: Key, size: Size }));
+  const truncated = response.IsTruncated ?? false;
+  const representatives = getRepresentativeFiles(files);
+  const filesWithSchema = await Promise.all(
+    files.map((fileObj) => maybeAddSchema(s3Client, resolvedBucket, representatives, fileObj)),
+  );
+
+  return {
+    content: [
+      { type: 'text', text: JSON.stringify({ directories, files: filesWithSchema, truncated }, bigintReplacer) },
+    ],
+  };
+}
+
+/** Helpers */
+
+async function addSchema(s3Client, bucket, { file, size }) {
+  if (!file.endsWith('.parquet')) return { file, size };
+  const columns = await readParquetColumns(s3Client, bucket, file).catch(() => null);
+  return { file, size, columns };
+}
+
+function maybeAddSchema(s3Client, bucket, representatives, fileObj) {
+  if (representatives.has(fileObj.file)) return addSchema(s3Client, bucket, fileObj);
+  return Promise.resolve(fileObj);
+}
+
+function getRepresentativeFiles(files) {
+  const parquetFiles = files.filter(({ file }) => file.endsWith('.parquet'));
+  const dirMap = parquetFiles.reduce(addFirstFilePerDir, new Map());
+  return new Set(dirMap.values());
+}
+
+function addFirstFilePerDir(acc, { file }) {
+  const dir = file.substring(0, file.lastIndexOf('/'));
+  if (!acc.has(dir)) acc.set(dir, file);
+  return acc;
+}

package/src/mcp/handlers/query.js

@@ -0,0 +1,34 @@
+import s3Querier, { bigintReplacer } from '../../s3-querier.js';
+
+const {
+  S3_ACCESS_KEY_ID,
+  S3_SECRET_ACCESS_KEY,
+  S3_API_KEY,
+  S3_ENDPOINT,
+  S3_BUCKET,
+  S3_BUCKETS_DIR = '/tmp/s3-querier',
+} = process.env;
+
+export async function handleQuery({ sql, from, to, endpoint, bucket }) {
+  const fromMs = from ? new Date(from).getTime() : undefined;
+  const toMs = to ? new Date(to).getTime() : undefined;
+  const resolvedEndpoint = endpoint || S3_ENDPOINT;
+  const resolvedBucket = bucket || S3_BUCKET;
+
+  const results = await s3Querier({
+    query: sql,
+    from: fromMs,
+    to: toMs,
+    defaultEndpoint: resolvedEndpoint,
+    defaultBucket: resolvedBucket,
+    bucketsDir: S3_BUCKETS_DIR,
+    apiKey: S3_API_KEY,
+    accessKeyId: S3_ACCESS_KEY_ID,
+    secretAccessKey: S3_SECRET_ACCESS_KEY,
+    format: 'jsonRecords',
+  });
+
+  return {
+    content: [{ type: 'text', text: JSON.stringify(results, bigintReplacer) }],
+  };
+}

package/src/mcp/s3querier-mcp.js

@@ -0,0 +1,119 @@
+import { readFileSync } from 'node:fs';
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
+import { z } from 'zod';
+import { handleListFiles } from './handlers/list-files.js';
+import { handleQuery } from './handlers/query.js';
+
+const pkg = JSON.parse(readFileSync(new URL('../../package.json', import.meta.url), 'utf8'));
+
+const toolDescription = readFileSync(new URL('./descriptions/tool.md', import.meta.url), 'utf8');
+const sqlDescription = readFileSync(new URL('./descriptions/sql-param.md', import.meta.url), 'utf8');
+const listFilesTemplate = readFileSync(new URL('./descriptions/list-files.md', import.meta.url), 'utf8');
+const docsContent = readFileSync(new URL('../../docs/s3-querier.md', import.meta.url), 'utf8');
+
+const QUERY_TOOL_SCHEMA = {
+  sql: z.string().describe(sqlDescription),
+  from: z
+    .string()
+    .optional()
+    .describe('Start of date range as ISO 8601 (e.g. "2025-01-01"). Required when the query uses date tokens.'),
+  to: z
+    .string()
+    .optional()
+    .describe('End of date range as ISO 8601 (e.g. "2025-01-31"). Required when the query uses date tokens.'),
+  endpoint: z.string().optional().describe('S3 endpoint URL. Overrides S3_ENDPOINT for this query.'),
+  bucket: z.string().optional().describe('S3 bucket name. Overrides S3_BUCKET for this query.'),
+};
+
+const LIST_FILES_TOOL_SCHEMA = {
+  prefix: z
+    .string()
+    .optional()
+    .describe('Path prefix to list under (e.g. "sales/" or ""). Defaults to empty string to list all files.'),
+  maxResults: z
+    .number()
+    .int()
+    .min(1)
+    .max(1000)
+    .optional()
+    .describe('Maximum number of files to return (default 100). Increase if the response is truncated.'),
+  endpoint: z.string().optional().describe('S3 endpoint URL. Overrides S3_ENDPOINT for this call.'),
+  bucket: z.string().optional().describe('S3 bucket name. Overrides S3_BUCKET for this call.'),
+};
+
+const DOCS_RESOURCE = {
+  title: 'S3 Querier Documentation',
+  description: 'Full documentation: query planning, file tokens, location tokens, and examples.',
+  mimeType: 'text/markdown',
+};
+
+export class S3QuerierMCP {
+  constructor(config = {}) {
+    this.config = config;
+  }
+
+  async start() {
+    const server = new McpServer({ name: 's3-querier', version: pkg.version });
+    const transport = new StdioServerTransport();
+    const listFilesDescription = buildListFilesDescription(this.config);
+    const enrichedToolDescription = buildToolDescription(this.config);
+
+    server.registerResource('s3-querier-docs', 's3-querier://docs', DOCS_RESOURCE, serveDocsHandler);
+    server.registerTool(
+      'list_files',
+      { description: listFilesDescription, inputSchema: LIST_FILES_TOOL_SCHEMA },
+      handleListFiles,
+    );
+    server.registerTool('query', { description: enrichedToolDescription, inputSchema: QUERY_TOOL_SCHEMA }, handleQuery);
+    (this.config.tools ?? []).forEach(({ name, description, inputSchema, handler }) => {
+      server.registerTool(name, { description, inputSchema }, handler);
+    });
+
+    await server.connect(transport);
+  }
+}
+
+/** Helpers */
+
+function serveDocsHandler(uri) {
+  return { contents: [{ uri: uri.href, text: docsContent, mimeType: 'text/markdown' }] };
+}
+
+function buildListFilesDescription(config) {
+  const today = new Date().toISOString().slice(0, 10);
+  const withDate = listFilesTemplate.replace('{{TODAY}}', today);
+  const datasetContext = buildDatasetContext(config.datasets);
+  return datasetContext ? `${withDate}\n${datasetContext}` : withDate;
+}
+
+function buildToolDescription(config) {
+  const datasetContext = buildDatasetContext(config.datasets);
+  return datasetContext ? `${toolDescription}\n\n${datasetContext}` : toolDescription;
+}
+
+function buildDatasetContext(datasets) {
+  if (!datasets?.length) return '';
+  const datasetLines = datasets.flatMap(formatDataset);
+  return ['CONFIGURED DATASETS', '', ...datasetLines].join('\n');
+}
+
+function formatDataset({ name, description, bucket, endpoint, prefix, partitioning, files }) {
+  const header = description ? `${name} — ${description}` : name;
+  const lines = [header];
+  if (bucket) lines.push(`  Bucket: ${bucket}`);
+  if (endpoint) lines.push(`  Endpoint: ${endpoint}`);
+  if (prefix) lines.push(`  Prefix: ${prefix}`);
+  if (partitioning) lines.push(`  Partitioning: ${partitioning}`);
+  if (files) {
+    const fileLines = Object.entries(files).map(formatFileLine);
+    lines.push('  Files:', ...fileLines);
+  }
+  lines.push('');
+  return lines;
+}
+
+function formatFileLine([fileName, { description: fileDesc }]) {
+  const label = fileDesc ? `${fileName} — ${fileDesc}` : fileName;
+  return `    ${label}`;
+}

package/src/mcp/server.js

@@ -0,0 +1,5 @@
+#!/usr/bin/env node
+
+import { S3QuerierMCP } from './s3querier-mcp.js';
+
+await new S3QuerierMCP().start();

package/src/utils/parquet-schema-reader.js

@@ -0,0 +1,33 @@
+import { GetObjectCommand } from '@aws-sdk/client-s3';
+import { parquetMetadata, parquetSchema } from 'hyparquet';
+
+/**
+ * Reads the column names from a parquet file using two ranged S3 GET requests
+ * (footer only — no full file download regardless of file size).
+ *
+ * @param {import('@aws-sdk/client-s3').S3Client} s3Client
+ * @param {string} bucket
+ * @param {string} key
+ * @returns {Promise<string[]>} Column names
+ */
+export async function readParquetColumns(s3Client, bucket, key) {
+  const tail = await fetchRange(s3Client, bucket, key, 'bytes=-8');
+  const footerLength = tail.readUInt32LE(0);
+  const footerBuffer = await fetchRange(s3Client, bucket, key, `bytes=-${footerLength + 8}`);
+  const { buffer, byteOffset, byteLength } = footerBuffer;
+  const arrayBuffer = buffer.slice(byteOffset, byteOffset + byteLength);
+  const metadata = parquetMetadata(arrayBuffer);
+  const schema = parquetSchema(metadata);
+  return schema.children.map((col) => col.element.name);
+}
+
+/** Helpers */
+
+async function fetchRange(s3Client, bucket, key, range) {
+  const response = await s3Client.send(new GetObjectCommand({ Bucket: bucket, Key: key, Range: range }));
+  const chunks = [];
+  for await (const chunk of response.Body) {
+    chunks.push(chunk);
+  }
+  return Buffer.concat(chunks);
+}