neptune-lambda-client 1.0.0 → 3.0.0

package/.claude/settings.local.json

@@ -0,0 +1,8 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(npm test *)",
+      "Bash(TESTCONTAINERS_RYUK_DISABLED=true npm test)"
+    ]
+  }
+}

package/.github/workflows/ci.yml

@@ -0,0 +1,24 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+env:
+  FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        node: [22, 24]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-node@v4
+        with:
+          node-version: ${{ matrix.node }}
+          cache: npm
+      - run: npm ci
+      - run: npm test

package/README.md

@@ -2,24 +2,61 @@
 
 ## Overview
 A very simple Gremlin client to robustly query AWS Neptune from AWS Lambda. The client will automatically
-reestablish a connection to the database if the web socket connection closes and will also automatically 
-retry (5 times) when it encounters `ConcurrentModificationException` and `ReadOnlyViolationException` errors.
+reestablish a connection to the database if the web socket connection closes and will also automatically
+retry (up to 5 times, with exponential backoff and jitter) when it encounters `ConcurrentModificationException`
+and `ReadOnlyViolationException` errors.
+
+## Installation
+
+`gremlin` is a peer dependency — install it alongside the client so your application controls the
+exact version that talks to your Neptune cluster:
+
+```sh
+npm install neptune-lambda-client gremlin
+```
 
 ## Usage
 
-This client is instantiated with a factory function and exposes a single function called `query`. `query` accepts 
-a single argument, which is a function that use the Gremlin `g` object.
+This client is instantiated with a factory function and exposes a `query` function that accepts a single
+argument: a function that uses the Gremlin `g` object. It also exposes a `close` function for graceful shutdown.
 
 ```js
-const gremlinClient = require('neptune-lambda-client');
+import { create } from 'neptune-lambda-client';
 
-const {query} = gremlinClient.create({
-    host: 'neptune-db-url',
-    port: '8182',
-    useIam: true
-});
+const { query } = create('neptune-db-url', '8182', { useIam: true });
 
 async function getNode(id) {
     return query(async g => g.V(id).next().then(x => x.value));
 }
 ```
