feat: implement session management with SessionService and SessionPlugin; refactor query execution to handle session cookies; update token and query handling for improved session tracking

src/query/base/base-query.controller.ts

@@ -5,10 +5,11 @@ import {
   Inject,
   Param,
   Post,
+  Req,
   Res,
   UseGuards,
 } from "@nestjs/common";
-import { Response } from "express";
+import { Response, Request } from "express";
 import { QueryHandlerService } from "../handler/query.handler.service";
 import { ApiTokenGuard } from "src/api/guards/api-token.guard";
 import { QueryExecuterService } from "../executer/query.executer.service";
@@ -52,11 +53,22 @@ export abstract class BaseQueryController {
     const queryResult = await this.queryExecuterService.runQueryQueued(
       id,
       query,
-      headers
+      headers,
+      headers.cookie.split("; ").reduce((acc, cookie) => {
+        const [key, value] = cookie.split("=");
+        acc[key] = value;
+        return acc;
+      }, {})
     );
 
     res.status(queryResult?.statusCode || 200);
 
+    if (queryResult?.cookies) {
+      for (const [key, value] of Object.entries(queryResult?.cookies || {})) {
+        res.cookie(key, value);
+      }
+    }
+
     if (
       queryResult?.redirect ||
       (queryResult?.statusCode === 302 &&

src/query/executer/query.executer.service.ts

@@ -15,6 +15,7 @@ import { InjectQueue } from "@nestjs/bullmq";
 import { QUEUE_NAMES } from "src/queue/constants";
 import { Queue, QueueEvents } from "bullmq";
 import { FunctionService } from "src/query/function/function.service";
+import { SessionService } from "../session/session.service";
 
 @Injectable()
 export class QueryExecuterService {
@@ -27,6 +28,7 @@ export class QueryExecuterService {
     private readonly functionService: FunctionService,
     readonly databaseManagerService: DatabaseManagerService,
     readonly redisNodeService: RedisNodeService,
+    readonly sessionService: SessionService,
     @InjectQueue(QUEUE_NAMES.QUERY) private queryQueue: Queue
   ) {
     this.queueEvents = new QueueEvents(this.queryQueue.name);
@@ -54,7 +56,8 @@ export class QueryExecuterService {
   async runQueryQueued(
     token: string,
     queryData: any,
-    headers: Record<string, any> = {}
+    headers: Record<string, any> = {},
+    cookies: Record<string, any> = {}
   ): Promise<QueryResponse> {
     const job = await this.queryQueue.add(
       `${new Date().getTime()}_${token}`,
@@ -62,6 +65,7 @@ export class QueryExecuterService {
         token,
         queryData,
         headers,
+        cookies,
       },
       {
         removeOnComplete: true,
@@ -77,7 +81,8 @@ export class QueryExecuterService {
   async runQuery(
     token: string,
     queryData: any,
-    headers: Record<string, any> = {}
+    headers: Record<string, any> = {},
+    cookies: Record<string, any> = {}
   ): Promise<QueryResponse> {
     const query = await this.queryRepository.findOne({
       where: { id: token },
@@ -88,7 +93,16 @@ export class QueryExecuterService {
       throw new Error("Query not found");
     }
 
-    const vm = await this.createVm(query);
+    const sessionId = cookies["x-session-id"] || null;
+
+    console.log("Session ID:", sessionId);
+
+    if (!sessionId) {
+      const session = await this.sessionService.create(query.project.id);
+      cookies["x-session-id"] = session.sessionId;
+    }
+
+    const vm = await this.createVm(query, cookies["x-session-id"]);
     const result = await vm.runScript(
       this.clearImports(query.source),
       queryData,
@@ -99,10 +113,20 @@ export class QueryExecuterService {
       throw new Error(`Error initializing VM: ${JSON.stringify(result)}`);
     }
 
+    if (!result?.cookies || !result?.cookies["x-session-id"]) {
+      if (result.cookies === undefined) {
+        result.cookies = {};
+      }
+
+      result.cookies["x-session-id"] = (
+        await this.sessionService.get(cookies["x-session-id"], query.project.id)
+      ).sessionId;
+    }
+
     return result;
   }
 
-  private async createVm(query: Query) {
+  private async createVm(query: Query, sessionId: string = null) {
     const imports = this.parseImports(query.source);
     const importsParsed = imports.map((imp) => {
       const item = imp.split("/");
@@ -140,7 +164,9 @@ export class QueryExecuterService {
         throw new Error(`Plugin ${plugin.name} not found`);
       }
 
-      plugins.push(await registeredPlugins[plugin.name](this, query));
+      plugins.push(
+        await registeredPlugins[plugin.name](this, query, sessionId)
+      );
     }
 
     const vm = new Vm({

src/query/query.module.ts

@@ -14,6 +14,7 @@ import { FunctionService } from "src/query/function/function.service";
 import { FunctionController } from "src/query/function/function.controller";
 import { RedisManagerModule } from "src/redisManager/redisManager.module";
 import { RedisModule } from "src/redis/redis.module";
+import { SessionService } from "./session/session.service";
 
 @Module({
   imports: [
@@ -26,7 +27,12 @@ import { RedisModule } from "src/redis/redis.module";
     TypeOrmModule.forFeature([Query, FunctionEntity]),
   ],
   controllers: [QueryController, CommandController, FunctionController],
-  providers: [QueryExecuterService, QueryHandlerService, FunctionService],
+  providers: [
+    QueryExecuterService,
+    SessionService,
+    QueryHandlerService,
+    FunctionService,
+  ],
   exports: [QueryExecuterService, TypeOrmModule, QueryHandlerService],
 })
 export class QueryModule {}

src/query/session/session.service.ts

@@ -0,0 +1,50 @@
+import { Inject, Injectable } from "@nestjs/common";
+import { RedisClient } from "src/redis/redis.service";
+
+@Injectable()
+export class SessionService {
+  constructor(@Inject(RedisClient) private readonly redisClient: RedisClient) {}
+
+  private generateSessionId(length: number = 16): string {
+    const characters =
+      "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
+    let result = "";
+    const charactersLength = characters.length;
+    for (let i = 0; i < length; i++) {
+      result += characters.charAt(Math.floor(Math.random() * charactersLength));
+    }
+    return `${result}_${new Date().getTime()}`;
+  }
+
+  async create(prefix: string): Promise<{ sessionId: string }> {
+    const sessionId = this.generateSessionId();
+    await this.set(sessionId, prefix, {
+      sessionId,
+    });
+
+    return { sessionId };
+  }
+
+  async get(sessionId: string | null, prefix: string): Promise<any> {
+    if (!sessionId) {
+      return await this.create(prefix);
+    }
+
+    const data = await this.redisClient.get(`${prefix}:${sessionId}`);
+
+    if (data) {
+      await this.redisClient.set(`${prefix}:${sessionId}`, data, 3600);
+      return data;
+    }
+
+    return await this.create(prefix);
+  }
+
+  async set(sessionId: string, prefix: string, data: any): Promise<void> {
+    await this.redisClient.set(`${prefix}:${sessionId}`, data, 3600);
+  }
+
+  async delete(sessionId: string, prefix: string): Promise<void> {
+    await this.redisClient.del(`${prefix}:${sessionId}`);
+  }
+}