Skip to content

Commit

Permalink
refactor(server): optimize workspace avatar handling
Browse files Browse the repository at this point in the history
  • Loading branch information
@fengmk2
fengmk2 committed Feb 24, 2025
1 parent 60b994f commit 13e315f
Show file tree
Hide file tree
Showing 5 changed files with 140 additions and 45 deletions.
70 changes: 66 additions & 4 deletions packages/backend/server/src/core/doc/__tests__/reader-from-database.spec.ts
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { randomUUID } from 'node:crypto';
import { mock } from 'node:test';

import { User, Workspace } from '@prisma/client';
import ava, { TestFn } from 'ava';
Expand All @@ -8,14 +9,16 @@ import { createTestingApp, type TestingApp } from '../../../__tests__/utils';
import { AppModule } from '../../../app.module';
import { ConfigModule } from '../../../base/config';
import { Models } from '../../../models';
import { WorkspaceBlobStorage } from '../../storage/wrappers/blob';
import { DocReader, PgWorkspaceDocStorageAdapter } from '..';
import { DatabaseDocReader } from '../reader';
import { DatabaseDocReader, DEFAULT_WORKSPACE_AVATAR } from '../reader';

const test = ava as TestFn<{
models: Models;
app: TestingApp;
docReader: DocReader;
adapter: PgWorkspaceDocStorageAdapter;
blobStorage: WorkspaceBlobStorage;
}>;

