Merge pull request #165 from import-ai/feat/trace

feat(trace): add trace module

Author: ycdzj

package.json

@@ -77,6 +77,7 @@
     "cookie-parser": "^1.4.7",
     "express": "^5.1.0",
     "handlebars": "^4.7.8",
+    "kafkajs": "^2.2.4",
     "lodash": "^4.17.21",
     "meilisearch": "^0.50.0",
     "minio": "^8.0.5",

pnpm-lock.yaml

@@ -47,6 +47,7 @@ import { UpdateAttachmentUrls1755499552000 } from 'omniboxd/migrations/175549955
 import { ScanResourceAttachments1755504936756 } from 'omniboxd/migrations/1755504936756-scan-resource-attachments';
 import { SharesAllResources1754471311959 } from 'omniboxd/migrations/1754471311959-shares-all-resources';
 import { ResourcesModule } from 'omniboxd/resources/resources.module';
+import { TraceModule } from 'omniboxd/trace/trace.module';
 
 @Module({})
 export class AppModule implements NestModule {
@@ -93,6 +94,7 @@ export class AppModule implements NestModule {
         AttachmentsModule,
         SharesModule,
         SharedResourcesModule,
+        TraceModule,
         // CacheModule.registerAsync({
         //   imports: [ConfigModule],
         //   inject: [ConfigService],

src/app/app.module.ts

@@ -0,0 +1,10 @@
+import { IsNotEmpty, IsString, IsObject } from 'class-validator';
+
+export class TraceEventDto {
+  @IsString()
+  @IsNotEmpty()
+  name: string;
+
+  @IsObject()
+  props: Record<string, any>;
+}

src/trace/dto/trace-event.dto.ts

@@ -0,0 +1,29 @@
+import { Expose } from 'class-transformer';
+import { TraceEventDto } from './trace-event.dto';
+
+export class TraceMessageDto {
+  @Expose()
+  timestamp: number;
+
+  @Expose({ name: 'event_name' })
+  eventName: string;
+
+  @Expose({ name: 'event_props' })
+  eventProps?: string;
+
+  @Expose({ name: 'user_id' })
+  userId?: string;
+
+  @Expose({ name: 'user_agent' })
+  userAgent?: string;
+
+  static fromEvent(event: TraceEventDto, userId?: string, userAgent?: string) {
+    const message = new TraceMessageDto();
+    message.timestamp = Date.now();
+    message.eventName = event.name;
+    message.eventProps = JSON.stringify(event.props);
+    message.userId = userId;
+    message.userAgent = userAgent;
+    return message;
+  }
+}

src/trace/dto/trace-message.dto.ts

@@ -0,0 +1,10 @@
+import { IsArray, ValidateNested } from 'class-validator';
+import { Type } from 'class-transformer';
+import { TraceEventDto } from 'omniboxd/trace/dto/trace-event.dto';
+
+export class TraceReqDto {
+  @IsArray()
+  @ValidateNested({ each: true })
+  @Type(() => TraceEventDto)
+  events: TraceEventDto[];
+}

src/trace/dto/trace-req.dto.ts

@@ -0,0 +1,21 @@
+import { Body, Controller, Post, Req } from '@nestjs/common';
+import { TraceService } from 'omniboxd/trace/trace.service';
+import { TraceReqDto } from './dto/trace-req.dto';
+import { CookieAuth } from 'omniboxd/auth/decorators';
+import { UserId } from 'omniboxd/decorators/user-id.decorator';
+
+@Controller('api/v1/trace')
+export class TraceController {
+  constructor(private traceService: TraceService) {}
+
+  @Post()
+  @CookieAuth({ onAuthFail: 'continue' })
+  async trace(
+    @Req() req: Request,
+    @Body() traceReq: TraceReqDto,
+    @UserId({ optional: true }) userId?: string,
+  ): Promise<void> {
+    const userAgent = req.headers['user-agent'];
+    await this.traceService.emitTraceEvents(traceReq.events, userId, userAgent);
+  }
+}

src/trace/trace.controller.ts

@@ -0,0 +1,12 @@
+import { Module } from '@nestjs/common';
+import { ConfigModule } from '@nestjs/config';
+import { TraceController } from 'omniboxd/trace/trace.controller';
+import { TraceService } from 'omniboxd/trace/trace.service';
+
+@Module({
+  imports: [ConfigModule],
+  controllers: [TraceController],
+  providers: [TraceService],
+  exports: [TraceService],
+})
+export class TraceModule {}

src/trace/trace.module.ts

@@ -0,0 +1,69 @@
+import {
+  Injectable,
+  Logger,
+  OnModuleDestroy,
+  OnModuleInit,
+} from '@nestjs/common';
+import { ConfigService } from '@nestjs/config';
+import { Kafka, Producer } from 'kafkajs';
+import { TraceEventDto } from './dto/trace-event.dto';
+import { instanceToPlain } from 'class-transformer';
+import { TraceMessageDto } from './dto/trace-message.dto';
+
+@Injectable()
+export class TraceService implements OnModuleInit, OnModuleDestroy {
+  private readonly logger = new Logger(TraceService.name);
+  private readonly topic?: string;
+  private readonly kafka?: Kafka;
+  private readonly producer?: Producer;
+
+  constructor(private readonly configService: ConfigService) {
+    const brokerUrl = this.configService.get<string>('OBB_KAFKA_BROKER');
+    const topic = this.configService.get<string>('OBB_KAFKA_TOPIC');
+    const clientId = this.configService.get<string>('OBB_KAFKA_CLIENT_ID');
+    if (!brokerUrl || !topic || !clientId) {
+      return;
+    }
+    const brokers = brokerUrl.split(',');
+    this.topic = topic;
+    this.kafka = new Kafka({
+      clientId,
+      brokers,
+    });
+    this.producer = this.kafka.producer();
+  }
+
+  async onModuleInit() {
+    if (this.producer) {
+      await this.producer.connect();
+      this.logger.log('Kafka producer connected successfully');
+    }
+  }
+
+  async onModuleDestroy() {
+    if (this.producer) {
+      await this.producer.disconnect();
+      this.logger.log('Kafka producer disconnected');
+    }
+  }
+
+  async emitTraceEvents(
+    events: TraceEventDto[],
+    userId?: string,
+    userAgent?: string,
+  ): Promise<void> {
+    if (!this.producer || !this.topic) {
+      return;
+    }
+    const messages = events.map((event) => {
+      const dto = TraceMessageDto.fromEvent(event, userId, userAgent);
+      return {
+        value: JSON.stringify(instanceToPlain(dto)),
+      };
+    });
+    await this.producer.send({
+      topic: this.topic,
+      messages,
+    });
+  }
+}