@teamkeel/functions-runtime 0.393.1 → 0.393.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@teamkeel/functions-runtime",
-  "version": "0.393.1",
+  "version": "0.393.3",
   "description": "Internal package used by @teamkeel/sdk",
   "main": "src/index.js",
   "scripts": {
@@ -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"

package/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
@@ -43,10 +43,11 @@ async function withDatabase(db, actionType, cb) {
   });
 }
 
-let db = null;
-
 const dbInstance = new AsyncLocalStorage();
 
+// used to establish a singleton for our vitest environment
+let vitestDb = null;
+
 // useDatabase will retrieve the database client set by withDatabase from the local storage
 function useDatabase() {
   // retrieve the instance of the database client from the store which is aware of
@@ -61,7 +62,10 @@ function useDatabase() {
   // which covers any test files ending in *.test.ts. Custom function code runs in a different node process which will not have this environment variable. Tests written using our testing
   // framework call actions (and in turn custom function code) over http using the ActionExecutor class
   if ("NODE_ENV" in process.env && process.env.NODE_ENV == "test") {
-    return getDatabaseClient();
+    if (!vitestDb) {
+      vitestDb = createDatabaseClient();
+    }
+    return vitestDb;
   }
 
   // If we've gotten to this point, then we know that we are in a custom function runtime server
@@ -69,17 +73,11 @@ function useDatabase() {
   throw new Error("useDatabase must be called within a function");
 }
 
-// getDatabaseClient will return a brand new instance of Kysely. Every instance of Kysely
+// createDatabaseClient will return a brand new instance of Kysely. Every instance of Kysely
 // represents an individual connection to the database.
 // not to be exported externally from our sdk - consumers should use useDatabase
-function getDatabaseClient() {
-  // 'db' represents the singleton connection to the database which is stored
-  // as a module scope variable.
-  if (db) {
-    return db;
-  }
-
-  db = new Kysely({
+function createDatabaseClient() {
+  const db = new Kysely({
     dialect: getDialect(),
     plugins: [
       // ensures that the audit context data is written to Postgres configuration parameters
@@ -107,7 +105,18 @@ 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"]);
+      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"]);
       return _super(...args);
     });
   }
@@ -134,6 +143,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 +175,41 @@ 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: 50000,
           connectionString: mustEnv("KEEL_DB_CONN"),
         }),
       });
     case "neon":
-      return new NeonDialect({
+      neonserverless.neonConfig.webSocketConstructor = ws;
+
+      const pool = new InstrumentedNeonServerlessPool({
         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,9 +224,6 @@ function mustEnv(key) {
   return v;
 }
 
-// initialise the database client at module scope level so the db variable is set
-getDatabaseClient();
-
-module.exports.getDatabaseClient = getDatabaseClient;
+module.exports.createDatabaseClient = createDatabaseClient;
 module.exports.useDatabase = useDatabase;
 module.exports.withDatabase = withDatabase;

package/src/handleJob.js

@@ -3,7 +3,7 @@ const {
   createJSONRPCSuccessResponse,
   JSONRPCErrorCode,
 } = require("json-rpc-2.0");
-const { getDatabaseClient } = require("./database");
+const { createDatabaseClient } = require("./database");
 const { errorToJSONRPCResponse, RuntimeErrors } = require("./errors");
 const opentelemetry = require("@opentelemetry/api");
 const { withSpan } = require("./tracing");
@@ -25,6 +25,8 @@ async function handleJob(request, config) {
   return opentelemetry.context.with(activeContext, () => {
     // Wrapping span for the whole request
     return withSpan(request.method, async (span) => {
+      let db = null;
+
       try {
         const { createJobContextAPI, jobs } = config;
 
@@ -51,7 +53,7 @@ async function handleJob(request, config) {
             ? true
             : null;
 
-        const db = getDatabaseClient();
+        db = createDatabaseClient();
         const jobFunction = jobs[request.method];
         const actionType = PROTO_ACTION_TYPES.JOB;
 
@@ -89,6 +91,10 @@ async function handleJob(request, config) {
           RuntimeErrors.UnknownError,
           message
         );
+      } finally {
+        if (db) {
+          await db.destroy();
+        }
       }
     });
   });

package/src/handleRequest.js

@@ -3,7 +3,7 @@ const {
   createJSONRPCSuccessResponse,
   JSONRPCErrorCode,
 } = require("json-rpc-2.0");
-const { getDatabaseClient } = require("./database");
+const { createDatabaseClient } = require("./database");
 const { tryExecuteFunction } = require("./tryExecuteFunction");
 const { errorToJSONRPCResponse, RuntimeErrors } = require("./errors");
 const opentelemetry = require("@opentelemetry/api");
@@ -28,6 +28,8 @@ async function handleRequest(request, config) {
   return opentelemetry.context.with(activeContext, () => {
     // Wrapping span for the whole request
     return withSpan(request.method, async (span) => {
+      let db = null;
+
       try {
         const { createContextAPI, functions, permissionFns, actionTypes } =
           config;
@@ -62,7 +64,7 @@ async function handleRequest(request, config) {
             ? true
             : null;
 
-        const db = getDatabaseClient();
+        db = createDatabaseClient();
         const customFunction = functions[request.method];
         const actionType = actionTypes[request.method];
 
@@ -122,6 +124,10 @@ async function handleRequest(request, config) {
           RuntimeErrors.UnknownError,
           message
         );
+      } finally {
+        if (db) {
+          await db.destroy();
+        }
       }
     });
   });

package/src/handleSubscriber.js

@@ -3,7 +3,7 @@ const {
   createJSONRPCSuccessResponse,
   JSONRPCErrorCode,
 } = require("json-rpc-2.0");
-const { getDatabaseClient } = require("./database");
+const { createDatabaseClient } = require("./database");
 const { errorToJSONRPCResponse, RuntimeErrors } = require("./errors");
 const opentelemetry = require("@opentelemetry/api");
 const { withSpan } = require("./tracing");
@@ -24,6 +24,8 @@ async function handleSubscriber(request, config) {
   return opentelemetry.context.with(activeContext, () => {
     // Wrapping span for the whole request
     return withSpan(request.method, async (span) => {
+      let db = null;
+
       try {
         const { createSubscriberContextAPI, subscribers } = config;
 
@@ -45,7 +47,7 @@ async function handleSubscriber(request, config) {
           meta: request.meta,
         });
 
-        const db = getDatabaseClient();
+        db = createDatabaseClient();
         const subscriberFunction = subscribers[request.method];
         const actionType = PROTO_ACTION_TYPES.SUBSCRIBER;
 
@@ -77,6 +79,10 @@ async function handleSubscriber(request, config) {
           RuntimeErrors.UnknownError,
           message
         );
+      } finally {
+        if (db) {
+          await db.destroy();
+        }
       }
     });
   });