refactor(server): optimize workspace avatar handling

packages/backend/server/src/core/doc-renderer/controller.ts

@@ -5,7 +5,7 @@ import { Controller, Get, Logger, Req, Res } from '@nestjs/common';
 import type { Request, Response } from 'express';
 import isMobile from 'is-mobile';
 
-import { Config, metrics, URLHelper } from '../../base';
+import { Config, metrics } from '../../base';
 import { htmlSanitize } from '../../native';
 import { Public } from '../auth';
 import { DocReader } from '../doc';
@@ -52,8 +52,7 @@ export class DocRendererController {
   constructor(
     private readonly doc: DocReader,
     private readonly permission: PermissionService,
-    private readonly config: Config,
-    private readonly url: URLHelper
+    private readonly config: Config
   ) {
     this.webAssets = this.readHtmlAssets(
       join(this.config.projectRoot, 'static')
@@ -132,11 +131,7 @@ export class DocRendererController {
         return {
           title: workspaceContent.name,
           summary: '',
-          avatar: workspaceContent.avatarKey
-            ? this.url.link(
-                `/api/workspaces/${workspaceId}/blobs/${workspaceContent.avatarKey}`
-              )
-            : undefined,
+          avatar: workspaceContent.avatarUrl,
         };
       }
     }

packages/backend/server/src/core/doc/__tests__/reader-from-database.spec.ts

@@ -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';
@@ -8,6 +9,7 @@ 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';
 
@@ -16,6 +18,7 @@ const test = ava as TestFn<{
   app: TestingApp;
   docReader: DocReader;
   adapter: PgWorkspaceDocStorageAdapter;
+  blobStorage: WorkspaceBlobStorage;
 }>;
 
 test.before(async t => {
@@ -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;
 });
 
@@ -40,6 +44,10 @@ test.beforeEach(async t => {
   workspace = await t.context.models.workspace.create(user.id);
 });
 
+test.afterEach.always(() => {
+  mock.reset();
+});
+
 test.after.always(async t => {
   await t.context.app.close();
 });
@@ -162,8 +170,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');
@@ -184,7 +193,66 @@ 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',
+    avatarKey: '',
+    avatarUrl: undefined,
+  });
+  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',
+    avatarKey,
+    avatarUrl: `http://localhost:3010/api/workspaces/${workspace.id}/blobs/${avatarKey}`,
+  });
+  t.is(track.mock.callCount(), 1);
 });

packages/backend/server/src/core/doc/__tests__/reader-from-rpc.spec.ts

@@ -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',
+        avatarKey: '',
+      };
+    }
+  );
 
   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',
+    avatarKey: '',
   });
 });
 

packages/backend/server/src/core/doc/index.ts

@@ -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';
@@ -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,

packages/backend/server/src/core/doc/reader.ts

@@ -12,19 +12,27 @@ import {
   Config,
   CryptoHelper,
   getOrGenRequestId,
+  URLHelper,
   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 interface WorkspaceDocInfo {
+  id: string;
+  name: string;
+  avatarKey: string;
+  avatarUrl?: string;
+}
+
 export abstract class DocReader {
   constructor(protected readonly cache: Cache) {}
 
@@ -60,9 +68,9 @@ 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);
+    const cachedResult = await this.cache.get<WorkspaceDocInfo>(cacheKey);
     if (cachedResult) {
       return cachedResult;
     }
@@ -91,7 +99,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;
@@ -107,7 +115,9 @@ 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,
+    protected readonly url: URLHelper
   ) {
     super(cache);
   }
@@ -146,13 +156,32 @@ 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 avatarUrl: string | undefined;
+    if (content.avatarKey) {
+      avatarUrl = this.url.link(
+        `/api/workspaces/${workspaceId}/blobs/${content.avatarKey}`
+      );
+    }
+    return {
+      id: workspaceId,
+      name: content.name,
+      avatarKey: content.avatarKey,
+      avatarUrl,
+    };
+  }
+
+  private parseWorkspaceContent(bin: Uint8Array) {
     const doc = new YDoc();
-    applyUpdate(doc, docRecord.bin);
+    applyUpdate(doc, bin);
     return parseWorkspaceDoc(doc);
   }
 }
@@ -165,9 +194,11 @@ 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,
+    protected override readonly url: URLHelper
   ) {
-    super(cache, workspace);
+    super(cache, workspace, blobStorage, url);
   }
 
   private async fetch(
@@ -302,15 +333,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;