Achieving Causal Consistency in your GraphQL API using NestJS with MongoDB Server Sessions

"Achieving what now?!"

In this article we will be looking at how to achieve causal consistency, a concept that captures causal relationships between operations in a software system. We'll do this by looking at code examples of a Node API with GraphQL endpoints. We will be using NestJS as our API framework and Mongoose as our MongoDB adapter. At the end of this article, you will:

• be familiar with the concept of causal consistency
• be able to recognize the need for causal consistency
• be able to build a Node implementation using a modern technology stack

Causal Consistency

The concept of causal consistency captures the causal relationship between operations. This means that if we have some operation A that causes operation B, causal consistency provides the assurance of the order of operations, regardless of the current process they are running in. If we don't put this assurance in place, processes may not agree on the execution order. For example, we may have an application that transfers money between party A and party B, with a transfer going from A to B. The money transfer consists of two parts:

1. Withdraw money from bank account A
2. Deposit money to bank account B

Besides the order the steps, it is important that if either one of these operations fails, the other should be nullified. Imagine, for example, that only step 1 succeeds. We then have withdrawn money from A, but B never got it, and vice versa.

In the next couple of sections, we'll be looking at a simple use case defined by a set of requirements, along with some of the technologies used to implement causal consistency.

🍕 Use Case

For demonstration purposes, we take a simple pizza delivery Node API with a single GraphQL endpoint called orderPizza. The endpoint is used by a front-end web app to, indeed, order a pizza. Once the endpoint is hit by a user, the API requirements state that two things should occur:

1. Create a new order document in a pizza-orders collection.

2. Add an order reference to the requesting user document in a users collection.

   The latter will be used for efficient querying of orders-per-user, as we want to avoid scanning the whole pizza-orders collection each time a user requests its own orders. This requirement is important, as the underlying motivation reveals a data access pattern of retrieving orders-per-user that should be enforced within the API. Furthermore, requirements state that:

3. The operations should be causally consistent. In other words: execute in order and if one operation fails, the other(s) should be aborted.

Now that we have our requirements set up, we can start thinking about the implementation.

NestJS

NestJS is one of the most robust Node API frameworks out there. It provides a powerful modular system, along with features like TypeScript support and Dependency Injection, as seen in frameworks like Angular. A typical modern API that uses MongoDB will likely reach a point where multiple documents need to be updated within a single operation, although the data structure may not provide such atomicity. Instead, the challenge of performing a multi-document write operation arises. Even though each of these operations itself is atomic, the operation as a whole is not. This may cause inconsistencies regarding updates to our data, as our data may be causally related. This is where MongoDB Server Sessions come in. Fortunately, NestJS provides the concepts of Interceptors and Custom Decorators to integrate MongoDB Server Sessions in a nice and clean fashion.

MongoDB Server Sessions

The concept of MongoDB Server Sessions is somewhat equivalent to that of Transactions in Mongoose, as Mongoose Transactions are built on top of MongoDB Server Sessions. From the official docs:

MongoDB’s server sessions, or logical sessions, are the underlying framework used by client sessions to support Causal Consistency and retryable writes.

Please note that MongoDB Server Sessions have their limitations. Again from the official docs:

Server sessions are available for replica sets and sharded clusters only.

If you are having trouble setting up a MongoDB Replica Set in your local environment, you can alternatively set up a free MongoDB Cluster using MongoDB Atlas.

Implementation Time

NestJS Mongoose & GraphQL Setup

In order to implement our API requirements, we start by utilizing the @nestjs/graphql module. It provides some nice abstractions for fast implementations of GraphQL endpoints. Furthermore, we will be using the MongooseModule taken from @nestjs/mongoose for database interactions. We will start by adding these modules to our main NestJS application module:

app.module.ts

import { Module } from '@nestjs/common';
import { GraphQLModule } from '@nestjs/graphql';
import { MongooseModule } from '@nestjs/mongoose';

/**
 * You can use your own connection string for connecting
 * to a MongoDB Replica Set using the `forRoot` method of
 * the `MongooseModule`.
 */
@Module({
  imports: [
    MongooseModule.forRoot('mongodb+srv://user:password@cluster/database'),
    GraphQLModule.forRoot()
  ]
})
export class AppModule {}

