neptune-lambda-client 1.0.0 → 2.0.0

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,31 @@
 
 ## 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.
 
 ## 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));
 }
 ```
+
+## 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,12 @@
-const aws4 = require('aws4');
-const gremlin = require('gremlin');
-const retry = require('async-retry');
+import aws4 from 'aws4';
+import gremlin from 'gremlin';
+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 +35,86 @@ 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'} = {}) {
+    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 => {
+        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;
+                    }
+                    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": "2.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",
@@ -27,5 +31,9 @@
     "async-retry": "1.3.3",
     "aws4": "1.11.0",
     "gremlin": "3.5.2"
+  },
+  "devDependencies": {
+    "vitest": "^2.1.9",
+    "ws": "^8.18.0"
   }
 }

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'
+        );
+    });
+});