rotor: mongodb tweaks

jitsucom · Dec 16, 2023 · 40598bd · 40598bd · vercel · Dec 16, 2023
1 parent c2cb1b1
commit 40598bd
Show file tree

Hide file tree

Showing 8 changed files with 112 additions and 18 deletions.
diff --git a/libs/core-functions/package.json b/libs/core-functions/package.json
@@ -40,6 +40,7 @@
     "juava": "workspace:*",
     "lodash": "^4.17.21",
     "mongodb": "^5.9.2",
+    "@mongodb-js/zstd": "^1.2.0",
     "ioredis": "^5.3.2",
     "posthog-node": "^2.4.0",
     "tslib": "^2.4.0",

diff --git a/libs/core-functions/src/functions/lib/mongodb.ts b/libs/core-functions/src/functions/lib/mongodb.ts
@@ -17,7 +17,7 @@ async function createClient() {
   log.atInfo().log(`Connecting to MongoDB server...`);
 
   // Create a new MongoClient
-  const client = new MongoClient(mongodbURL);
+  const client = new MongoClient(mongodbURL, { compressors: ["zstd"] });
   // Connect the client to the server (optional starting in v4.7)
   await client.connect();
   // Establish and verify connection

diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/services/rotor/dist_package.json b/services/rotor/dist_package.json
@@ -3,6 +3,7 @@
   "version": "0.0.0",
   "dependencies": {
     "isolated-vm": "4.6.0",
-    "@sensejs/kafkajs-zstd-support": "^0.11.0"
+    "@sensejs/kafkajs-zstd-support": "^0.11.0",
+    "@mongodb-js/zstd": "^1.2.0"
   }
 }
diff --git a/services/rotor/src/http/udf.ts b/services/rotor/src/http/udf.ts
@@ -1,12 +1,22 @@
-import { UDFTestRun, UDFTestRequest, mongodb, defaultTTL, createMongoStore } from "@jitsu/core-functions";
+import {
+  UDFTestRun,
+  UDFTestRequest,
+  mongodb,
+  defaultTTL,
+  createMongoStore,
+  createTtlStore,
+} from "@jitsu/core-functions";
 import { getLog } from "juava";
+import { redis } from "@jitsu-internal/console/lib/server/redis";
 
 const log = getLog("udf_run");
 
 export const UDFRunHandler = async (req, res) => {
   const body = req.body as UDFTestRequest;
-  log.atInfo().log(`Running function: ${body?.functionId} workspace: ${body?.workspaceId}`, JSON.stringify(body));
-  body.store = createMongoStore(body?.workspaceId, await mongodb.waitInit(), defaultTTL);
+  //log.atInfo().log(`Running function: ${body?.functionId} workspace: ${body?.workspaceId}`, JSON.stringify(body));
+  body.store = process.env.MONGODB_URL
+    ? createMongoStore(body?.workspaceId, mongodb(), defaultTTL)
+    : createTtlStore(body?.workspaceId, redis(), defaultTTL);
   const result = await UDFTestRun(body);
   if (result.error) {
     log

diff --git a/services/rotor/src/index.ts b/services/rotor/src/index.ts
@@ -6,7 +6,7 @@ import {
   rotorConsumerGroupId,
 } from "./lib/kafka-config";
 import { kafkaRotor } from "./lib/rotor";
-
+import { mongodb } from "@jitsu/core-functions";
 import minimist from "minimist";
 import { glob } from "glob";
 import fs from "fs";
@@ -16,6 +16,7 @@ import Prometheus from "prom-client";
 import { FunctionsHandler, FunctionsHandlerMulti } from "./http/functions";
 import { initMaxMindClient } from "./lib/maxmind";
 import { rotorMessageHandler } from "./lib/message-handler";
+import { redis } from "@jitsu-internal/console/lib/server/redis";
 
 export const log = getLog("rotor");
 
@@ -54,6 +55,8 @@ async function main() {
     });
     log.atInfo().log("Starting kafka processing");
     Prometheus.collectDefaultMetrics();
+    await mongodb.waitInit();
+    await redis.waitInit();
     rotor
       .start()
       .then(chMetrics => {

diff --git a/services/rotor/src/lib/message-handler.ts b/services/rotor/src/lib/message-handler.ts
@@ -108,8 +108,12 @@ export async function rotorMessageHandler(
     },
   };
   const oldStore = createTtlStore(connection.workspaceId, redis(), defaultTTL);
-  const newStore = createMongoStore(connection.workspaceId, await mongodb.waitInit(), defaultTTL);
-  const store = newStoreWorskpaceId.includes(connection.workspaceId) ? newStore : createMultiStore(newStore, oldStore);
+  const newStore = process.env.MONGODB_URL ? createMongoStore(connection.workspaceId, mongodb(), defaultTTL) : oldStore;
+  const store = newStoreWorskpaceId.includes(connection.workspaceId)
+    ? newStore
+    : newStore != oldStore
+    ? createMultiStore(newStore, oldStore)
+    : oldStore;
 
   const rl = await redisLogger.waitInit();
 

diff --git a/services/rotor/webpack.config.js b/services/rotor/webpack.config.js
@@ -6,6 +6,8 @@ const config = {
   target: "node",
   externals: {
     "isolated-vm": "require('isolated-vm')",
+    "@sensejs/kafkajs-zstd-support": "require('@sensejs/kafkajs-zstd-support')",
+    "@mongodb-js/zstd": "require('@mongodb-js/zstd')",
   },
   node: {
     __dirname: false,