Following the NestJS approach, our next objective is to create a PizzaOrder module in which we can define our MongoDB PizzaOrderSchema, our GraphQL PizzaOrderResolver and our PizzaOrderService. Using a feature-based file and folder structure, we need something like this:

.
📁 pizza-order/
  +-- 📁 models/
    +-- pizza-order.model.ts
  +-- 📁 schemas/
    +-- pizza-order.schema.ts
  +-- pizza-order.module.ts
  +-- pizza-order.resolver.ts
  +-- pizza-order.service.ts
+-- app.module.ts

This approach allows for a nice separation of concerns and divides responsibilities among different layers of the API. The PizzaOrderModule will look something like this:

pizza-order.module.ts

import { Module } from '@nestjs/common';
import { MongooseModule } from '@nestjs/mongoose';
import { PizzaOrderResolver } from './pizza-order.resolver';
import { PizzaOrderService } from './pizza-order.service';
import { PizzaOrderSchema } from './schemas/pizza-order.schema';

/**
 * Mongoose automatically registers a Model based
 * on the provided Schema using the `forFeature`
 * method on the `MongooseModule`.
 */
@Module({
  imports: [
    MongooseModule.forFeature([
      {
        name: 'PizzaOrder',
        schema: PizzaOrderSchema,
        collection: 'pizza-orders'
      }
    ])
  ],
  providers: [PizzaOrderResolver, PizzaOrderService]
})
export class PizzaOrderModule {}

Now we need to provide the actual GraphQL endpoint. We do this by defining our GraphQL Resolver:

pizza-order.resolver.ts

import { Args, Mutation, Resolver } from '@nestjs/graphql';
import { User } from '@pizza-api/user/models/user.model';
import { CurrentUser } from '@pizza-api/user/decorators/current-user';
import { PizzaOrderInput } from './dto/pizza-order.input';
import { PizzaOrder } from './models/pizza-order.model';
import { PizzaOrderService } from './pizza-order.service';

/**
 * The GraphQL Resolver takes care of all resolver-related
 * responsibilities. It should therefore **never** perform
 * any database interactions directly.
 */
@Resolver(() => PizzaOrder)
export class PizzaOrderResolver {
  /**
   * We inject the `PizzaOrderService`, which is used to
   * interact with MongoDB through Mongoose.
   */
  constructor(private pizzaOrderService: PizzaOrderService) {}

  /**
   * The mutation that our API should perform, based
   * on the requirements.
   */
  @Mutation()
  public orderPizza(
    @CurrentUser() currUser: User,
    @Args() newOrder: PizzaOrderInput
  ) {
    /**
     * This seems like a good place for our requirements:
     *
     * 1. Create a new order document in a `pizza-orders` collection.
     * 2. Add an order reference to the requesting user document in a `users` collection.
     * 3. The operations should be **causally consistent**.
     */
  }
}

Finally, we have our service to handle the database interaction through Mongoose:

pizza-order.service.ts

import { Injectable } from '@nestjs/common';
import { InjectModel } from '@nestjs/mongoose';
import { UserService } from '@pizza-api/user/user.service';
import { Document, Model } from 'mongoose';
import { PizzaOrder } from './interfaces/pizza-order';
import { PizzaOrderInput } from './dto/pizza-order.input';

/**
 * The Service is used to provide a layer through which
 * database interactions ought to be performed. Results are
 * then passed back to the GraphQL Resolver for optional
 * further processing.
 */
@Injectable()
export class PizzaOrderService {
  /**
   * The injected `UserService` is used to perform
   * a mutation on the User's document for adding a
   * reference to a Pizza Order.
   */
  constructor(
    @InjectModel('PizzaOrder')
    private pizzaOrderModel: Model<PizzaOrder & Document>,
    private userService: UserService
  ) {}

  public createOrder(userId: string, newOrder: PizzaOrderInput) {
    /**
     * This is the place where the actual creation
     * of the document takes place:
     *
     * 1. Create a new order document in a `pizza-orders` collection.
     */
  }

  public addOrderRefToUser(userId: string, orderId: string) {
    /**
     * This is where we add the order to the user:
     *
     * 2. Add an order reference to the requesting user document in a `users` collection.
     */
  }
}

NestJS Interceptors & Parameter Decorators

