NestJS và Prisma: Áp dụng Database read replicas, hỗ trợ Typescript

Như các bài Ant của mình đã viết, Prisma là một Typescript ORM phổ biến và được ứng dụng khá nhiều trong các dự án vừa và nhỏ, đặc biệt là các dự án NestJS.

Tuy nhiên, khi dự án bắt đã đi hoạt động và có nhiều lượng truy cập, server sẽ có dấu hiện quá tải. Chúng ta cần phải áp dụng các biện pháp để nâng tải server: caching, autoscale,… và áp dụng read replicas là một trong số đó.

Prisma service & read replica plugin

Dưới đây là prisma.service.ts thường ứng dụng trong project NestJS

import { Injectable, OnModuleDestroy, OnModuleInit } from '@nestjs/common';
import { PrismaClient } from '@prisma/client';

@Injectable()
export class PrismaService
  extends PrismaClient
  implements OnModuleInit, OnModuleDestroy
{
  async onModuleInit() {
    await this.$connect();
  }

  async onModuleDestroy() {
    await this.$disconnect();
  }
}

Thêm @prisma/extension-read-replicas vào project như hướng dẫn https://www.prisma.io/docs/orm/prisma-client/setup-and-configuration/read-replicas

Tài liệu hướng dẫn của Prisma sẽ trông như dưới đây:

// init single
const prisma = new PrismaClient().$extends(
  readReplicas({
    url: process.env.DATABASE_URL_REPLICA,
  })
)

// init multiple replicas
const prisma = new PrismaClient().$extends(
  readReplicas({
    url: [
      process.env.DATABASE_URL_REPLICA_1,
      process.env.DATABASE_URL_REPLICA_2,
    ],
  })
)

// query
const posts = await prisma.$primary().post.findMany()
const result = await prisma.$replica().user.findFirst(...)

Tuy nhiên, với NestJS, Prisma hoạt động dưới dạng Service Module, nên chúng ta không thể init PrismaClient ở global được.

Đầu tiên, trong Prisma Service, chúng ta ngăn việc Prisma Module luôn mở kết nối tới database khi Prisma Module init, chúng ta sẽ để Prisma tự quyết định connection string của primary hay replica sẽ cần kết nối khi query.

  async onModuleInit() {
    // lazy connect to db, to determine which primary/replica to use
    // await this.$connect();
  }

Sau đó, chúng ta lấy type Prisma Client sau khi extends read replicas plugin, dưới tên ExtendedPrismaClient

// extendedPrismaClient is for type only
const extendedPrismaClient = () =>
  new PrismaClient().$extends(
    readReplicas({
      url: [],
    }),
  );

export type ExtendedPrismaClient = ReturnType<typeof extendedPrismaClient>;

Chúng ta thêm vào Prisma Service:

• private readReplicasUrl: nhận env là 1 connection string sang read replica
• private extendedPrismaClient: ExtendedPrismaClient
• get property client, sẽ extend Prisma service hiện tại this vào private property extendedPrismaClient lần đầu tiên, sau đó sẽ luôn trả về extendedPrismaClient.

  extendsClient = (prismaClient: PrismaClient, url?: string) => {
    if (url) {
      const replicaClient = new PrismaClient({
        datasources: {
         db: { url },
        },
      });

      const extendedPrismaClient = prismaClient.$extends(
        readReplicas({
          replicas: [replicaClient]
        }),
      );

      return extendedPrismaClient;
    }
    return prismaClient.$extends({});
  };

  get client() {
    if (this.customPrismaClient) {
      return this.customPrismaClient;
    }
    this.customPrismaClient = this.extendsClient(
      this,
      this.readReplicasUrl,
    ) as unknown as ExtendedPrismaClient;
    return this.customPrismaClient;
  }

Như vậy, bằng cách sử dụng

await this.prismaService.client.findMany(...)
await this.prismaService.client.findOne(...)

Chúng ta đã có thể áp dụng Prisma read replica plugin vào NestJS project rồi.

Hỗ trợ nhiều read replicas.

Bạn có thể để ý ở implementation code ở trên, `[replicaClient]` được để vào trong array chứ không phải 1 client duy nhất.

Lý do nên sử dụng new PrismaClient vì

• các Client này sẽ không tự kết nối đến database nếu không được chỉ định .connect bởi replica plugin
• nếu dùng url string thay cho Client, read replica sẽ không hoạt động.

