refactor(server): separate s3 & r2 storage to plugin (#5805)

packages/backend/server/src/app.module.ts

@@ -29,6 +29,7 @@ import { MailModule } from './fundamentals/mailer';
 import { MetricsModule } from './fundamentals/metrics';
 import { PrismaModule } from './fundamentals/prisma';
 import { SessionModule } from './fundamentals/session';
+import { StorageProviderModule } from './fundamentals/storage';
 import { RateLimiterModule } from './fundamentals/throttler';
 import { WebSocketModule } from './fundamentals/websocket';
 import { pluginsMap } from './plugins';
@@ -43,6 +44,7 @@ export const FunctionalityModules = [
   RateLimiterModule,
   SessionModule,
   MailModule,
+  StorageProviderModule,
 ];
 
 export class AppModuleBuilder {

packages/backend/server/src/config/affine.self.ts

@@ -20,19 +20,19 @@ const env = process.env;
 AFFiNE.metrics.enabled = !AFFiNE.node.test;
 
 if (env.R2_OBJECT_STORAGE_ACCOUNT_ID) {
-  AFFiNE.storage.providers.r2 = {
+  AFFiNE.plugins.use('cloudflare-r2', {
     accountId: env.R2_OBJECT_STORAGE_ACCOUNT_ID,
     credentials: {
       accessKeyId: env.R2_OBJECT_STORAGE_ACCESS_KEY_ID!,
       secretAccessKey: env.R2_OBJECT_STORAGE_SECRET_ACCESS_KEY!,
     },
-  };
-  AFFiNE.storage.storages.avatar.provider = 'r2';
+  });
+  AFFiNE.storage.storages.avatar.provider = 'cloudflare-r2';
   AFFiNE.storage.storages.avatar.bucket = 'account-avatar';
   AFFiNE.storage.storages.avatar.publicLinkFactory = key =>
     `https://avatar.affineassets.com/${key}`;
 
-  AFFiNE.storage.storages.blob.provider = 'r2';
+  AFFiNE.storage.storages.blob.provider = 'cloudflare-r2';
   AFFiNE.storage.storages.blob.bucket = `workspace-blobs-${
     AFFiNE.affine.canary ? 'canary' : 'prod'
   }`;

packages/backend/server/src/config/affine.ts

@@ -87,8 +87,31 @@ AFFiNE.port = 3010;
 AFFiNE.plugins.use('redis', {
   /* override options */
 });
+//
+//
 // /* Payment Plugin */
 AFFiNE.plugins.use('payment', {
   stripe: { keys: {}, apiVersion: '2023-10-16' },
 });
 //
+//
+// /* Cloudflare R2 Plugin */
+// /* Enable if you choose to store workspace blobs or user avatars in Cloudflare R2 Storage Service */
+// AFFiNE.plugins.use('cloudflare-r2', {
+//   accountId: '',
+//   credentials: {
+//     accessKeyId: '',
+//     secretAccessKey: '',
+//   },
+// });
+//
+// /* AWS S3 Plugin */
+// /* Enable if you choose to store workspace blobs or user avatars in AWS S3 Storage Service */
+// AFFiNE.plugins.use('aws-s3', {
+//  credentials: {
+//    accessKeyId: '',
+//    secretAccessKey: '',
+// })
+// /* Update the provider of storages */
+// AFFiNE.storage.storages.blob.provider = 'r2';
+// AFFiNE.storage.storages.avatar.provider = 'r2';

packages/backend/server/src/core/storage/wrappers/avatar.ts

@@ -6,15 +6,18 @@ import type {
   PutObjectMetadata,
   StorageProvider,
 } from '../../../fundamentals';
-import { Config, createStorageProvider, OnEvent } from '../../../fundamentals';
+import { Config, OnEvent, StorageProviderFactory } from '../../../fundamentals';
 
 @Injectable()
 export class AvatarStorage {
   public readonly provider: StorageProvider;
   private readonly storageConfig: Config['storage']['storages']['avatar'];
 
-  constructor(private readonly config: Config) {
-    this.provider = createStorageProvider(this.config.storage, 'avatar');
+  constructor(
+    private readonly config: Config,
+    private readonly storageFactory: StorageProviderFactory
+  ) {
+    this.provider = this.storageFactory.create('avatar');
     this.storageConfig = this.config.storage.storages.avatar;
   }
 

packages/backend/server/src/core/storage/wrappers/blob.ts

@@ -6,10 +6,9 @@ import type {
   StorageProvider,
 } from '../../../fundamentals';
 import {
-  Config,
-  createStorageProvider,
   EventEmitter,
   OnEvent,
+  StorageProviderFactory,
 } from '../../../fundamentals';
 
 @Injectable()
@@ -18,9 +17,9 @@ export class WorkspaceBlobStorage {
 
   constructor(
     private readonly event: EventEmitter,
-    private readonly config: Config
+    private readonly storageFactory: StorageProviderFactory
   ) {
-    this.provider = createStorageProvider(this.config.storage, 'blob');
+    this.provider = this.storageFactory.create('blob');
   }
 
   async put(workspaceId: string, key: string, blob: BlobInputType) {

packages/backend/server/src/fundamentals/config/storage/index.ts

@@ -1,37 +1,34 @@
 import { homedir } from 'node:os';
 import { join } from 'node:path';
 
-import { S3ClientConfigType } from '@aws-sdk/client-s3';
-
-export type StorageProviderType = 'fs' | 'r2' | 's3';
 export interface FsStorageConfig {
   path: string;
 }
-export type R2StorageConfig = S3ClientConfigType & {
-  accountId: string;
-};
-export type S3StorageConfig = S3ClientConfigType;
 
-export type StorageTargetConfig<Ext = unknown> = {
+export interface StorageProvidersConfig {
+  fs: FsStorageConfig;
+}
+
+export type StorageProviderType = keyof StorageProvidersConfig;
+
+export type StorageConfig<Ext = unknown> = {
   provider: StorageProviderType;
   bucket: string;
 } & Ext;
 
+export interface StoragesConfig {
+  avatar: StorageConfig<{ publicLinkFactory: (key: string) => string }>;
+  blob: StorageConfig;
+}
+
 export interface AFFiNEStorageConfig {
   /**
    * All providers for object storage
    *
    * Support different providers for different usage at the same time.
    */
-  providers: {
-    fs?: FsStorageConfig;
-    s3?: S3StorageConfig;
-    r2?: R2StorageConfig;
-  };
-  storages: {
-    avatar: StorageTargetConfig<{ publicLinkFactory: (key: string) => string }>;
-    blob: StorageTargetConfig;
-  };
+  providers: StorageProvidersConfig;
+  storages: StoragesConfig;
 }
 
 export type StorageProviders = AFFiNEStorageConfig['providers'];

packages/backend/server/src/fundamentals/index.ts

@@ -24,6 +24,7 @@ export {
 export { PrismaService } from './prisma';
 export { SessionService } from './session';
 export * from './storage';
+export { type StorageProvider, StorageProviderFactory } from './storage';
 export { AuthThrottlerGuard, CloudThrottlerGuard, Throttle } from './throttler';
 export {
   getRequestFromHost,

packages/backend/server/src/fundamentals/storage/index.ts

@@ -1,36 +1,24 @@
-import { createRequire } from 'node:module';
+import { Global, Module } from '@nestjs/common';
 
-export const StorageProvide = Symbol('Storage');
+import { registerStorageProvider, StorageProviderFactory } from './providers';
+import { FsStorageProvider } from './providers/fs';
 
-let storageModule: typeof import('@affine/storage');
-try {
-  storageModule = await import('@affine/storage');
-} catch {
-  const require = createRequire(import.meta.url);
-  storageModule =
-    process.arch === 'arm64'
-      ? require('../../../storage.arm64.node')
-      : process.arch === 'arm'
-        ? require('../../../storage.armv7.node')
-        : require('../../../storage.node');
-}
+registerStorageProvider('fs', (config, bucket) => {
+  if (!config.storage.providers.fs) {
+    throw new Error('Missing fs storage provider configuration');
+  }
 
-export const mergeUpdatesInApplyWay = storageModule.mergeUpdatesInApplyWay;
+  return new FsStorageProvider(config.storage.providers.fs, bucket);
+});
 
-export const verifyChallengeResponse = async (
-  response: any,
-  bits: number,
-  resource: string
-) => {
-  if (typeof response !== 'string' || !response || !resource) return false;
-  return storageModule.verifyChallengeResponse(response, bits, resource);
-};
-
-export const mintChallengeResponse = async (resource: string, bits: number) => {
-  if (!resource) return null;
-  return storageModule.mintChallengeResponse(resource, bits);
-};
+@Global()
+@Module({
+  providers: [StorageProviderFactory],
+  exports: [StorageProviderFactory],
+})
+export class StorageProviderModule {}
 
+export * from './native';
 export type {
   BlobInputType,
   BlobOutputType,
@@ -39,5 +27,5 @@ export type {
   PutObjectMetadata,
   StorageProvider,
 } from './providers';
-export { createStorageProvider } from './providers';
-export { toBuffer } from './providers/utils';
+export { registerStorageProvider, StorageProviderFactory } from './providers';
+export { autoMetadata, toBuffer } from './providers/utils';

packages/backend/server/src/fundamentals/storage/native.ts

@@ -0,0 +1,30 @@
+import { createRequire } from 'node:module';
+
+let storageModule: typeof import('@affine/storage');
+try {
+  storageModule = await import('@affine/storage');
+} catch {
+  const require = createRequire(import.meta.url);
+  storageModule =
+    process.arch === 'arm64'
+      ? require('../../../storage.arm64.node')
+      : process.arch === 'arm'
+        ? require('../../../storage.armv7.node')
+        : require('../../../storage.node');
+}
+
+export const mergeUpdatesInApplyWay = storageModule.mergeUpdatesInApplyWay;
+
+export const verifyChallengeResponse = async (
+  response: any,
+  bits: number,
+  resource: string
+) => {
+  if (typeof response !== 'string' || !response || !resource) return false;
+  return storageModule.verifyChallengeResponse(response, bits, resource);
+};
+
+export const mintChallengeResponse = async (resource: string, bits: number) => {
+  if (!resource) return null;
+  return storageModule.mintChallengeResponse(resource, bits);
+};

packages/backend/server/src/fundamentals/storage/providers/index.ts

@@ -1,34 +1,37 @@
-import { AFFiNEStorageConfig, Storages } from '../../config/storage';
-import { FsStorageProvider } from './fs';
-import type { StorageProvider } from './provider';
-import { R2StorageProvider } from './r2';
-import { S3StorageProvider } from './s3';
+import { Injectable } from '@nestjs/common';
 
-export function createStorageProvider(
-  config: AFFiNEStorageConfig,
-  storage: Storages
-): StorageProvider {
-  const storageConfig = config.storages[storage];
-  const providerConfig = config.providers[storageConfig.provider] as any;
-  if (!providerConfig) {
+import { Config } from '../../config';
+import type { StorageProviderType, Storages } from '../../config/storage';
+import type { StorageProvider } from './provider';
+
+const availableProviders = new Map<
+  StorageProviderType,
+  (config: Config, bucket: string) => StorageProvider
+>();
+
+export function registerStorageProvider(
+  type: StorageProviderType,
+  providerFactory: (config: Config, bucket: string) => StorageProvider
+) {
+  availableProviders.set(type, providerFactory);
+}
+
+@Injectable()
+export class StorageProviderFactory {
+  constructor(private readonly config: Config) {}
+
+  create(storage: Storages): StorageProvider {
+    const storageConfig = this.config.storage.storages[storage];
+    const providerFactory = availableProviders.get(storageConfig.provider);
+
+    if (!providerFactory) {
       throw new Error(
-      `Failed to create ${storageConfig.provider} storage, configuration not correctly set`
+        `Unknown storage provider type: ${storageConfig.provider}`
       );
     }
 
-  if (storageConfig.provider === 's3') {
-    return new S3StorageProvider(providerConfig, storageConfig.bucket);
+    return providerFactory(this.config, storageConfig.bucket);
   }
-
-  if (storageConfig.provider === 'r2') {
-    return new R2StorageProvider(providerConfig, storageConfig.bucket);
-  }
-
-  if (storageConfig.provider === 'fs') {
-    return new FsStorageProvider(providerConfig, storageConfig.bucket);
-  }
-
-  throw new Error(`Unknown storage provider type: ${storageConfig.provider}`);
 }
 
 export type * from './provider';

packages/backend/server/src/plugins/config.ts

@@ -1,12 +1,15 @@
 import { GCloudConfig } from './gcloud/config';
 import { PaymentConfig } from './payment';
 import { RedisOptions } from './redis';
+import { R2StorageConfig, S3StorageConfig } from './storage';
 
 declare module '../fundamentals/config' {
   interface PluginsConfig {
     readonly payment: PaymentConfig;
     readonly redis: RedisOptions;
     readonly gcloud: GCloudConfig;
+    readonly 'cloudflare-r2': R2StorageConfig;
+    readonly 'aws-s3': S3StorageConfig;
   }
 
   export type AvailablePlugins = keyof PluginsConfig;

packages/backend/server/src/plugins/index.ts

@@ -2,9 +2,12 @@ import type { AvailablePlugins } from '../fundamentals/config';
 import { GCloudModule } from './gcloud';
 import { PaymentModule } from './payment';
 import { RedisModule } from './redis';
+import { AwsS3Module, CloudflareR2Module } from './storage';
 
 export const pluginsMap = new Map<AvailablePlugins, AFFiNEModule>([
   ['payment', PaymentModule],
   ['redis', RedisModule],
   ['gcloud', GCloudModule],
+  ['cloudflare-r2', CloudflareR2Module],
+  ['aws-s3', AwsS3Module],
 ]);

packages/backend/server/src/plugins/storage/index.ts

@@ -0,0 +1,40 @@
+import { OptionalModule } from '../../fundamentals';
+import { registerStorageProvider } from '../../fundamentals/storage';
+import { R2StorageProvider } from './providers/r2';
+import { S3StorageProvider } from './providers/s3';
+
+registerStorageProvider('cloudflare-r2', (config, bucket) => {
+  if (!config.plugins['cloudflare-r2']) {
+    throw new Error('Missing cloudflare-r2 storage provider configuration');
+  }
+
+  return new R2StorageProvider(config.plugins['cloudflare-r2'], bucket);
+});
+registerStorageProvider('aws-s3', (config, bucket) => {
+  if (!config.plugins['aws-s3']) {
+    throw new Error('Missing aws-s3 storage provider configuration');
+  }
+
+  return new S3StorageProvider(config.plugins['aws-s3'], bucket);
+});
+
+@OptionalModule({
+  requires: [
+    'plugins.cloudflare-r2.accountId',
+    'plugins.cloudflare-r2.credentials.accessKeyId',
+    'plugins.cloudflare-r2.credentials.secretAccessKey',
+  ],
+  if: config => config.flavor.graphql,
+})
+export class CloudflareR2Module {}
+
+@OptionalModule({
+  requires: [
+    'plugins.aws-s3.credentials.accessKeyId',
+    'plugins.aws-s3.credentials.secretAccessKey',
+  ],
+  if: config => config.flavor.graphql,
+})
+export class AwsS3Module {}
+
+export type { R2StorageConfig, S3StorageConfig } from './types';

packages/backend/server/src/plugins/storage/providers/r2.ts

@@ -1,10 +1,10 @@
 import { Logger } from '@nestjs/common';
 
-import { R2StorageConfig } from '../../config/storage';
+import type { R2StorageConfig } from '../types';
 import { S3StorageProvider } from './s3';
 
 export class R2StorageProvider extends S3StorageProvider {
-  override readonly type = 'r2' as any /* cast 'r2' to 's3' */;
+  override readonly type = 'cloudflare-r2' as any /* cast 'r2' to 's3' */;
 
   constructor(config: R2StorageConfig, bucket: string) {
     super(

packages/backend/server/src/plugins/storage/providers/s3.ts

@@ -11,21 +11,22 @@ import {
 } from '@aws-sdk/client-s3';
 import { Logger } from '@nestjs/common';
 
-import { S3StorageConfig } from '../../config/storage';
 import {
+  autoMetadata,
   BlobInputType,
   GetObjectMetadata,
   ListObjectsMetadata,
   PutObjectMetadata,
   StorageProvider,
-} from './provider';
-import { autoMetadata, toBuffer } from './utils';
+  toBuffer,
+} from '../../../fundamentals/storage';
+import type { S3StorageConfig } from '../types';
 
 export class S3StorageProvider implements StorageProvider {
   protected logger: Logger;
   protected client: S3Client;
 
-  readonly type = 's3';
+  readonly type = 'aws-s3';
 
   constructor(
     config: S3StorageConfig,

packages/backend/server/src/plugins/storage/types.ts

@@ -0,0 +1,16 @@
+import { S3ClientConfigType } from '@aws-sdk/client-s3';
+
+type WARNING = '__YOU_SHOULD_NOT_MANUALLY_CONFIGURATE_THIS_TYPE__';
+export type R2StorageConfig = S3ClientConfigType & {
+  accountId: string;
+};
+
+export type S3StorageConfig = S3ClientConfigType;
+
+declare module '../../fundamentals/config/storage' {
+  interface StorageProvidersConfig {
+    // the type here is only existing for extends [StorageProviderType] with better type inference and checking.
+    'cloudflare-r2'?: WARNING;
+    'aws-s3'?: WARNING;
+  }
+}