test.before(async t => {
Expand All @@ -26,6 +29,7 @@ test.before(async t => {
t.context.models = app.get(Models);
t.context.docReader = app.get(DocReader);
t.context.adapter = app.get(PgWorkspaceDocStorageAdapter);
t.context.blobStorage = app.get(WorkspaceBlobStorage);
t.context.app = app;
});

Expand Down Expand Up @@ -162,8 +166,9 @@ test('should get doc content', async t => {
t.is(docContent, null);
});

test('should get workspace content', async t => {
test('should get workspace content with default avatar', async t => {
const { docReader } = t.context;
t.true(docReader instanceof DatabaseDocReader);

const doc = new YDoc();
const text = doc.getText('content');
Expand All @@ -184,7 +189,64 @@ test('should get workspace content', async t => {
user.id
);

// @ts-expect-error parseWorkspaceContent is private
const track = mock.method(docReader, 'parseWorkspaceContent', () => ({
name: 'Test Workspace',
avatarKey: '',
}));

const workspaceContent = await docReader.getWorkspaceContent(workspace.id);
// TODO(@fengmk2): should create a test ydoc with blocks
t.is(workspaceContent, null);
t.truthy(workspaceContent);
t.deepEqual(workspaceContent, {
id: workspace.id,
name: 'Test Workspace',
avatar: DEFAULT_WORKSPACE_AVATAR,
});
t.is(track.mock.callCount(), 1);
});

test('should get workspace content with custom avatar', async t => {
const { docReader, blobStorage } = t.context;
t.true(docReader instanceof DatabaseDocReader);

const doc = new YDoc();
const text = doc.getText('content');
const updates: Buffer[] = [];

doc.on('update', update => {
updates.push(Buffer.from(update));
});

text.insert(0, 'hello');
text.insert(5, 'world');
text.insert(5, ' ');

await t.context.adapter.pushDocUpdates(
workspace.id,
workspace.id,
updates,
user.id
);

const avatarKey = randomUUID();
await blobStorage.put(
workspace.id,
avatarKey,
Buffer.from('mock avatar image data here')
);

// @ts-expect-error parseWorkspaceContent is private
const track = mock.method(docReader, 'parseWorkspaceContent', () => ({
name: 'Test Workspace',
avatarKey,
}));

const workspaceContent = await docReader.getWorkspaceContent(workspace.id);
t.truthy(workspaceContent);
t.deepEqual(workspaceContent, {
id: workspace.id,
name: 'Test Workspace',
avatar: 'bW9jayBhdmF0YXIgaW1hZ2UgZGF0YSBoZXJl',
});
t.is(track.mock.callCount(), 1);
});
23 changes: 15 additions & 8 deletions packages/backend/server/src/core/doc/__tests__/reader-from-rpc.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ import { Config, UserFriendlyError } from '../../../base';
import { ConfigModule } from '../../../base/config';
import { Models } from '../../../models';
import { DatabaseDocReader, DocReader, PgWorkspaceDocStorageAdapter } from '..';
import { RpcDocReader } from '../reader';
import { DEFAULT_WORKSPACE_AVATAR, RpcDocReader } from '../reader';

const test = ava as TestFn<{
models: Models;
Expand Down Expand Up @@ -302,17 +302,24 @@ test('should return null when doc content not exists', async t => {

test('should get workspace content from doc service rpc', async t => {
const { docReader, databaseDocReader } = t.context;
mock.method(databaseDocReader, 'getWorkspaceContent', async () => {
return {
name: 'test name',
avatarKey: 'avatar key',
};
});
const track = mock.method(
databaseDocReader,
'getWorkspaceContent',
async () => {
return {
id: workspace.id,
name: 'test name',
avatar: DEFAULT_WORKSPACE_AVATAR,
};
}
);

const workspaceContent = await docReader.getWorkspaceContent(workspace.id);
t.is(track.mock.callCount(), 1);
t.deepEqual(workspaceContent, {
id: workspace.id,
name: 'test name',
avatarKey: 'avatar key',
avatar: DEFAULT_WORKSPACE_AVATAR,
});
});

Expand Down
3 changes: 2 additions & 1 deletion packages/backend/server/src/core/doc/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import { Module } from '@nestjs/common';

import { PermissionModule } from '../permission';
import { QuotaModule } from '../quota';
import { StorageModule } from '../storage';
import { PgUserspaceDocStorageAdapter } from './adapters/userspace';
import { PgWorkspaceDocStorageAdapter } from './adapters/workspace';
import { DocEventsListener } from './event';
Expand All @@ -12,7 +13,7 @@ import { DocStorageOptions } from './options';
import { DatabaseDocReader, DocReader, DocReaderProvider } from './reader';

@Module({
imports: [QuotaModule, PermissionModule],
imports: [QuotaModule, PermissionModule, StorageModule],
providers: [
DocStorageOptions,
PgWorkspaceDocStorageAdapter,
Expand Down
63 changes: 51 additions & 12 deletions packages/backend/server/src/core/doc/reader.ts
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { FactoryProvider, Injectable, Logger } from '@nestjs/common';
import { ModuleRef } from '@nestjs/core';
import { getStreamAsBuffer } from 'get-stream';
import {
applyUpdate,
diffUpdate,
Expand All @@ -14,17 +15,29 @@ import {
getOrGenRequestId,
UserFriendlyError,
} from '../../base';
import { WorkspaceBlobStorage } from '../storage';
import {
type PageDocContent,
parsePageDoc,
parseWorkspaceDoc,
type WorkspaceDocContent,
} from '../utils/blocksuite';
import { PgWorkspaceDocStorageAdapter } from './adapters/workspace';
import { type DocDiff, type DocRecord } from './storage';

const DOC_CONTENT_CACHE_7_DAYS = 7 * 24 * 60 * 60 * 1000;

export const DEFAULT_WORKSPACE_AVATAR =
'iVBORw0KGgoAAAANSUhEUgAAABgAAAAYCAYAAADgdz34AAAACXBIWXMAAAsTAAALEwEAmpwYAAAAAXNSR0IArs4c6QAAAARnQU1BAACxjwv8YQUAAAQtSURBVHgBfVa9jhxFEK6q7rkf+4T2AgdIIC0ZoXkBuNQJtngBuIzs1hIRye1FhL438D0CRgKRGUeE6wwkhHYlkE2AtGdkbN/MdJe/qu7Z27PWnnG5Znq7v/rqd47pHddkNh/918tR1/FBamXc9zxOPVFKfJ4yP86qD1LD3/986/3F2zB40+LXv83HrHq/6+gAoNS1kF4odUz2nhJRTkI5E6mD6Bk1crLJkLy5cHc+P4ohzxLng8RKLqKUq6hkUtBSe8Zvdmfir7TT2a0fnkzeaeCbv/44ztSfZskjP2ygVRM0mbYTpgHMMMS8CsIIj/c+//Hp8UYD3z758whQUwdeEwPjAZQLqJhI0VxB2MVco+kXP/0zuZKD6dP5uM397ELzqEtMba/UJ4t7iXeq8U94z52Q+js09qjlIXMxAEsRDJpI59dVPzlDTooHko7BdlR2FcYmAtbGMmAt2mFI4yDQkIjtEQkxUAMKAPD9SiOK4b578N0S7Nt+fqFKbTbmRD1YGXurEmdtnjjz4kFuIV0gtWewV62hMHBY2gpEOw3Rnmztx9jnO72xzTV/YkzgNmgkiypeYJdCLjonqyAAg7VCshVpjTbD08HbxrySdhKxcDvoJTA5gLvpeXVQ+K340WKea9UkNeZVqGSba/IbF6athj+LUeRmRCyiAVnlAKhJJQfmugGZ28ZWna24RGzwNUNUqpWGf6HkajvAgNA4NsSjHgcb9obx+k5c3DUttcwd3NcHxpVurXQ2d4MZACGw9TwEHsdtbEwytL1xywAGcxavjoH1quLVywuGi+aBhFWexRilFSwK0QzgdUdkkVMeKw4wijrgxjzz2CefCRZn+21ViOWW4Ym9nNnyFLMbMS8ivNhGP8RdlgUojBkuBLDpEPi+5LpWiDURgFkKOIIckJTgN/sZ84KtKkKpDnsOZiTQ47jD4ZGwHghbw6AXIL3lo5Zg6Tp2AwIAyYJ8BRzGfmfPl6kI7HOLUdN2LIg+4IfL5SiFdvkK4blI6h50qda7jQI0CUMLdEhFIkqtQciMvXsgpaZ1pWtVUfrIa+TX5/8+RBcftAhTa91r8ycXA5ZxBqhAh2zgVagUAddxMkxfF/JxfvbpB+8d2jhBtsPhtuqsE0HJlhxYeHKdkCU8xUCos8dmkDdnGaOlJ1yy9dM52J2spqldvz9fTgB4z+aQd2kqjUY2KU2s4dTT7ezD0AqDAbvZiKF/VO9+fGPv9IoBu+b/P5ti6djDY+JlSg4ug1jc6fJbMAx9/3b4CNGTD/evT698D9avv188m4gKvko8MiMeJC3jmOvU9MSuHXZohAVpOrmxd+10HW/jR3/58uU45TRFt35ZR2XpY61DzW+tH3z/7xdM8sP93d3Fm1gbDawbEtU7CMtt/JVxEw01Kh7RAmoBE4+u7eycYv38bRivAZbdHBtPrwOHAAAAAElFTkSuQmCC';

export interface WorkspaceDocInfo {
id: string;
name: string;
/**
* Base64 encoded avatar
*/
avatar: string;
}

export abstract class DocReader {
constructor(protected readonly cache: Cache) {}

Expand Down Expand Up @@ -60,10 +73,11 @@ export abstract class DocReader {

async getWorkspaceContent(
workspaceId: string
): Promise<WorkspaceDocContent | null> {
): Promise<WorkspaceDocInfo | null> {
const cacheKey = this.cacheKey(workspaceId, workspaceId);
const cachedResult = await this.cache.get<WorkspaceDocContent>(cacheKey);
if (cachedResult) {
const cachedResult = await this.cache.get<WorkspaceDocInfo>(cacheKey);
// if avatar not exists, means the cache result is the old one, override it
if (cachedResult && cachedResult.avatar) {
return cachedResult;
}

Expand Down Expand Up @@ -91,7 +105,7 @@ export abstract class DocReader {

protected abstract getWorkspaceContentWithoutCache(
workspaceId: string
): Promise<WorkspaceDocContent | null>;
): Promise<WorkspaceDocInfo | null>;

protected docDiff(update: Uint8Array, stateVector?: Uint8Array) {
const missing = stateVector ? diffUpdate(update, stateVector) : update;
Expand All @@ -107,7 +121,8 @@ export abstract class DocReader {
export class DatabaseDocReader extends DocReader {
constructor(
protected override readonly cache: Cache,
protected readonly workspace: PgWorkspaceDocStorageAdapter
protected readonly workspace: PgWorkspaceDocStorageAdapter,
protected readonly blobStorage: WorkspaceBlobStorage
) {
super(cache);
}
Expand Down Expand Up @@ -146,13 +161,36 @@ export class DatabaseDocReader extends DocReader {

protected override async getWorkspaceContentWithoutCache(
workspaceId: string
): Promise<WorkspaceDocContent | null> {
): Promise<WorkspaceDocInfo | null> {
const docRecord = await this.workspace.getDoc(workspaceId, workspaceId);
if (!docRecord) {
return null;
}
const content = this.parseWorkspaceContent(docRecord.bin);
if (!content) {
return null;
}
let avatar: string | undefined;
if (content.avatarKey) {
const avatarBlob = await this.blobStorage.get(
workspaceId,
content.avatarKey
);

if (avatarBlob.body) {
avatar = (await getStreamAsBuffer(avatarBlob.body)).toString('base64');
}
}
return {
id: workspaceId,
name: content.name,
avatar: avatar ?? DEFAULT_WORKSPACE_AVATAR,
};
}

private parseWorkspaceContent(bin: Uint8Array) {
const doc = new YDoc();
applyUpdate(doc, docRecord.bin);
applyUpdate(doc, bin);
return parseWorkspaceDoc(doc);
}
}
Expand All @@ -165,9 +203,10 @@ export class RpcDocReader extends DatabaseDocReader {
private readonly config: Config,
private readonly crypto: CryptoHelper,
protected override readonly cache: Cache,
protected override readonly workspace: PgWorkspaceDocStorageAdapter
protected override readonly workspace: PgWorkspaceDocStorageAdapter,
protected override readonly blobStorage: WorkspaceBlobStorage
) {
super(cache, workspace);
super(cache, workspace, blobStorage);
}

private async fetch(
Expand Down Expand Up @@ -302,15 +341,15 @@ export class RpcDocReader extends DatabaseDocReader {

protected override async getWorkspaceContentWithoutCache(
workspaceId: string
): Promise<WorkspaceDocContent | null> {
): Promise<WorkspaceDocInfo | null> {
const url = `${this.config.docService.endpoint}/rpc/workspaces/${workspaceId}/content`;
const accessToken = this.crypto.sign(workspaceId);
try {
const res = await this.fetch(accessToken, url, 'GET');
if (!res) {
return null;
}
return (await res.json()) as WorkspaceDocContent;
return (await res.json()) as WorkspaceDocInfo;
} catch (e) {
if (e instanceof UserFriendlyError) {
throw e;
Expand Down
26 changes: 6 additions & 20 deletions packages/backend/server/src/core/workspaces/resolvers/service.ts
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import { Injectable, Logger } from '@nestjs/common';
import { PrismaClient } from '@prisma/client';
import { getStreamAsBuffer } from 'get-stream';

import {
Cache,
Expand All @@ -11,11 +10,8 @@ import {
} from '../../../base';
import { Models } from '../../../models';
import { DocReader } from '../../doc';
import { DEFAULT_WORKSPACE_AVATAR } from '../../doc/reader';
import { PermissionService, WorkspaceRole } from '../../permission';
import { WorkspaceBlobStorage } from '../../storage';

export const defaultWorkspaceAvatar =
'iVBORw0KGgoAAAANSUhEUgAAABgAAAAYCAYAAADgdz34AAAACXBIWXMAAAsTAAALEwEAmpwYAAAAAXNSR0IArs4c6QAAAARnQU1BAACxjwv8YQUAAAQtSURBVHgBfVa9jhxFEK6q7rkf+4T2AgdIIC0ZoXkBuNQJtngBuIzs1hIRye1FhL438D0CRgKRGUeE6wwkhHYlkE2AtGdkbN/MdJe/qu7Z27PWnnG5Znq7v/rqd47pHddkNh/918tR1/FBamXc9zxOPVFKfJ4yP86qD1LD3/986/3F2zB40+LXv83HrHq/6+gAoNS1kF4odUz2nhJRTkI5E6mD6Bk1crLJkLy5cHc+P4ohzxLng8RKLqKUq6hkUtBSe8Zvdmfir7TT2a0fnkzeaeCbv/44ztSfZskjP2ygVRM0mbYTpgHMMMS8CsIIj/c+//Hp8UYD3z758whQUwdeEwPjAZQLqJhI0VxB2MVco+kXP/0zuZKD6dP5uM397ELzqEtMba/UJ4t7iXeq8U94z52Q+js09qjlIXMxAEsRDJpI59dVPzlDTooHko7BdlR2FcYmAtbGMmAt2mFI4yDQkIjtEQkxUAMKAPD9SiOK4b578N0S7Nt+fqFKbTbmRD1YGXurEmdtnjjz4kFuIV0gtWewV62hMHBY2gpEOw3Rnmztx9jnO72xzTV/YkzgNmgkiypeYJdCLjonqyAAg7VCshVpjTbD08HbxrySdhKxcDvoJTA5gLvpeXVQ+K340WKea9UkNeZVqGSba/IbF6athj+LUeRmRCyiAVnlAKhJJQfmugGZ28ZWna24RGzwNUNUqpWGf6HkajvAgNA4NsSjHgcb9obx+k5c3DUttcwd3NcHxpVurXQ2d4MZACGw9TwEHsdtbEwytL1xywAGcxavjoH1quLVywuGi+aBhFWexRilFSwK0QzgdUdkkVMeKw4wijrgxjzz2CefCRZn+21ViOWW4Ym9nNnyFLMbMS8ivNhGP8RdlgUojBkuBLDpEPi+5LpWiDURgFkKOIIckJTgN/sZ84KtKkKpDnsOZiTQ47jD4ZGwHghbw6AXIL3lo5Zg6Tp2AwIAyYJ8BRzGfmfPl6kI7HOLUdN2LIg+4IfL5SiFdvkK4blI6h50qda7jQI0CUMLdEhFIkqtQciMvXsgpaZ1pWtVUfrIa+TX5/8+RBcftAhTa91r8ycXA5ZxBqhAh2zgVagUAddxMkxfF/JxfvbpB+8d2jhBtsPhtuqsE0HJlhxYeHKdkCU8xUCos8dmkDdnGaOlJ1yy9dM52J2spqldvz9fTgB4z+aQd2kqjUY2KU2s4dTT7ezD0AqDAbvZiKF/VO9+fGPv9IoBu+b/P5ti6djDY+JlSg4ug1jc6fJbMAx9/3b4CNGTD/evT698D9avv188m4gKvko8MiMeJC3jmOvU9MSuHXZohAVpOrmxd+10HW/jR3/58uU45TRFt35ZR2XpY61DzW+tH3z/7xdM8sP93d3Fm1gbDawbEtU7CMtt/JVxEw01Kh7RAmoBE4+u7eycYv38bRivAZbdHBtPrwOHAAAAAElFTkSuQmCC';

export type InviteInfo = {
workspaceId: string;
Expand All @@ -28,7 +24,6 @@ export class WorkspaceService {
private readonly logger = new Logger(WorkspaceService.name);

constructor(
private readonly blobStorage: WorkspaceBlobStorage,
private readonly cache: Cache,
private readonly doc: DocReader,
private readonly mailer: MailService,
Expand Down Expand Up @@ -64,24 +59,15 @@ export class WorkspaceService {
}

async getWorkspaceInfo(workspaceId: string) {
const workspaceContent = await this.doc.getWorkspaceContent(workspaceId);

let avatar = defaultWorkspaceAvatar;
if (workspaceContent?.avatarKey) {
const avatarBlob = await this.blobStorage.get(
workspaceId,
workspaceContent.avatarKey
);

if (avatarBlob.body) {
avatar = (await getStreamAsBuffer(avatarBlob.body)).toString('base64');
}
const workspace = await this.doc.getWorkspaceContent(workspaceId);
if (workspace) {
return workspace;
}

return {
avatar,
id: workspaceId,
name: workspaceContent?.name ?? 'Untitled Workspace',
name: 'Untitled Workspace',
avatar: DEFAULT_WORKSPACE_AVATAR,
};
}

Expand Down

0 comments on commit 13e315f

Please sign in to comment.