Chúng ta sẽ thay đổi các implementation đưới đây để hỗ trợ nhiều hơn 1 read replica url

  constructor() {
    const urls = process.env.DATABASE_URL_READ_REPLICAS?.split(',') || [];
    this.readReplicasUrls = urls;
  }

  extendsClient = (prismaClient: PrismaClient, urls?: string) => {
    if (urls?.length) {
      const replicaClients = urls.map((url) => 
        new PrismaClient({
         datasources: {
          db: { url },
         },
       })
      );

      const extendedPrismaClient = prismaClient.$extends(
        readReplicas({
          replicas: replicaClients
        }),
      );

      return extendedPrismaClient;
    }
    return prismaClient.$extends({});
  };

  get client() {
    if (this.customPrismaClient) {
      return this.customPrismaClient;
    }
    this.customPrismaClient = this.extendsClient(
      this,
      this.readReplicasUrls,
    ) as unknown as ExtendedPrismaClient;
    return this.customPrismaClient;
  }

Như vậy, khi truyền env DATABASE_URL_READ_REPLICAS là các connection string cách nhau bởi dấu phẩy, thì Prisma Service sẽ tự tạo đúng số read replicas và khi query, Prisma Read replica plugin sẽ tự connect đến primary database hoặc read replica tương ứng.

Dưới đây là full implementation của Prisma Service

import { Injectable, OnModuleDestroy, OnModuleInit, INestApplication } from '@nestjs/common';
import { PrismaClient } from '@prisma/client';
import { DefaultArgs } from '@prisma/client/runtime/library';
import { readReplicas } from '@prisma/extension-read-replicas';

// get type purpose, do not use
const extendedPrismaClient = () =>
  new PrismaClient().$extends(
    readReplicas({
      url: [],
    }),
  );

export type ExtendedPrismaClient = ReturnType<typeof extendedPrismaClient>;

@Injectable()
export class PrismaService extends PrismaClient<{}, 'error', DefaultArgs> implements OnModuleInit, OnModuleDestroy {
  readReplicasUrls: string[] = [];
  customPrismaClient: ExtendedPrismaClient;
  constructor() {
    super({});
    const DATABASE_URL_READ_REPLICAS = process.env.DATABASE_URL_READ_REPLICAS;
    const urls = DATABASE_URL_READ_REPLICAS?.split(',') || [];
    this.readReplicasUrls = urls;
  }

  extendsClient = (prismaClient: PrismaClient, urls?: string[]) => {
    if (urls?.length) {
      const extendedPrismaClient = prismaClient.$extends(
        readReplicas({
          replicas: urls.map((url) => {
            return new PrismaClient({
              datasources: {
                db: { url },
              },
            });
          }),
        }),
      );

      return extendedPrismaClient;
    }
    return prismaClient.$extends({});
  };

  async onModuleInit() {
    // lazy connect to db, to determine which primary/replica to use
    // await this.$connect();
  }

  async onModuleDestroy() {
    await this.$disconnect();
  }

  get client() {
    if (this.customPrismaClient) {
      return this.customPrismaClient;
    }
    this.customPrismaClient = this.extendsClient(this, this.readReplicasUrls) as unknown as ExtendedPrismaClient;
    return this.customPrismaClient;
  }

  $primary() {
    if (this.readReplicasUrls.length) {
      return (this.client as unknown as ExtendedPrismaClient).$primary();
    }
    return this;
  }

  $replica() {
    // replica should be used when queryRaw or queryRawUnsafe
    // default prisma.client won't use replicas for them
    if (this.readReplicasUrls.length) {
      return (this.client as unknown as ExtendedPrismaClient).$replica();
    }
    return this;
  }

  async enableShutdownHooks(app: INestApplication) {
    process.on('beforeExit', async () => {
      await this.$disconnect();
      app.close();
    });
  }
}

Lưu ý

• Cách áp dụng Prisma read replicas extension trong bài viết phụ thuộc vào custom prisma module/service, không tương thích với nestjs-prisma
• Prisma $queryRaw và $queryRawUnsafe sẽ luôn gọi primary database, để bắt query raw phải call read replica, hãy chỉ định replica bằng cách dùng $replica từ client: `prisma.client.$replica().queryRawUnsafe(…)`