+
+### Partition strategy
+
+To scope every query through the client to a Gremlin
+[`PartitionStrategy`](https://tinkerpop.apache.org/docs/current/reference/#partitionstrategy)
+— useful for multi-tenant graphs and per-suite test isolation — pass a `partition`
+option. Its fields are forwarded directly to `PartitionStrategy`:
+
+```js
+const { query } = create('neptune-db-url', '8182', {
+    useIam: true,
+    partition: {
+        partitionKey: '_partition',
+        writePartition: 'tenant-a',
+        readPartitions: ['tenant-a', 'shared'],
+    },
+});
+```
+
+The strategy is reapplied automatically when the client reconnects after a dropped
+WebSocket, so partition isolation holds across the full Lambda lifetime.
+
+## Known limitations
+
+Per the [AWS Neptune Lambda guidance](https://docs.aws.amazon.com/neptune/latest/userguide/lambda-functions-examples.html#lambda-functions-examples-javascript),
+if the underlying WebSocket is closed after the driver sends a request but before the response arrives,
+the query resolves with `undefined` rather than throwing. Because this state cannot be turned into an
+exception on the request/response path, we instead throw from the socket's `close` handler when the close
+code is `1006` (abnormal closure) — this surfaces as an unhandled exception that fails the Lambda
+invocation, so the client invoking the Lambda can retry. We recommend implementing retry logic on the
+caller side as well as using the built-in retry in this library.

package/index.js

@@ -1,10 +1,13 @@
-const aws4 = require('aws4');
-const gremlin = require('gremlin');
-const retry = require('async-retry');
+import aws4 from 'aws4';
+import gremlin from 'gremlin';
+import { PartitionStrategy } from 'gremlin/lib/process/traversal-strategy.js';
+import retry from 'async-retry';
 
 const traversal = gremlin.process.AnonymousTraversalSource.traversal;
 const {driver: {DriverRemoteConnection}} = gremlin;
 
+const RETRYABLE_ERRORS = ['ConcurrentModificationException', 'ReadOnlyViolationException'];
+
 function createHeaders(host, port, path, options) {
     if (!host || !port) {
         throw new Error('Host and port are required');
@@ -33,75 +36,87 @@ function createHeaders(host, port, path, options) {
     }).headers;
 }
 
-module.exports = {
-    create: function(host, port, {useIam = true} = {useIam: true}) {
-        let conn = null;
-        let g = null;
-
-        const path = "/gremlin"
-        const url = `wss://${host}:${port}${path}`
+export function create(host, port, {useIam = true, protocol = 'wss', partition} = {}) {
+    let conn = null;
+    let g = null;
 
-        const createRemoteConnection = () => {
+    const path = "/gremlin"
+    const url = `${protocol}://${host}:${port}${path}`
 
-            const c = new DriverRemoteConnection(
-                url,
-                {
-                    mimeType: 'application/vnd.gremlin-v2.0+json',
-                    pingEnabled: false,
-                    headers: useIam ? createHeaders(host, port, path, {}) : {}
-                });
+    const createRemoteConnection = () => {
 
-            c._client._connection.on('log', message => {
-                console.info(`connection message - ${message}`);
+        const c = new DriverRemoteConnection(
+            url,
+            {
+                // Lambda freezes between invocations; heartbeats fired just before freeze
+                // time out after thaw, causing noisy disconnects. Our reconnect-on-error path
+                // handles dead connections, so skip the liveness pings.
+                pingEnabled: false,
+                headers: useIam ? createHeaders(host, port, path, {}) : {}
             });
 
-            c._client._connection.on('close', (code, message) => {
-                console.info(`close - ${code} ${message}`);
-                if (code == 1006){
-                    console.error('Connection closed prematurely');
-                    throw new Error('Connection closed prematurely');
-                }
-            });
+        // gremlin opens the WebSocket eagerly on construction; swallow a potential
+        // orphan rejection here. The real error resurfaces on the next submit().
+        c.open().catch(() => {});
+
+        c._client._connection.on('log', message => {
+            console.info(`connection message - ${message}`);
+        });
 
-            return c;
-        };
-
-        const createGraphTraversalSource = conn => {
-            return traversal().withRemote(conn);
-        };
-
-        return {
-            query: async f => {
-                if (conn == null){
-                    console.info('Initializing connection')
-                    conn = createRemoteConnection();
-                    g = createGraphTraversalSource(conn);
-                }
-
-                return retry(async (bail, count) => {
-                    return f(g).catch(err => {
-                        if(count > 0) console.log('Retry attempt no: ' + count);
-                        if (err.message.startsWith('WebSocket is not open')){
-                            console.warn('Reopening connection');
-                            conn.close();
-                            conn = createRemoteConnection();
-                            g = createGraphTraversalSource(conn);
-                            throw err;
-                        } else if (err.message.includes('ConcurrentModificationException')){
-                            console.warn('Retrying query because of ConcurrentModificationException');
-                            throw err;
-                        } else if (err.message.includes('ReadOnlyViolationException')){
-                            console.warn('Retrying query because of ReadOnlyViolationException');
-                            throw err;
-                        } else {
-                            console.warn('Unrecoverable error: ' + err);
-                            return bail(err);
-                        }
-                    })
-                }, {
-                    factor: 1,
-                    retries: 5
-                });
+        c._client._connection.on('close', (code, message) => {
+            console.info(`close - ${code} ${message}`);
+            if (code == 1006){
+                console.error('Connection closed prematurely');
+                throw new Error('Connection closed prematurely');
+            }
+        });
+
+        return c;
+    };
+
+    const createGraphTraversalSource = conn => {
+        const g = traversal().withRemote(conn);
+        return partition ? g.withStrategies(new PartitionStrategy(partition)) : g;
+    };
+
+    return {
+        query: async f => {
+            if (conn == null){
+                console.info('Initializing connection')
+                conn = createRemoteConnection();
+                g = createGraphTraversalSource(conn);
+            }
+
+            return retry(async (bail, count) => {
+                return f(g).catch(err => {
+                    if(count > 0) console.log('Retry attempt no: ' + count);
+                    if (err.message.startsWith('WebSocket is not open')){
+                        console.warn('Reopening connection');
+                        conn.close();
+                        conn = createRemoteConnection();
+                        g = createGraphTraversalSource(conn);
+                        throw err;
+                    }
+                    if (RETRYABLE_ERRORS.some(name => err.message.includes(name))) {
+                        console.warn('Retrying query: ' + err.message);
+                        throw err;
+                    }
+                    console.warn('Unrecoverable error: ' + err);
+                    return bail(err);
+                })
+            }, {
+                retries: 5,
+                factor: 2,
+                minTimeout: 1000,
+                maxTimeout: 10000,
+                randomize: true
+            });
+        },
+        close: async () => {
+            if (conn != null) {
+                await conn.close();
+                conn = null;
+                g = null;
             }
         }
     }

package/package.json

@@ -1,10 +1,14 @@
 {
   "name": "neptune-lambda-client",
-  "version": "1.0.0",
+  "version": "3.0.0",
   "description": "Gremlin client to robustly query AWS Neptune from AWS Lambda",
-  "main": "index.js",
+  "type": "module",
+  "exports": {
+    ".": "./index.js"
+  },
   "scripts": {
-    "test": "mocha"
+    "test": "vitest run",
+    "test:watch": "vitest"
   },
   "repository": {
     "type": "git",
@@ -25,7 +29,15 @@
   "homepage": "https://github.com/svozza/neptune-lambda-client#readme",
   "dependencies": {
     "async-retry": "1.3.3",
-    "aws4": "1.11.0",
-    "gremlin": "3.5.2"
+    "aws4": "1.11.0"
+  },
+  "peerDependencies": {
+    "gremlin": "^3.5.0"
+  },
+  "devDependencies": {
+    "gremlin": "3.5.2",
+    "testcontainers": "^12.0.0",
+    "vitest": "^2.1.9",
+    "ws": "^8.18.0"
   }
 }

package/test/helpers/gremlin-container.js

@@ -0,0 +1,14 @@
+import { GenericContainer, Wait } from 'testcontainers';
+
+export async function startGremlinServer() {
+    const container = await new GenericContainer('tinkerpop/gremlin-server:3.7.3')
+        .withExposedPorts(8182)
+        .withWaitStrategy(Wait.forLogMessage(/Channel started at port 8182/))
+        .start();
+
+    return {
+        host: container.getHost(),
+        port: container.getMappedPort(8182),
+        stop: () => container.stop()
+    };
+}

package/test/helpers/gremlin-server.js

@@ -0,0 +1,32 @@
+import { WebSocketServer } from 'ws';
+
+export async function startFakeGremlin() {
+    const upgrades = [];
+    let resolveNextUpgrade;
+    let nextUpgrade = new Promise(resolve => { resolveNextUpgrade = resolve; });
+
+    const wss = new WebSocketServer({ port: 0, path: '/gremlin' });
+
+    wss.on('connection', (_socket, request) => {
+        const record = { headers: { ...request.headers }, url: request.url };
+        upgrades.push(record);
+        const resolve = resolveNextUpgrade;
+        nextUpgrade = new Promise(r => { resolveNextUpgrade = r; });
+        resolve(record);
+    });
+
+    await new Promise(resolve => wss.on('listening', resolve));
+    const { port } = wss.address();
+
+    return {
+        port,
+        upgrades,
+        waitForNextUpgrade: () => nextUpgrade,
+        close: () =>
+            new Promise(resolve => {
+                for (const client of wss.clients) client.close(1000, 'shutdown');
+                wss.close(() => resolve());
+            })
+    };
+}
+

package/test/index.test.js

@@ -0,0 +1,147 @@
+import { describe, it, expect, beforeEach, afterEach, onTestFinished, vi } from 'vitest';
+import { create } from '../index.js';
+import { startFakeGremlin } from './helpers/gremlin-server.js';
+
+let server;
+
+beforeEach(async () => {
+    server = await startFakeGremlin();
+});
+
+afterEach(async () => {
+    await server.close();
+});
+
+function createClient(opts = {}) {
+    const client = create('localhost', server.port, { useIam: false, protocol: 'ws', ...opts });
+    onTestFinished(() => client.close());
+    return client;
+}
+
+describe('create() / query()', () => {
+    it('resolves with the user function result on the happy path', async () => {
+        const client = createClient();
+        const result = await client.query(async () => 42);
+        expect(result).toBe(42);
+    });
+
+    describe('with faked retry timers', () => {
+        beforeEach(() => {
+            vi.useFakeTimers({ toFake: ['setTimeout'] });
+        });
+        afterEach(() => {
+            vi.useRealTimers();
+        });
+
+        it('retries on ConcurrentModificationException and eventually succeeds', async () => {
+            const client = createClient();
+            let calls = 0;
+            const p = client.query(async () => {
+                calls++;
+                if (calls < 3) throw new Error('ConcurrentModificationException');
+                return 'ok';
+            });
+            await vi.runAllTimersAsync();
+            expect(await p).toBe('ok');
+            expect(calls).toBe(3);
+        });
+
+        it('retries on ReadOnlyViolationException and eventually succeeds', async () => {
+            const client = createClient();
+            let calls = 0;
+            const p = client.query(async () => {
+                calls++;
+                if (calls < 2) throw new Error('ReadOnlyViolationException: stale primary');
+                return 'ok';
+            });
+            await vi.runAllTimersAsync();
+            expect(await p).toBe('ok');
+            expect(calls).toBe(2);
+        });
+
+        it('reconnects when the error is "WebSocket is not open"', async () => {
+            const client = createClient();
+            let calls = 0;
+            const p = client.query(async () => {
+                calls++;
+                if (calls === 1) throw new Error('WebSocket is not open: readyState 3 (CLOSED)');
+                return 'reconnected';
+            });
+            await vi.runAllTimersAsync();
+            expect(await p).toBe('reconnected');
+            expect(calls).toBe(2);
+
+            // The new WebSocket handshake runs on real I/O; wait for it before asserting.
+            vi.useRealTimers();
+            await vi.waitFor(() => expect(server.upgrades.length).toBeGreaterThanOrEqual(2), { timeout: 2000 });
+        });
+
+        it('gives up after the retry cap', async () => {
+            const client = createClient();
+            let calls = 0;
+            const p = client.query(async () => {
+                calls++;
+                throw new Error('ConcurrentModificationException always');
+            });
+            const assertion = expect(p).rejects.toThrow('ConcurrentModificationException');
+            await vi.runAllTimersAsync();
+            await assertion;
+            expect(calls).toBe(6);
+        });
+    });
+
+    it('does not retry on unrecoverable errors', async () => {
+        const client = createClient();
+        let calls = 0;
+        await expect(
+            client.query(async () => {
+                calls++;
+                throw new Error('some random failure');
+            })
+        ).rejects.toThrow('some random failure');
+        expect(calls).toBe(1);
+    });
+});
+
+describe('IAM signing', () => {
+    afterEach(() => {
+        vi.unstubAllEnvs();
+    });
+
+    it('sends a SigV4-signed Authorization header when useIam is true', async () => {
+        vi.stubEnv('AWS_ACCESS_KEY_ID', 'AKIATEST');
+        vi.stubEnv('AWS_SECRET_ACCESS_KEY', 'secret');
+        vi.stubEnv('AWS_REGION', 'eu-west-1');
+        vi.stubEnv('AWS_SESSION_TOKEN', undefined);
+
+        const client = createClient({ useIam: true });
+        const upgradePromise = server.waitForNextUpgrade();
+        await client.query(async () => 'done');
+        const upgrade = await upgradePromise;
+
+        expect(upgrade.headers.authorization).toMatch(/^AWS4-HMAC-SHA256/);
+        expect(upgrade.headers['x-amz-date']).toBeDefined();
+    });
+
+    it('sends no Authorization header when useIam is false', async () => {
+        const client = createClient();
+        const upgradePromise = server.waitForNextUpgrade();
+        await client.query(async () => 'done');
+        const upgrade = await upgradePromise;
+
+        expect(upgrade.headers.authorization).toBeUndefined();
+        expect(upgrade.headers['x-amz-date']).toBeUndefined();
+    });
+
+    it('rejects when useIam is true but AWS env vars are missing', async () => {
+        vi.stubEnv('AWS_ACCESS_KEY_ID', undefined);
+        vi.stubEnv('AWS_SECRET_ACCESS_KEY', undefined);
+        vi.stubEnv('AWS_REGION', undefined);
+        vi.stubEnv('AWS_SESSION_TOKEN', undefined);
+
+        const client = createClient({ useIam: true });
+        await expect(client.query(async () => 'done')).rejects.toThrow(
+            'AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY and AWS_REGION are required'
+        );
+    });
+});

package/test/integration.test.js

@@ -0,0 +1,155 @@
+import { beforeAll, afterAll, describe, it, expect, onTestFinished } from 'vitest';
+import { randomUUID } from 'node:crypto';
+import { create } from '../index.js';
+import { startGremlinServer } from './helpers/gremlin-container.js';
+
+let server;
+
+beforeAll(async () => {
+    server = await startGremlinServer();
+}, 120_000);
+
+afterAll(async () => {
+    await server?.stop();
+});
+
+const KEY = '_partition';
+
+function mkClient(write, reads) {
+    const client = create(server.host, String(server.port), {
+        useIam: false,
+        protocol: 'ws',
+        partition: { partitionKey: KEY, writePartition: write, readPartitions: reads }
+    });
+    onTestFinished(() => client.close());
+    return client;
+}
+
+describe('partition strategy', () => {
+    it('isolates writes to writePartition and gates reads on readPartitions', async () => {
+        const a = `a-${randomUUID()}`;
+        const b = `b-${randomUUID()}`;
+        const name = `iso-${randomUUID()}`;
+
+        const writerA = mkClient(a, [a]);
+        const readerA = mkClient(a, [a]);
+        const readerB = mkClient(b, [b]);
+
+        await writerA.query(g => g.addV('Widget').property('name', name).next());
+
+        const fromB = await readerB.query(g => g.V().has('Widget', 'name', name).toList());
+        expect(fromB).toEqual([]);
+
+        const fromA = await readerA.query(g => g.V().has('Widget', 'name', name).toList());
+        expect(fromA.length).toBe(1);
+    });
+
+    it('lets readPartitions span multiple partitions', async () => {
+        const a = `a-${randomUUID()}`;
+        const shared = `shared-${randomUUID()}`;
+        const nameA = `multi-a-${randomUUID()}`;
+        const nameShared = `multi-s-${randomUUID()}`;
+
+        const writerA = mkClient(a, [a]);
+        const writerShared = mkClient(shared, [shared]);
+        const readerBoth = mkClient(a, [a, shared]);
+
+        await writerA.query(g => g.addV('Widget').property('name', nameA).next());
+        await writerShared.query(g => g.addV('Widget').property('name', nameShared).next());
+
+        const fromA = await readerBoth.query(g => g.V().has('Widget', 'name', nameA).toList());
+        const fromShared = await readerBoth.query(g => g.V().has('Widget', 'name', nameShared).toList());
+        expect(fromA.length).toBe(1);
+        expect(fromShared.length).toBe(1);
+    });
+
+    it('preserves the partition strategy across a reconnect', async () => {
+        const a = `a-${randomUUID()}`;
+        const b = `b-${randomUUID()}`;
+        const name = `reconnect-${randomUUID()}`;
+
+        const writerA = mkClient(a, [a]);
+        const readerA = mkClient(a, [a]);
+        const readerB = mkClient(b, [b]);
+
+        let calls = 0;
+        await writerA.query(g => {
+            calls++;
+            if (calls === 1) {
+                // Triggers index.js's reconnect path: close + rebuild g via
+                // createGraphTraversalSource. If the strategy is dropped on
+                // rebuild, the write below leaks into other partitions.
+                throw new Error('WebSocket is not open: readyState 3 (CLOSED)');
+            }
+            return g.addV('Widget').property('name', name).next();
+        });
+        expect(calls).toBe(2);
+
+        const fromB = await readerB.query(g => g.V().has('Widget', 'name', name).toList());
+        expect(fromB).toEqual([]);
+
+        const fromA = await readerA.query(g => g.V().has('Widget', 'name', name).toList());
+        expect(fromA.length).toBe(1);
+    });
+});
+
+describe('end-to-end query', () => {
+    function plainClient() {
+        const client = create(server.host, String(server.port), { useIam: false, protocol: 'ws' });
+        onTestFinished(() => client.close());
+        return client;
+    }
+
+    it('round-trips a vertex through bytecode + GraphSON', async () => {
+        const client = plainClient();
+        const name = `e2e-${randomUUID()}`;
+
+        await client.query(g => g.addV('Widget').property('name', name).property('count', 7).next());
+
+        const fetched = await client.query(g =>
+            g.V().has('Widget', 'name', name).valueMap(true).toList()
+        );
+
+        expect(fetched.length).toBe(1);
+        const v = fetched[0];
+        expect(v.get('name')).toEqual([name]);
+        expect(v.get('count')).toEqual([7]);
+    });
+});
+
+describe('connection lifecycle', () => {
+    it('reuses one connection across multiple queries, then rebuilds after close()', async () => {
+        const client = create(server.host, String(server.port), { useIam: false, protocol: 'ws' });
+        onTestFinished(() => client.close());
+
+        for (let i = 0; i < 3; i++) {
+            const out = await client.query(g => g.inject(i).next());
+            expect(out.value).toBe(i);
+        }
+
+        await client.close();
+        await client.close();
+
+        const out = await client.query(g => g.inject('after-reopen').next());
+        expect(out.value).toBe('after-reopen');
+    });
+});
+
+describe('server-side errors', () => {
+    it('bails immediately on a malformed traversal (no retry)', async () => {
+        const client = create(server.host, String(server.port), { useIam: false, protocol: 'ws' });
+        onTestFinished(() => client.close());
+
+        let calls = 0;
+        await expect(
+            client.query(g => {
+                calls++;
+                // addE() without from/to context is invalid at execution time
+                // on the server; the resulting error doesn't match any
+                // RETRYABLE_ERRORS entry, so we expect a single attempt.
+                return g.addE('orphan').next();
+            })
+        ).rejects.toThrow();
+        expect(calls).toBe(1);
+    });
+});