Class TraceRepository

backend/src/trace/trace.repository.ts:27–206

@Injectable()
export class TraceRepository implements OnModuleDestroy {
  private pool: Pool;

  constructor(
    @Inject(DRIZZLE_TOKEN)
    private readonly db: NodePgDatabase,
  ) {
    // Create a separate pool for LISTEN/NOTIFY to avoid conflicts
    this.pool = new Pool({
      connectionString: process.env.DATABASE_URL,
    });
  }

  async onModuleDestroy() {
    await this.pool.end();
  }

  /**
   * Subscribe to real-time trace events for a specific run ID using Postgres LISTEN/NOTIFY
   */
  async subscribeToRun(
    runId: string,
    callback: (payload: string) => void,
  ): Promise<() => Promise<void>> {
    const client = await this.pool.connect();
    const channel = `trace_events_${runId}`;

    try {
      await client.query(`LISTEN "${channel}"`);

      client.on('notification', (msg) => {
        if (msg.channel === channel && msg.payload) {
          callback(msg.payload);
        }
      });

      // Return unsubscribe function
      return async () => {
        try {
          await client.query(`UNLISTEN "${channel}"`);
        } finally {
          client.release();
        }
      };
    } catch (error) {
      client.release();
      throw error;
    }
  }

  /**
   * Notify subscribers of new trace events
   */
  async notifyRun(runId: string, payload: string): Promise<void> {
    const channel = `trace_events_${runId}`;
    await this.pool.query('SELECT pg_notify($1, $2)', [channel, payload]);
  }