slaypni · capt4ce · Feb 13, 2025
diff --git a/README.md b/README.md
@@ -114,3 +114,23 @@ export default class UserResolver {
   }
 }
 ```
+
+### Automatically use specific replica for mutation operation
+
+When using database with replica and fetching data after performing data mutation, sometimes the loader returns stale data because the typeorm data fetching uses `slave` replica. To overcome this, we need to setup `mutationReplica` to use when fetching data after performing mutation.
+
+```ts
+...
+
+const apollo = new ApolloServer({
+  schema,
+  plugins: [
+    ApolloServerLoaderPlugin({
+      typeormGetConnection: getConnection,  // for use with TypeORM
+      mutationReplica: 'master'
+    }),
+  ],
+});
+
+...
+```
diff --git a/src/decorators/typeorm/ExplicitLoaderImpl.ts b/src/decorators/typeorm/ExplicitLoaderImpl.ts
@@ -3,7 +3,7 @@ import DataLoader from "dataloader";
 import { Dictionary, groupBy, keyBy } from "lodash";
 import { UseMiddleware } from "type-graphql";
 import Container from "typedi";
-import type { Connection } from "typeorm";
+import type { Connection, QueryRunner } from "typeorm";
 import type { ColumnMetadata } from "typeorm/metadata/ColumnMetadata";
 import type { RelationMetadata } from "typeorm/metadata/RelationMetadata";
 import { TypeormLoaderOption } from "./TypeormLoader";
@@ -15,7 +15,7 @@ export function ExplicitLoaderImpl<V>(
   option?: TypeormLoaderOption
 ): PropertyDecorator {
   return (target: Object, propertyKey: string | symbol) => {
-    UseMiddleware(async ({ root, context }, next) => {
+    UseMiddleware(async ({ root, context, info }, next) => {
       const tgdContext = context._tgdContext as TgdContext;
       if (tgdContext.typeormGetConnection == null) {
         throw Error("typeormGetConnection is not set");
@@ -52,7 +52,17 @@ export function ExplicitLoaderImpl<V>(
         relation.isManyToMany ?
           handleToMany :
         () => next();
-      return await handle<V>(keyFunc, root, tgdContext, relation);
+      return await handle<V>(
+        keyFunc,
+        root,
+        tgdContext,
+        relation,
+        tgdContext.mutationReplica && String(info.operation) === "mutation"
+          ? tgdContext
+              .typeormGetConnection()
+              .createQueryRunner(tgdContext.mutationReplica)
+          : undefined
+      );
     })(target, propertyKey);
   };
 }
@@ -88,13 +98,14 @@ async function handleToMany<V>(
   foreignKeyFunc: (root: any) => any | undefined,
   root: any,
   tgdContext: TgdContext,
-  relation: RelationMetadata
+  relation: RelationMetadata,
+  queryRunner?: QueryRunner
 ) {
   return handler(
     tgdContext,
     relation,
     relation.inverseEntityMetadata.primaryColumns,
-    (connection) => new ToManyDataloader<V>(relation, connection),
+    (connection) => new ToManyDataloader<V>(relation, connection, queryRunner),
     async (dataloader) => {
       const fks = foreignKeyFunc(root);
       return await dataloader.loadMany(fks);
@@ -106,13 +117,14 @@ async function handleToOne<V>(
   foreignKeyFunc: (root: any) => any | undefined,
   root: any,
   tgdContext: TgdContext,
-  relation: RelationMetadata
+  relation: RelationMetadata,
+  queryRunner?: QueryRunner
 ) {
   return handler(
     tgdContext,
     relation,
     relation.inverseEntityMetadata.primaryColumns,
-    (connection) => new ToOneDataloader<V>(relation, connection),
+    (connection) => new ToOneDataloader<V>(relation, connection, queryRunner),
     async (dataloader) => {
       const fk = foreignKeyFunc(root);
       return fk != null ? await dataloader.load(fk) : null;
@@ -157,12 +169,17 @@ async function handleOneToOneNotOwnerWithSelfKey<V>(
 function directLoader<V>(
   relation: RelationMetadata,
   connection: Connection,
-  grouper: string | ((entity: V) => any)
+  grouper: string | ((entity: V) => any),
+  queryRunner?: QueryRunner
 ) {
   return async (ids: readonly any[]) => {
     const entities = keyBy(
       await connection
-        .createQueryBuilder<V>(relation.type, relation.propertyName)
+        .createQueryBuilder<V>(
+          relation.type,
+          relation.propertyName,
+          queryRunner
+        )
         .whereInIds(ids)
         .getMany(),
       grouper
@@ -172,22 +189,37 @@ function directLoader<V>(
 }
 
 class ToManyDataloader<V> extends DataLoader<any, V> {
-  constructor(relation: RelationMetadata, connection: Connection) {
+  constructor(
+    relation: RelationMetadata,
+    connection: Connection,
+    queryRunner?: QueryRunner
+  ) {
     super(
-      directLoader(relation, connection, (entity) =>
-        relation.inverseEntityMetadata.primaryColumns[0].getEntityValue(entity)
+      directLoader(
+        relation,
+        connection,
+        (entity) =>
+          relation.inverseEntityMetadata.primaryColumns[0].getEntityValue(
+            entity
+          ),
+        queryRunner
       )
     );
   }
 }
 
 class ToOneDataloader<V> extends DataLoader<any, V> {
-  constructor(relation: RelationMetadata, connection: Connection) {
+  constructor(
+    relation: RelationMetadata,
+    connection: Connection,
+    queryRunner?: QueryRunner
+  ) {
     super(
       directLoader(
         relation,
         connection,
-        relation.inverseEntityMetadata.primaryColumns[0].propertyName
+        relation.inverseEntityMetadata.primaryColumns[0].propertyName,
+        queryRunner
       )
     );
   }

diff --git a/src/decorators/typeorm/ImplicitLoaderImpl.ts b/src/decorators/typeorm/ImplicitLoaderImpl.ts
@@ -2,17 +2,18 @@ import type { TgdContext } from "#/types/TgdContext";
 import DataLoader from "dataloader";
 import { UseMiddleware } from "type-graphql";
 import Container from "typedi";
-import type { Connection } from "typeorm";
+import type { Connection, QueryRunner } from "typeorm";
 import type { ColumnMetadata } from "typeorm/metadata/ColumnMetadata";
 import type { RelationMetadata } from "typeorm/metadata/RelationMetadata";
 
 export function ImplicitLoaderImpl<V>(): PropertyDecorator {
   return (target: Object, propertyKey: string | symbol) => {
-    UseMiddleware(async ({ root, context }, next) => {
+    UseMiddleware(async ({ root, context, info }, next) => {
       const tgdContext = context._tgdContext as TgdContext;
       if (tgdContext.typeormGetConnection == null) {
         throw Error("typeormGetConnection is not set");
       }
+
       const relation = tgdContext
         .typeormGetConnection()
         .getMetadata(target.constructor)
@@ -38,7 +39,18 @@ export function ImplicitLoaderImpl<V>(): PropertyDecorator {
       if (dataloaderCls == null) {
         return await next();
       }
-      return await handler<V>(root, tgdContext, relation, dataloaderCls);
+
+      return await handler<V>(
+        root,
+        tgdContext,
+        relation,
+        dataloaderCls,
+        tgdContext.mutationReplica && String(info.operation) === "mutation"
+          ? tgdContext
+              .typeormGetConnection()
+              .createQueryRunner(tgdContext.mutationReplica)
+          : undefined
+      );
     })(target, propertyKey);
   };
 }
@@ -48,8 +60,15 @@ async function handler<V>(
   { requestId, typeormGetConnection }: TgdContext,
   relation: RelationMetadata,
   dataloaderCls:
-    | (new (r: RelationMetadata, c: Connection) => DataLoader<string, V | null>)
-    | (new (r: RelationMetadata, c: Connection) => DataLoader<string, V[]>)
+    | (new (r: RelationMetadata, c: Connection, qr?: QueryRunner) => DataLoader<
+        string,
+        V | null
+      >)
+    | (new (r: RelationMetadata, c: Connection, qr?: QueryRunner) => DataLoader<
+        string,
+        V[]
+      >),
+  queryRunner?: QueryRunner
 ) {
   if (typeormGetConnection == null) {
     throw Error("Connection is not available");
@@ -60,7 +79,7 @@ async function handler<V>(
   if (!container.has(serviceId)) {
     container.set(
       serviceId,
-      new dataloaderCls(relation, typeormGetConnection())
+      new dataloaderCls(relation, typeormGetConnection(), queryRunner)
     );
   }
 
@@ -73,7 +92,11 @@ async function handler<V>(
 }
 
 class ToOneOwnerDataloader<V> extends DataLoader<string, V | null> {
-  constructor(relation: RelationMetadata, connection: Connection) {
+  constructor(
+    relation: RelationMetadata,
+    connection: Connection,
+    queryRunner?: QueryRunner
+  ) {
     super(async (pks) => {
       const relationName = relation.inverseRelation!.propertyName;
       const columns = relation.entityMetadata.primaryColumns;
@@ -83,7 +106,8 @@ class ToOneOwnerDataloader<V> extends DataLoader<string, V | null> {
         connection,
         pks,
         relationName,
-        columns
+        columns,
+        queryRunner
       );
       const referencedColumnNames = columns.map((c) => c.propertyPath);
       const entitiesByRelationKey = await getEntitiesByRelationKey(
@@ -97,7 +121,11 @@ class ToOneOwnerDataloader<V> extends DataLoader<string, V | null> {
 }
 
 class ToOneNotOwnerDataloader<V> extends DataLoader<string, V | null> {
-  constructor(relation: RelationMetadata, connection: Connection) {
+  constructor(
+    relation: RelationMetadata,
+    connection: Connection,
+    queryRunner?: QueryRunner
+  ) {
     super(async (pks) => {
       const inverseRelation = relation.inverseRelation!;
       const relationName = relation.propertyName;
@@ -108,7 +136,8 @@ class ToOneNotOwnerDataloader<V> extends DataLoader<string, V | null> {
         connection,
         pks,
         relationName,
-        columns
+        columns,
+        queryRunner
       );
       const referencedColumnNames = columns.map(
         (c) => c.referencedColumn!.propertyPath
@@ -124,7 +153,11 @@ class ToOneNotOwnerDataloader<V> extends DataLoader<string, V | null> {
 }
 
 class OneToManyDataloader<V> extends DataLoader<string, V[]> {
-  constructor(relation: RelationMetadata, connection: Connection) {
+  constructor(
+    relation: RelationMetadata,
+    connection: Connection,
+    queryRunner?: QueryRunner
+  ) {
     super(async (pks) => {
       const inverseRelation = relation.inverseRelation!;
       const columns = inverseRelation.joinColumns;
@@ -134,7 +167,8 @@ class OneToManyDataloader<V> extends DataLoader<string, V[]> {
         connection,
         pks,
         relation.propertyName,
-        columns
+        columns,
+        queryRunner
       );
       const referencedColumnNames = columns.map(
         (c) => c.referencedColumn!.propertyPath
@@ -150,7 +184,11 @@ class OneToManyDataloader<V> extends DataLoader<string, V[]> {
 }
 
 class ManyToManyDataloader<V> extends DataLoader<string, V[]> {
-  constructor(relation: RelationMetadata, connection: Connection) {
+  constructor(
+    relation: RelationMetadata,
+    connection: Connection,
+    queryRunner?: QueryRunner
+  ) {
     super(async (pks) => {
       const inversePropName = relation.inverseRelation!.propertyName;
       const { ownerColumns, inverseColumns } = relation.junctionEntityMetadata!;
@@ -163,7 +201,8 @@ class ManyToManyDataloader<V> extends DataLoader<string, V[]> {
         connection,
         pks,
         relationName,
-        columns
+        columns,
+        queryRunner
       );
       const referencedColumnNames = columns.map(
         (c) => c.referencedColumn!.propertyPath
@@ -183,13 +222,15 @@ async function findEntities<V>(
   connection: Connection,
   stringifiedPrimaryKeys: readonly string[],
   relationName: string,
-  columnMetas: ColumnMetadata[]
+  columnMetas: ColumnMetadata[],
+  queryRunner?: QueryRunner
 ): Promise<V[]> {
   const { Brackets } = await import("typeorm");
 
   const qb = connection.createQueryBuilder<V>(
     relation.type,
-    relation.propertyName
+    relation.propertyName,
+    queryRunner
   );
 
   if (relation.isOneToOneOwner || relation.isManyToOne) {

diff --git a/src/plugins/apollo-server/ApolloServerLoaderPlugin.ts b/src/plugins/apollo-server/ApolloServerLoaderPlugin.ts
@@ -7,6 +7,7 @@ import { v4 as uuidv4 } from "uuid";
 
 interface ApolloServerLoaderPluginOption {
   typeormGetConnection?: () => Connection;
+  mutationReplica?: boolean;
 }
 
 const ApolloServerLoaderPlugin = (
@@ -18,6 +19,7 @@ const ApolloServerLoaderPlugin = (
         _tgdContext: {
           requestId: uuidv4(),
           typeormGetConnection: option?.typeormGetConnection,
+          mutationReplica: option?.mutationReplica,
         } as TgdContext,
       });
     },

diff --git a/src/types/TgdContext.ts b/src/types/TgdContext.ts
@@ -1,6 +1,7 @@
-import type { Connection } from "typeorm";
+import type { Connection, ReplicationMode } from "typeorm";
 
 export interface TgdContext {
   requestId: string;
   typeormGetConnection?: () => Connection;
+  mutationReplica?: ReplicationMode;
 }