neptune-lambda-client 2.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/README.md

@@ -6,6 +6,15 @@ reestablish a connection to the database if the web socket connection closes and
 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 `query` function that accepts a single
@@ -21,6 +30,27 @@ async function getNode(id) {
 }
 ```
 
+### 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),

package/index.js

@@ -1,5 +1,6 @@
 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;
@@ -35,7 +36,7 @@ function createHeaders(host, port, path, options) {
     }).headers;
 }
 
-export function create(host, port, {useIam = true, protocol = 'wss'} = {}) {
+export function create(host, port, {useIam = true, protocol = 'wss', partition} = {}) {
     let conn = null;
     let g = null;
 
@@ -74,7 +75,8 @@ export function create(host, port, {useIam = true, protocol = 'wss'} = {}) {
     };
 
     const createGraphTraversalSource = conn => {
-        return traversal().withRemote(conn);
+        const g = traversal().withRemote(conn);
+        return partition ? g.withStrategies(new PartitionStrategy(partition)) : g;
     };
 
     return {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "neptune-lambda-client",
-  "version": "2.0.0",
+  "version": "3.0.0",
   "description": "Gremlin client to robustly query AWS Neptune from AWS Lambda",
   "type": "module",
   "exports": {
@@ -29,10 +29,14 @@
   "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/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);
+    });
+});