feat(plugins): add BucketResolverPlugin for multi-bucket dataset queries (#51)

* feat(plugins): add BucketResolverPlugin for multi-bucket dataset queries

Resolves a missing bucket/endpoint per file by matching parsed file paths
against configured dataset prefixes. Runs after QueryParserPlugin so each
file in a query is resolved independently, letting a single query join
across datasets that live in different buckets. Opt-in via the plugins
array, consistent with FSPurgePlugin/StatsPlugin/AvroPlugin.

* fix(s3): avoid day-of-month collision in evictTodayFromListingCache test

The "past day" cache fixture was hardcoded to day=01, colliding with
today's own day-level prefix whenever the suite runs on the 1st of the
month (UTC) and making the "past day entries should remain" assertion
fail. Derive it relative to the actual current day instead.

* feat(plugins): resolve list_files bucket via BucketResolverPlugin

list_files talks to S3 directly and never runs through the query pipeline
that BucketResolverPlugin's processQuery hooks into, so it couldn't
auto-resolve a bucket from datasets. Add a resolveListFiles hook to the
plugin (sharing its prefix-matching logic with processQuery) and have
ListFilesTool reduce over config.plugins to call it — kept local to the
MCP tool rather than lifecycle.js, since that module is scoped to the
s3quoia() query/download pipeline which list_files never touches.

Adding BucketResolverPlugin to config.plugins now covers both the query
and list_files tools with no extra configuration.

Author: grommett

README.md

@@ -233,6 +233,27 @@ const results = await s3quoia({
 });
 ```
 
+### BucketResolverPlugin
+
+When a query references files by relative path (no `{bucket:...}` token) and no `defaultBucket` is set, `BucketResolverPlugin` resolves each file's bucket and endpoint from a list of dataset configs (see [Dataset options](#dataset-options)), matching the file path against each dataset's `prefix`. Resolution runs per file, so a single query can join across datasets in different buckets. Files that don't match any prefix are left unresolved — the existing "no bucket configured" error still applies to them.
+
+```js
+import s3quoia, { BucketResolverPlugin } from 's3quoia';
+
+const datasets = [
+  { name: 'sales', prefix: 'sales/', bucket: 'sales-bucket' },
+  { name: 'logs', prefix: 'logs/raw/', bucket: 'logs-bucket', endpoint: 'https://s3.logs.example.com' },
+];
+
+const results = await s3quoia({
+  // ...
+  plugins: [new BucketResolverPlugin(datasets)],
+  query: `SELECT * FROM read_parquet('sales/year={yyyy}/data.parquet')`,
+});
+```
+
+When used with `S3QuoiaMCP` (see [Plugins](#plugins-1)), the same `BucketResolverPlugin` instance also resolves the bucket for the `list_files` tool — it calls a second hook, `resolveListFiles`, since `list_files` talks to S3 directly and doesn't go through the query pipeline that `processQuery` hooks into.
+
 ## MCP Server
 
 s3quoia ships a [Model Context Protocol](https://modelcontextprotocol.io/) server that exposes three tools to any MCP-compatible client (Claude Desktop, Claude Code, IBM Bob etc.):
@@ -405,22 +426,25 @@ new S3QuoiaMCP({
 
 #### Plugins
 
-Pass a `plugins` array to enable `FSPurgePlugin`, `StatsPlugin`, or any custom plugin for every query the server handles:
+Pass a `plugins` array to enable `FSPurgePlugin`, `StatsPlugin`, `BucketResolverPlugin`, or any custom plugin for every query the server handles:
 
 ```js
 import { S3QuoiaMCP } from 's3quoia/mcp';
-import { FSPurgePlugin, StatsPlugin } from 's3quoia';
+import { FSPurgePlugin, StatsPlugin, BucketResolverPlugin } from 's3quoia';
+
+const datasets = [ /* ... */ ];
 
 new S3QuoiaMCP({
-  datasets: [ /* ... */ ],
+  datasets,
   plugins: [
     new FSPurgePlugin({ bucketsDir: '/tmp/s3quoia', lastAccessTTLMinutes: 120 }),
     new StatsPlugin((event) => console.error(event)),
+    new BucketResolverPlugin(datasets),
   ],
 }).start();
 ```
 
-The built-in server (`npx s3quoia`) runs `FSPurgePlugin` and `StatsPlugin` by default. When extending with `S3QuoiaMCP`, plugins are opt-in.
+The built-in server (`npx s3quoia`) runs `FSPurgePlugin` and `StatsPlugin` by default. When extending with `S3QuoiaMCP`, plugins are opt-in — add `BucketResolverPlugin` explicitly if you want both `query` and `list_files` to auto-resolve a bucket per file from your `datasets` config. `list_files` doesn't run through the query pipeline, so the plugin resolves it via a separate `resolveListFiles` hook rather than `processQuery` — no extra config needed, adding the plugin once covers both tools.
 
 ### Adding custom tools
 

e2e/list-files.e2e.js

@@ -0,0 +1,64 @@
+import { describe, it } from 'node:test';
+import assert from 'node:assert';
+
+// ListFilesTool reads S3 credentials/endpoint from process.env at module load
+// time, so these must be set before it's imported. S3_BUCKET is deliberately
+// left unset — bucket resolution for these tests comes entirely from
+// BucketResolverPlugin via the datasets config.
+process.env.S3_ACCESS_KEY_ID = 'test-access-key';
+process.env.S3_SECRET_ACCESS_KEY = 'test-secret-key';
+process.env.S3_ENDPOINT = 'http://localhost:9000';
+delete process.env.S3_BUCKET;
+
+const { default: ListFilesTool } = await import('../src/mcp/tools/list-files/list-files.js');
+const { BucketResolverPlugin } = await import('../src/s3quoia.js');
+
+const BUCKET = 'test-bucket';
+const BUCKET_2 = 'test-bucket-2';
+
+describe('list_files e2e', () => {
+  it('resolves the bucket via BucketResolverPlugin when no bucket param or S3_BUCKET is set', async () => {
+    const datasets = [{ name: 'reports', prefix: 'reports/', bucket: BUCKET }];
+    const tool = new ListFilesTool({ plugins: [new BucketResolverPlugin(datasets)] });
+
+    const result = await tool.handler({ prefix: 'reports/' });
+    const parsed = JSON.parse(result.content[0].text);
+
+    assert.ok(parsed.files.some((file) => file.file === 'reports/summary.parquet'));
+    assert.ok(parsed.files.some((file) => file.file === 'reports/summary.csv'));
+  });
+
+  it('resolves different prefixes to different buckets in the same server process', async () => {
+    const datasets = [
+      { name: 'reports', prefix: 'reports/', bucket: BUCKET },
+      { name: 'reference', prefix: 'reference', bucket: BUCKET_2 },
+    ];
+    const tool = new ListFilesTool({ plugins: [new BucketResolverPlugin(datasets)] });
+
+    const reportsResult = await tool.handler({ prefix: 'reports/' });
+    const referenceResult = await tool.handler({ prefix: 'reference' });
+
+    const reportsParsed = JSON.parse(reportsResult.content[0].text);
+    const referenceParsed = JSON.parse(referenceResult.content[0].text);
+
+    assert.ok(reportsParsed.files.some((file) => file.file === 'reports/summary.parquet'));
+    assert.ok(referenceParsed.files.some((file) => file.file === 'reference.parquet'));
+  });
+
+  it('rejects when no dataset prefix matches and no bucket/S3_BUCKET is set — no silent fallback', async () => {
+    const datasets = [{ name: 'reports', prefix: 'reports/', bucket: BUCKET }];
+    const tool = new ListFilesTool({ plugins: [new BucketResolverPlugin(datasets)] });
+
+    await assert.rejects(() => tool.handler({ prefix: 'unrelated/' }));
+  });
+
+  it('prefers an explicit bucket param over plugin resolution', async () => {
+    const datasets = [{ name: 'reports', prefix: 'reports/', bucket: 'wrong-bucket' }];
+    const tool = new ListFilesTool({ plugins: [new BucketResolverPlugin(datasets)] });
+
+    const result = await tool.handler({ prefix: 'reference.parquet', bucket: BUCKET_2 });
+    const parsed = JSON.parse(result.content[0].text);
+
+    assert.ok(parsed.files.some((file) => file.file === 'reference.parquet'));
+  });
+});

e2e/s3quoia.e2e.js

@@ -4,7 +4,7 @@ import { mkdtemp, rm, writeFile } from 'node:fs/promises';
 import { tmpdir } from 'node:os';
 import { join } from 'node:path';
 
-import s3quoia, { StatsPlugin, AvroPlugin } from '../src/s3quoia.js';
+import s3quoia, { StatsPlugin, AvroPlugin, BucketResolverPlugin } from '../src/s3quoia.js';
 
 const ENDPOINT = 'http://localhost:9000';
 const BUCKET = 'test-bucket';
@@ -433,4 +433,63 @@ describe('s3quoia e2e', () => {
     assert.strictEqual(result[0].description, 'user login event');
     assert.strictEqual(result[2].description, 'purchase completed');
   });
+
+  it('resolves a bucket from datasets via BucketResolverPlugin when no defaultBucket is set', async () => {
+    const datasetDir = await mkdtemp(join(tmpdir(), 's3-e2e-bucket-resolver-'));
+    const datasets = [{ name: 'reports', prefix: 'reports/', bucket: BUCKET }];
+
+    try {
+      const result = await s3quoia({
+        accessKeyId: ACCESS_KEY,
+        secretAccessKey: SECRET_KEY,
+        defaultEndpoint: ENDPOINT,
+        bucketsDir: datasetDir,
+        plugins: [new BucketResolverPlugin(datasets)],
+        from: FROM,
+        to: TO,
+        query: `SELECT * FROM read_parquet('reports/summary.parquet') ORDER BY id`,
+        format: 'jsonRecords',
+      });
+
+      assert.strictEqual(result.length, 3);
+      assert.deepStrictEqual(result[0], { id: 1, event_type: 'login', region: 'us-east', value: 42.5 });
+    } finally {
+      await rm(datasetDir, { recursive: true });
+    }
+  });
+
+  it('joins files from two different buckets resolved independently via BucketResolverPlugin, with no bucket tokens in the query', async () => {
+    const datasetDir = await mkdtemp(join(tmpdir(), 's3-e2e-bucket-resolver-join-'));
+    const datasets = [
+      { name: 'reports', prefix: 'reports/', bucket: BUCKET },
+      { name: 'reference', prefix: 'reference', bucket: BUCKET_2 },
+    ];
+
+    try {
+      const result = await s3quoia({
+        accessKeyId: ACCESS_KEY,
+        secretAccessKey: SECRET_KEY,
+        defaultEndpoint: ENDPOINT,
+        bucketsDir: datasetDir,
+        plugins: [new BucketResolverPlugin(datasets)],
+        from: FROM,
+        to: TO,
+        query: `
+          SELECT s.id, s.event_type, r.description
+          FROM read_parquet('reports/summary.parquet') s
+          JOIN read_parquet('reference.parquet') r ON s.id = r.id
+          ORDER BY s.id
+        `,
+        format: 'jsonRecords',
+      });
+
+      assert.strictEqual(result.length, 3);
+      assert.strictEqual(result[0].id, 1);
+      assert.strictEqual(result[0].event_type, 'login');
+      assert.strictEqual(result[0].description, 'user login event');
+      assert.strictEqual(result[2].description, 'purchase completed');
+    } finally {
+      await rm(datasetDir, { recursive: true });
+    }
+  });
 });

src/mcp/tools/list-files/list-files.js

@@ -41,8 +41,13 @@ export default class ListFilesTool extends BaseTool {
   }
 
   async handler({ prefix = '', maxResults = 100, endpoint, bucket }) {
-    const resolvedEndpoint = endpoint || S3_ENDPOINT;
-    const resolvedBucket = bucket || S3_BUCKET;
+    const resolved = resolveBucket(this.config.plugins, {
+      prefix,
+      bucket: bucket || S3_BUCKET,
+      endpoint: endpoint || S3_ENDPOINT,
+    });
+    const resolvedEndpoint = resolved.endpoint;
+    const resolvedBucket = resolved.bucket;
     const s3Client = buildS3Client({
       apiKey: S3_API_KEY,
       accessKeyId: S3_ACCESS_KEY_ID,
@@ -70,6 +75,10 @@ export default class ListFilesTool extends BaseTool {
 
 /** Helpers */
 
+function resolveBucket(plugins = [], context) {
+  return plugins.reduce((result, plugin) => plugin.resolveListFiles?.(result) ?? result, context);
+}
+
 function getRepresentativeFiles(files) {
   const parquetFiles = files.filter(({ file }) => file.endsWith('.parquet'));
   const dirMap = parquetFiles.reduce(addFirstFilePerDir, new Map());

src/mcp/tools/list-files/list-files.test.js

@@ -1,8 +1,51 @@
 import { describe, it } from 'node:test';
 import assert from 'node:assert';
 import esmock from 'esmock';
+import BucketResolverPlugin from '../../../plugins/bucket-resolver/bucket-resolver-plugin.js';
 
 describe('handleListFiles', () => {
+  it('resolves the bucket from a BucketResolverPlugin in config.plugins when no bucket/S3_BUCKET is set', async () => {
+    const sentParams = [];
+    const datasets = [{ name: 'sales', prefix: 'sales/', bucket: 'sales-bucket' }];
+    const tool = await getMockedTool({
+      s3Response: { Contents: [], CommonPrefixes: [], IsTruncated: false },
+      config: { plugins: [new BucketResolverPlugin(datasets)] },
+      sentParams,
+    });
+
+    await tool.handler({ prefix: 'sales/year=2024/' });
+
+    assert.strictEqual(sentParams[0].Bucket, 'sales-bucket');
+  });
+
+  it('prefers an explicit bucket param over plugin resolution', async () => {
+    const sentParams = [];
+    const datasets = [{ name: 'sales', prefix: 'sales/', bucket: 'plugin-bucket' }];
+    const tool = await getMockedTool({
+      s3Response: { Contents: [], CommonPrefixes: [], IsTruncated: false },
+      config: { plugins: [new BucketResolverPlugin(datasets)] },
+      sentParams,
+    });
+
+    await tool.handler({ prefix: 'sales/', bucket: 'explicit-bucket' });
+
+    assert.strictEqual(sentParams[0].Bucket, 'explicit-bucket');
+  });
+
+  it('leaves the bucket undefined when no plugin matches the prefix, with no fallback', async () => {
+    const sentParams = [];
+    const datasets = [{ name: 'sales', prefix: 'sales/', bucket: 'sales-bucket' }];
+    const tool = await getMockedTool({
+      s3Response: { Contents: [], CommonPrefixes: [], IsTruncated: false },
+      config: { plugins: [new BucketResolverPlugin(datasets)] },
+      sentParams,
+    });
+
+    await tool.handler({ prefix: 'unrelated/' });
+
+    assert.strictEqual(sentParams[0].Bucket, undefined);
+  });
+
   it('returns directories, files, and truncated flag from S3', async () => {
     const tool = await getMockedTool({
       s3Response: {
@@ -85,8 +128,13 @@ describe('handleListFiles', () => {
   });
 });
 
-async function getMockedTool({ s3Response, columns = [] }) {
-  const mockS3Client = { send: () => Promise.resolve(s3Response) };
+async function getMockedTool({ s3Response, columns = [], config = {}, sentParams = [] }) {
+  const mockS3Client = {
+    send: (command) => {
+      sentParams.push(command);
+      return Promise.resolve(s3Response);
+    },
+  };
   const { default: ListFilesTool } = await esmock('./list-files.js', {
     '@aws-sdk/client-s3': {
       ListObjectsV2Command: class ListObjectsV2Command {
@@ -99,5 +147,5 @@ async function getMockedTool({ s3Response, columns = [] }) {
     '../../../utils/parquet-schema-reader.js': { readParquetColumns: () => Promise.resolve(columns) },
     '../../../utils/bigint-replacer.js': { bigintReplacer: (_, val) => val },
   });
-  return new ListFilesTool({});
+  return new ListFilesTool(config);
 }

src/plugins/bucket-resolver/bucket-resolver-plugin.js

@@ -0,0 +1,51 @@
+/**
+ * Fills in a missing `bucket`/`endpoint` by matching a file path against
+ * configured dataset prefixes. Exposes two hooks over the same matching
+ * logic:
+ *
+ * - `processQuery` — runs after `QueryParserPlugin`, so `context.settings`
+ *   already holds one parsed entry per file reference (e.g.
+ *   `sales/year={yyyy}/data.parquet`). Matching is done per file rather than
+ *   once for the whole query, so a join across two datasets resolves each
+ *   file to its own bucket independently.
+ * - `resolveListFiles` — called directly by the `list_files` MCP tool, which
+ *   talks to S3 without going through the query pipeline, so it resolves its
+ *   own `{ prefix, bucket, endpoint }` context the same way.
+ *
+ * A path with no matching dataset prefix is left unchanged in both cases —
+ * no fallback to the first configured dataset — so the existing "no bucket"
+ * error surfaces normally instead of silently querying the wrong bucket.
+ */
+export default class BucketResolverPlugin {
+  name = 'BucketResolverPlugin';
+
+  constructor(datasets = []) {
+    this.datasets = datasets;
+  }
+
+  processQuery(context) {
+    const settings = context.settings.map((setting) => this.resolveSetting(setting));
+    return { ...context, settings };
+  }
+
+  resolveListFiles(context) {
+    return this.mergeDataset(context, context.prefix);
+  }
+
+  resolveSetting(setting) {
+    return this.mergeDataset(setting, setting.file);
+  }
+
+  mergeDataset(context, path) {
+    if (context.bucket) return context;
+    const dataset = this.findDataset(path);
+    if (!dataset) return context;
+    return { ...context, bucket: dataset.bucket, endpoint: context.endpoint ?? dataset.endpoint };
+  }
+
+  findDataset(path) {
+    const candidates = this.datasets.filter(({ prefix }) => prefix && path?.startsWith(prefix));
+    if (!candidates.length) return undefined;
+    return candidates.reduce((best, dataset) => (dataset.prefix.length > best.prefix.length ? dataset : best));
+  }
+}