mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-07-19 19:16:29 +08:00
feat: add queue management for admin panel
This commit is contained in:
@@ -40,6 +40,7 @@ import { MailModule } from './core/mail';
|
||||
import { MonitorModule } from './core/monitor';
|
||||
import { NotificationModule } from './core/notification';
|
||||
import { PermissionModule } from './core/permission';
|
||||
import { QueueDashboardModule } from './core/queue-dashboard';
|
||||
import { QuotaModule } from './core/quota';
|
||||
import { SelfhostModule } from './core/selfhost';
|
||||
import { StorageModule } from './core/storage';
|
||||
@@ -189,7 +190,8 @@ export function buildAppModule(env: Env) {
|
||||
OAuthModule,
|
||||
CustomerIoModule,
|
||||
CommentModule,
|
||||
AccessTokenModule
|
||||
AccessTokenModule,
|
||||
QueueDashboardModule
|
||||
)
|
||||
// doc service only
|
||||
.useIf(() => env.flavors.doc, DocServiceModule)
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
import { getQueueToken } from '@nestjs/bullmq';
|
||||
import { Injectable, Logger, Module, OnModuleInit } from '@nestjs/common';
|
||||
import { HttpAdapterHost, ModuleRef } from '@nestjs/core';
|
||||
import { createQueueDashExpressMiddleware } from '@queuedash/api';
|
||||
import type { Queue as BullMQQueue } from 'bullmq';
|
||||
import type { Application, NextFunction, Request, Response } from 'express';
|
||||
|
||||
import { Config } from '../../base/config';
|
||||
import { QUEUES } from '../../base/job/queue/def';
|
||||
import { AuthGuard, AuthModule } from '../auth';
|
||||
import { FeatureModule, FeatureService } from '../features';
|
||||
|
||||
type QueueDashQueue = {
|
||||
queue: BullMQQueue;
|
||||
displayName: string;
|
||||
type: 'bullmq';
|
||||
};
|
||||
|
||||
@Injectable()
|
||||
class QueueDashboardService implements OnModuleInit {
|
||||
private readonly logger = new Logger(QueueDashboardService.name);
|
||||
|
||||
constructor(
|
||||
private readonly adapterHost: HttpAdapterHost,
|
||||
private readonly config: Config,
|
||||
private readonly feature: FeatureService,
|
||||
private readonly authGuard: AuthGuard,
|
||||
private readonly moduleRef: ModuleRef
|
||||
) {}
|
||||
|
||||
async onModuleInit() {
|
||||
const httpAdapter = this.adapterHost.httpAdapter;
|
||||
if (!httpAdapter) {
|
||||
return;
|
||||
}
|
||||
|
||||
const app = httpAdapter.getInstance<Application>();
|
||||
const mountPath = `${this.config.server.path}/api/queue`;
|
||||
|
||||
const queues = this.collectQueues();
|
||||
if (!queues.length) {
|
||||
this.logger.warn('QueueDash not mounted: no queues available');
|
||||
app.use(mountPath, (_req, res) => {
|
||||
res.status(404).end();
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const guardMiddleware = async (
|
||||
req: Request,
|
||||
res: Response,
|
||||
next: NextFunction
|
||||
): Promise<void> => {
|
||||
try {
|
||||
const session = await this.authGuard.signIn(req, res);
|
||||
const userId = session?.user?.id;
|
||||
const isAdmin = userId ? await this.feature.isAdmin(userId) : false;
|
||||
if (!isAdmin) {
|
||||
res.status(404).end();
|
||||
return;
|
||||
}
|
||||
} catch (error) {
|
||||
this.logger.warn('QueueDash auth failed', error as Error);
|
||||
res.status(404).end();
|
||||
return;
|
||||
}
|
||||
|
||||
next();
|
||||
};
|
||||
|
||||
app.use(
|
||||
mountPath,
|
||||
guardMiddleware,
|
||||
createQueueDashExpressMiddleware({ ctx: { queues } })
|
||||
);
|
||||
this.logger.log(`QueueDash mounted on ${mountPath}`);
|
||||
}
|
||||
|
||||
private collectQueues(): QueueDashQueue[] {
|
||||
const queues: QueueDashQueue[] = [];
|
||||
|
||||
for (const name of QUEUES) {
|
||||
const queue = this.moduleRef.get<BullMQQueue>(getQueueToken(name), {
|
||||
strict: false,
|
||||
});
|
||||
|
||||
if (queue) {
|
||||
queues.push({ queue, displayName: name, type: 'bullmq' });
|
||||
}
|
||||
}
|
||||
|
||||
return queues;
|
||||
}
|
||||
}
|
||||
|
||||
@Module({
|
||||
imports: [AuthModule, FeatureModule],
|
||||
providers: [QueueDashboardService],
|
||||
})
|
||||
export class QueueDashboardModule {}
|
||||
@@ -129,10 +129,12 @@ export class WorkspaceStatsJob {
|
||||
private async withAdvisoryLock<T>(
|
||||
callback: (tx: Prisma.TransactionClient) => Promise<T>
|
||||
): Promise<T | null> {
|
||||
const lockIdSql = Prisma.sql`(${LOCK_NAMESPACE}::bigint << 32) + ${LOCK_KEY}::bigint`;
|
||||
|
||||
return await this.prisma.$transaction(
|
||||
async tx => {
|
||||
const [lock] = await tx.$queryRaw<{ locked: boolean }[]>`
|
||||
SELECT pg_try_advisory_lock(${LOCK_NAMESPACE}, ${LOCK_KEY}) AS locked
|
||||
SELECT pg_try_advisory_lock(${lockIdSql}) AS locked
|
||||
`;
|
||||
|
||||
if (!lock?.locked) {
|
||||
@@ -142,7 +144,7 @@ export class WorkspaceStatsJob {
|
||||
try {
|
||||
return await callback(tx);
|
||||
} finally {
|
||||
await tx.$executeRaw`SELECT pg_advisory_unlock(${LOCK_NAMESPACE}, ${LOCK_KEY})`;
|
||||
await tx.$executeRaw`SELECT pg_advisory_unlock(${lockIdSql})`;
|
||||
}
|
||||
},
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user