From 929d762aee4f438794eb112e9ec2d368070ca318 Mon Sep 17 00:00:00 2001 From: Dave New Date: Thu, 8 Aug 2024 13:19:57 +0200 Subject: [PATCH] fix: database smart timeout chore: only destroy if neon fix: flexi db conn suspend timeout fix: if env doesnt exist --- packages/functions-runtime/package.json | 3 +- packages/functions-runtime/pnpm-lock.yaml | 20 +------ packages/functions-runtime/src/database.js | 61 ++++++++++++++++++++-- 3 files changed, 58 insertions(+), 26 deletions(-) diff --git a/packages/functions-runtime/package.json b/packages/functions-runtime/package.json index 703f89427..f655720d9 100644 --- a/packages/functions-runtime/package.json +++ b/packages/functions-runtime/package.json @@ -21,7 +21,7 @@ "dependencies": { "@aws-sdk/client-s3": "^3.617.0", "@aws-sdk/credential-providers": "^3.617.0", - "@neondatabase/serverless": "^0.9.3", + "@neondatabase/serverless": "^0.9.4", "@opentelemetry/api": "^1.7.0", "@opentelemetry/exporter-trace-otlp-proto": "^0.46.0", "@opentelemetry/resources": "^1.19.0", @@ -31,7 +31,6 @@ "json-rpc-2.0": "^1.7.0", "ksuid": "^3.0.0", "kysely": "^0.25.0", - "kysely-neon": "^1.3.0", "pg": "^8.11.3", "traceparent": "^1.0.0", "ws": "^8.17.1" diff --git a/packages/functions-runtime/pnpm-lock.yaml b/packages/functions-runtime/pnpm-lock.yaml index 8f2f5ef6b..a32bd9751 100644 --- a/packages/functions-runtime/pnpm-lock.yaml +++ b/packages/functions-runtime/pnpm-lock.yaml @@ -12,7 +12,7 @@ dependencies: specifier: ^3.617.0 version: 3.621.0(@aws-sdk/client-sso-oidc@3.621.0) '@neondatabase/serverless': - specifier: ^0.9.3 + specifier: ^0.9.4 version: 0.9.4 '@opentelemetry/api': specifier: ^1.7.0 @@ -41,9 +41,6 @@ dependencies: kysely: specifier: ^0.25.0 version: 0.25.0 - kysely-neon: - specifier: ^1.3.0 - version: 1.3.0(@neondatabase/serverless@0.9.4)(kysely@0.25.0)(ws@8.18.0) pg: specifier: ^8.11.3 version: 8.12.0 @@ -2124,21 +2121,6 @@ packages: base-convert-int-array: 1.0.1 dev: false - /kysely-neon@1.3.0(@neondatabase/serverless@0.9.4)(kysely@0.25.0)(ws@8.18.0): - resolution: {integrity: sha512-CIIlbmqpIXVJDdBEYtEOwbmALag0jmqYrGfBeM4cHKb9AgBGs+X1SvXUZ8TqkDacQEqEZN2XtsDoUkcMIISjHw==} - peerDependencies: - '@neondatabase/serverless': ^0.4.3 - kysely: 0.x.x - ws: ^8.13.0 - peerDependenciesMeta: - ws: - optional: true - dependencies: - '@neondatabase/serverless': 0.9.4 - kysely: 0.25.0 - ws: 8.18.0 - dev: false - /kysely@0.25.0: resolution: {integrity: sha512-srn0efIMu5IoEBk0tBmtGnoUss4uwvxtbFQWG/U2MosfqIace1l43IFP1PmEpHRDp+Z79xIcKEqmHH3dAvQdQA==} engines: {node: '>=14.0.0'} diff --git a/packages/functions-runtime/src/database.js b/packages/functions-runtime/src/database.js index eb251e36c..f1761dcc4 100644 --- a/packages/functions-runtime/src/database.js +++ b/packages/functions-runtime/src/database.js @@ -1,10 +1,10 @@ const { Kysely, PostgresDialect, CamelCasePlugin } = require("kysely"); +const neonserverless = require("@neondatabase/serverless"); const { AsyncLocalStorage } = require("async_hooks"); const { AuditContextPlugin } = require("./auditing"); const pg = require("pg"); const { PROTO_ACTION_TYPES } = require("./consts"); const { withSpan } = require("./tracing"); -const { NeonDialect } = require("kysely-neon"); const ws = require("ws"); // withDatabase is responsible for setting the correct database client in our AsyncLocalStorage @@ -107,7 +107,20 @@ function getDatabaseClient() { class InstrumentedPool extends pg.Pool { async connect(...args) { const _super = super.connect.bind(this); - return withSpan("Database Connect", function () { + return withSpan("Database Connect", function (span) { + span.setAttribute("dialect", process.env["KEEL_DB_CONN_TYPE"]); + span.setAttribute("timeout", connectionTimeout()); + return _super(...args); + }); + } +} + +class InstrumentedNeonServerlessPool extends neonserverless.Pool { + async connect(...args) { + const _super = super.connect.bind(this); + return withSpan("Database Connect", function (span) { + span.setAttribute("dialect", process.env["KEEL_DB_CONN_TYPE"]); + span.setAttribute("timeout", connectionTimeout()); return _super(...args); }); } @@ -134,6 +147,7 @@ class InstrumentedClient extends pg.Client { return withSpan(spanName, function (span) { if (sqlAttribute) { span.setAttribute("sql", args[0]); + span.setAttribute("dialect", process.env["KEEL_DB_CONN_TYPE"]); } return _super(...args); }); @@ -165,14 +179,42 @@ function getDialect() { // Although I doubt we will run into these freeze/thaw issues if idleTimeoutMillis is always shorter than the // time is takes for a lambda to freeze (which is not a constant, but could be as short as several minutes, // https://www.pluralsight.com/resources/blog/cloud/how-long-does-aws-lambda-keep-your-idle-functions-around-before-a-cold-start) - idleTimeoutMillis: 120000, + idleTimeoutMillis: connectionTimeout(), connectionString: mustEnv("KEEL_DB_CONN"), }), }); case "neon": - return new NeonDialect({ + neonserverless.neonConfig.webSocketConstructor = ws; + + const pool = new InstrumentedNeonServerlessPool({ + idleTimeoutMillis: connectionTimeout(), connectionString: mustEnv("KEEL_DB_CONN"), - webSocketConstructor: ws, + }); + + pool.on("connect", (client) => { + const originalQuery = client.query; + client.query = function (...args) { + const sql = args[0]; + + let sqlAttribute = false; + let spanName = txStatements[sql.toLowerCase()]; + if (!spanName) { + spanName = "Database Query"; + sqlAttribute = true; + } + + return withSpan(spanName, function (span) { + if (sqlAttribute) { + span.setAttribute("sql", args[0]); + span.setAttribute("dialect", process.env["KEEL_DB_CONN_TYPE"]); + } + return originalQuery.apply(client, args); + }); + }; + }); + + return new PostgresDialect({ + pool: pool, }); default: throw Error("unexpected KEEL_DB_CONN_TYPE: " + dbConnType); @@ -187,6 +229,15 @@ function mustEnv(key) { return v; } + +function connectionTimeout() { + const v = Number(process.env["KEEL_DB_CONN_TIMEOUT"]); + if (!v || isNaN(v)) { + return 45000; // 60s is our current neon suspend default + } + return v; +} + // initialise the database client at module scope level so the db variable is set getDatabaseClient();