Now that the Pizza Order feature is set up, we continue by using NestJS Interceptors and Custom Decorators. Interceptors in NestJS are inspired by the Aspect Oriented Programming technique. One of the reasons for using Interceptors is to bind extra logic before or after the execution of a method. When we consider our requirements, we see that this concept fits our use case:

We want to start a session before the GraphQL Resolver method executes, and end a session after the method is executed.

Using this approach, we can create a single session per request and use that for all database interactions within that request. Therefore, the responsibility of managing the session is to be assigned to the GraphQL Resolver. We will expand our file structure as follows:

📁 decorators/
  +-- mongo-session.decorator.ts
📁 interceptors/
  +-- mongo-session.interceptor.ts
📁 pizza-order/
  +-- ...
+-- app.module.ts

This brings us to the NestJS Interceptor implementation:

mongo-session.interceptor.ts

import {
  CallHandler,
  ExecutionContext,
  Injectable,
  NestInterceptor
} from '@nestjs/common';
import { GqlExecutionContext } from '@nestjs/graphql';
import { InjectConnection } from '@nestjs/mongoose';
import { Connection } from 'mongoose';
import { from, Observable, throwError } from 'rxjs';
import { catchError, mapTo, switchMap, tap } from 'rxjs/operators';

export const MONGO_SESSION_KEY = 'MONGO_SESSION_KEY';

@Injectable()
export class MongoSessionInterceptor implements NestInterceptor {
  /**
   * A Mongoose Connection is injected, so we can use
   * it to create and start a session.
   */
  constructor(@InjectConnection() private readonly connection: Connection) {}

  /**
   * The `intercept` method return a Promise, since we
   * are `await`-ing the `startSession` method on the
   * connection. The returned Promise holds an Observable
   * we use to pipe different operators together.
   */
  async intercept(
    context: ExecutionContext,
    next: CallHandler
  ): Promise<Observable<any>> {
    const graphQlCtx = GqlExecutionContext.create(context);
    const ctx = graphQlCtx.getContext();
    const session = await this.connection.startSession();

    /**
     * We assign the newly created session to the
     * the GraphQL context, which allows us to access
     * the session later on via a `Custom Decorator`.
     */
    ctx[MONGO_SESSION_KEY] = session;
    session.startTransaction();

    /**
     * A chain of RxJS Observables is used to `pipe` the
     * operations together. In the end, we either commit
     * or abort the transactions, followed by a final statement
     * to end the transaction. 
     */
    return next.handle().pipe(
      switchMap(data =>
        from(
          session.inTransaction()
            ? session.commitTransaction()
            : Promise.resolve()
        ).pipe(mapTo(data))
      ),
      tap(() => session.inTransaction() && session.endSession()),
      catchError(async err => {
        if (session.inTransaction()) {
          await session.abortTransaction();
          session.endSession();
        }

        throw err;
      })
    );
  }
}

In the code above, we have some important steps to make the Interceptor behave as intended:

• The startSession() method is asynchronous, so we need to use await there.
• Since commitTransaction() returns void, we use a mapTo() operator to return the actual data passed along within the request.
• The catchError() callback is asynchronous, because we need to await the abortTransaction() method as it returns a Promise.
• At the end of the catchError() operator, the caught error is thrown again. This is an important thing to note, as we want our Interceptor to abide by the Single Responsibility Principle (SPR). Any transformation or formatting of error(s) should therefore be placed inside a(nother) dedicated Interceptor.
• We make sure that, before committing or aborting the transaction, we perform a check on whether the session is actually in transaction by using the inTransaction() method.

The last thing we now need to do is implement the Custom Decorator in order to access the session we created in the MongoSessionInterceptor, as we need to access the session inside our GraphQL Resolver. We therefore implement the MongoSession decorator as follows:

mongo-session.decorator.ts

import { createParamDecorator } from '@nestjs/common';
import { CustomParamFactory } from '@nestjs/common/interfaces';
import { GraphQLResolveInfo } from 'graphql';
import { MONGO_SESSION_KEY } from '../interceptors/mongo-session.interceptor';

export const mongoSessionFactory: CustomParamFactory = (
  data: any,
  [root, args, ctx, info]: [object, any, any, GraphQLResolveInfo]
) => ctx[MONGO_SESSION_KEY];

export const MongoSession = createParamDecorator(mongoSessionFactory);

You may notice that the above factory could also be directly implemented inside createParamDecorator(), which would, however, make testing the factory a lot harder. We therefore separate the two here.

We are now ready to start using our implementation in the PizzaOrderModule and friends.

Usage

Now that we have all the implementation for client usage in place, we can start adding our Interceptor and Custom Decorator to the GraphQL Resolver:

pizza-order.resolver.ts

import { ClientSession } from 'mongoose';
import { MongoSession } from '../decorators/mongo-session.decorator';
// ...

@Resolver(() => PizzaOrder)
export class PizzaOrderResolver {

  // ...

  /**
   * We can now start using our `MongoSession` decorator
   * inside this endpoint and pass that `session` along
   * with our service calls.
   */
  @Mutation()
  public async orderPizza(
    @CurrentUser() currUser: User,
    @Args() newOrder: PizzaOrderInput,
    @MongoSession() session: ClientSession
  ) {
    const order = await this.pizzaOrderService.createOrder(currUser.id, newOrder, session);
    await this.pizzaOrderService.addOrderRefToUser(currUser.id, order.id);

    return { success: true };
  }
}

The signature of the service methods now need to change in order to accept the session argument. We can then use the session for further database interactions:

// ...

@Injectable()
export class PizzaOrderService {

  // ...

  /**
   * We now use the `session` as an argument to
   * the `save()` method on a document. This ensures
   * that the operation becomes part of the current
   * transaction.
   */
  public createOrder(
    userId: string, 
    newOrder: PizzaOrderInput, 
    session: ClientSession
  ) {
    const doc = new this.pizzaOrderModel({
      ...newOrder,
      owner: userId
    });

    return doc.save({ session });
  }

  /**
   * The implementation of the `addPizzaOrderRef()` method
   * of the `UsersService` is left out for brevity here. You
   * can see that the `session` should also be used there, such
   * that the operation is added to the current session. Note
   * that sessions can be added using currying of Mongoose methods:
   *
   * @example
   *
   *   model
   *     .find()
   *     .session(session)
   */
  public addOrderRefToUser(
    userId: string, 
    orderId: string, 
    session: ClientSession
  ) {
    return this.usersService.addPizzaOrderRef(userId, orderId, session);
  }
}

It is also perfectly fine to make the addOrderRefToUser method private and call it from inside the createOrder function body. This allows the PizzaOrderResolver to make only a single call to the PizzaOrderService method createOrder without any additional calls. Please remember that the above implementation is for demonstration purposes.

Conclusion

This concludes our dive into causal consistency along with some code examples for implementation. You have now seen:

• what the concept of causal consistency consists of (pun intended)
• how to recognize the need for causal consistency from a set of requirements
• how to implement causal consistency in modern technology stack

Thank you reading this article. Hopefully it allows you to improve your development on Node API's, or really any software project for that matter. Feel free to leave any questions or comments below.

You can also follow me on Twitter

Comments

[bm-stschneider]

There is a small mistake at the pizza-order.module.ts right at the beginning. You import the service from the resolver.

[Nicky Lenaers]

Thanks again, this one is also updated.

[bm-stschneider]

in pizza-order.resolver.ts with decorators async and the decorator import is missing

[Nicky Lenaers]

Thanks for the feedback, I've updated the code examples.

[bm-stschneider]

Also there seems to be a dto and interfaces folder also? and what is this @pizza-api?

[Nicky Lenaers]

Yes, I use the interfaces folder to store interfaces that can be used across the API. The dto folder stores the specific GraphQL classes implementing those interfaces.

Regarding @pizza-api, this simply refers to some place in the API for demonstration purposes. I usually put a path in my tsconfig.json file to make these imports possible. So for example:

apps/pizza-order/tsconfig.json

{
  "extends": "../../tsconfig.json",
  "compilerOptions": {
    "types": ["node", "jest"],
    "paths": {
      "@pizza-api/user/*": ["apps/pizza-order/src/app/user/*"],
      // Other paths...
    }
  },
  "include": ["**/*.ts"]
}

[bm-stschneider]

in the pizza-order.service.ts the import { PizzaOrderInput } from './dto/pizza-order.input'; is missing

[Nicky Lenaers]

Thanks, this one is also updated.