Query caching using Nest.js and Typeorm

Recently one of our engineers was curious about something - how often does our API code runs the same SQL query multiple times while handling a single http request? A bit of local log analysis showed one query fired 29 times during a single page render. A page load didn’t equate to one HTTP request, but they knew it also didn’t call the API dozens of times.

Context: We run a monorepo with shared service layers, so a single HTTP path can traverse dozens of helpers.

Now our services code is pretty well decomposed and DRY, with shared helpers extracted into common packages, and clear boundaries that prevent teams from re-implementing the same logic in different places.

There aren’t multiple modules interacting with the same database tables, each table has a single repository layer that owns all reads and writes, and the database acts as a single source of truth.

But cross-cutting concerns still weave through many layers of the stack, sometimes triggering repeated database look-ups. For example, what are a user’s team memberships, what entitlements do they inherit from teams, and what are the settings for their environment?

So we set out to introduce a request-scoped in-memory query cache to our database adapter, so if the same query was run by multiple helpers at different points during the request, the results would be returned from the cache instead of sending the query to the database. Our hope was that this would reduce overall query load on the database.

The setup

The tech stack relevant to this project is:

• NestJS - our HTTP backend, handling dependency injection of our database adapter across services code.
• TypeORM - the library that issues queries to the database, with handy builder functions and typed object hydration.
• PostgreSQL - not super relevant to this project, except that it’s the source of truth for state in our application.

To limit the scope of how we’d actually implement this cache, there were some ground rules for the project:

1. No mass code rewrites – it must work without touching thousands of query builders.
2. Request-scoped only – the cache dies when the HTTP request ends.
3. Never return stale data – any write to the database blows away the cache immediately.

Initial optimism

With those constraints in mind, we started our scoping. At first glance, this looked like an easy project! NestJS can inject request-scoped instances of a provider, which could hold the contents of our cache. And TypeORM has built-in support for caching queries, with the ability to create custom cache providers. A little bit of glue code, and we could ship the feature by lunch!

Immovable objects meet TypeORM

Implementation started and immediately the plan ran into problems. Reviewing the TypeORM docs on caching in more detail indicated that the only way to enable the cache was to include cache(true) as part of every query builder. This violated our first rule of the project - don’t rewrite every query. Some sleuthing of the cache code in TypeORM’s GitHub repository uncovered an undocumented option that enabled the cache for every query. Phew!The next problem though was that TypeORM only supported cache invalidation on a time interval.

There were no explicit calls to remove entries from the cache in TypeORM’s code, even though the cache provider interface had methods corresponding to expiration. It expected the provider to check for stale entries when reading the cache. This violated our third rule of never returning stale data from the cache - no period of time could guarantee the data wasn’t stale. It’s entirely reasonable that code would update the database then immediate read that data back out.Here we had some flexibility in the use case we were using cacheing.

It’s nice for our app to avoid redundant database reads, but not strictly necessary. So we could extend our custom TypeORM repository code (this was in place to correct TypeORM behavior that drifted across versions) to simply drop the entire cache before save/update/delete queries. Any HTTP requests that only read data wasn’t at risk of staleness, and any requests that mutated data during its lifecycle would always fetch freshly from the database.

CLS to the rescue

Solving the TypeORM issues we moved on to NestJS, but here too there were hurdles. Because of the way our app instantiated our database configuration, it wasn’t possible to inject a request-scoped custom cache provider to pass to TypeORM, without writing custom middleware.

This felt dangerously close to violating rule 1 - minimize change surface area. While middleware for NestJS isn’t that complex (it’s running on top of express), if there was a bug introduced in middleware, it would affect the entire app. But also we found a much better approach, a library called nestjs-cls which is purpose built for our use case -

Continuation-local storage allows to store state and propagate it throughout callbacks and promise chains. It allows storing data throughout the lifetime of a web request or any other asynchronous duration. It is similar to thread-local storage in other languages.

NestJS CLS comes with a middleware which instantiates the provider and initializes the local storage. We can add the module to the root app and give it a setup function to create an empty cache object. Then we wrote a custom TypeORM query cache provider that wrapped the context-local storage, and injected the cache provider into our database configuration.

Going deeper

Time to look at the code, first the TypeORM query result cache. Note that we cache queries in a simple map of the query string to the query result object -

//file: src/.../cls-query-result-cache.ts
import { Inject, Injectable } from '@nestjs/common';
import { ClsService, ClsStore } from 'nestjs-cls';
import { QueryRunner } from 'typeorm';
import { QueryResultCache } from 'typeorm/cache/QueryResultCache';
import { QueryResultCacheOptions } from 'typeorm/cache/QueryResultCacheOptions';

// Add a type interface, so that we don't need to
// add type assertions when reading/writing the cache
interface QueryCache extends ClsStore {
  queryCache: Map<string | undefined, QueryResultCacheOptions | undefined>;
}

@Injectable()
export class ClsQueryResultCache implements QueryResultCache {
  constructor(
    @Inject(ClsService) private readonly cls: ClsService<QueryCache>,
  ) { }

  // nothing to do here
  async connect() { }
  async disconnect() { }
  async synchronize(_queryRunner?: QueryRunner) { }

  // TypeORM provides all of the query options when reading
  // from the cache, but we only care about the query string
  async getFromCache(
    options: QueryResultCacheOptions,
    _queryRunner?: QueryRunner,
  ): Promise<QueryResultCacheOptions | undefined> {
    if (!this.cls.isActive()) {
      return;
    }

    const cache = this.cls.get('queryCache');
    return cache.get(options.query);
  }

  // Just like reading, when writing to the cache
  // we key off of the query string
  async storeInCache(
    options: QueryResultCacheOptions,
    _savedCache: QueryResultCacheOptions | undefined,
    _queryRunner?: QueryRunner,
  ) {
    if (!this.cls.isActive()) {
      return;
    }

    if (options.query) {
      const cache = this.cls.get('queryCache');
      cache.set(options.query, options);
    }
  }

  // This function is intended to tell TypeORM if a ttl
  // has expired for the query, but we aren't using
  // time-based cache expiration
  isExpired(_savedCache: QueryResultCacheOptions): boolean {
    return false;
  }

  // We invoke this function to expire the entire cache on writes
  async clear(_queryRunner?: QueryRunner) {
    if (!this.cls.isActive()) {
      return;
    }

    const cache = this.cls.get('queryCache');
    cache.clear();
  }

  // This method is used by TypeORM's CLI for manually
  // interacting with the cache. We don't need it but
  // implement it for completeness
  async remove(identifiers: string[], _queryRunner?: QueryRunner) {
    if (!this.cls.isActive()) {
      return;
    }
    
    const cache = this.cls.get('queryCache');
    identifiers.forEach((identifier) => {
      cache.delete(identifier);
    });
  }
}

The query cache map gets instantiated when we define our root app’s module imports -

//file: src/app.module.ts
import { ClsModule } from 'nestjs-cls';
import { QueryResultCacheOptions } from 'typeorm/cache/QueryResultCacheOptions';
//... a whole bunch of other imports

// Relevant CLS configuration
@Module({
  imports: [
    ...,
    ClsModule.forRoot({
      global: true,
      middleware: {
        mount: true,
        setup: (cls) => {
          cls.set('queryCache', new Map<string, QueryResultCacheOptions>());
        },
      },
    }),
    ...,
  ],
  providers: [...]
})
export class AppModule implements NestModule {
  constructor() {
    initializeTransactionalContext();
  }
  //... other configuration and setup
}

Now that we have the cache set up, we need to add it to our TypeORM configuration. We also use the alwaysEnabled parameter to instruct TypeORM to always check the cache when executing queries -

//file: src/.../database-config.ts
import { ObjectLiteral } from 'typeorm';
import { ClsQueryResultCache } from '../database';
//... a whole bunch of other imports

interface MakeDatabaseTypeOrmOptionsOptions {
  ...,
  // need to set the cache option to false
  // for connections that don't need it
  cache: boolean | ObjectLiteral;
}

@Injectable()
export class DatabaseConfig {
  constructor(
    private readonly queryCache: ClsQueryResultCache,
  ) {}

  get mainDatabaseTypeOrmOptions(): TypeOrmModuleOptions {
    const cacheProvider = () => this.queryCache;

    return this.makeDatabaseTypeOrmOptions({
      ...,
      cache: {
        provider: cacheProvider,
        // undocumented TypeORM configuration that
        // eliminates the need to add cache option
        // to every query builder
        alwaysEnabled: true,
      },
    });
  }
  
  get otherDatabaseTypeOrmOptions(): TypeOrmModuleOptions {
    return this.makeDatabaseTypeOrmOptions({
      ...,
      cache: false,
    });
  }
  
  //... other stuff
}

Finally, we extended our custom TypeORM repository to include overrides that cleared the cache on writes. Here’s a sample -

//file: src/.../custom-typeorm-repository.ts
//... bunch of imports

export class CustomOrmRepository<
  Entity extends ObjectLiteral,
> extends Repository<Entity> {
  private readonly dataSource: DataSource;

  constructor(entity: EntityTarget<Entity>, dataSource: DataSource) {
    super(entity, dataSource.createEntityManager());
    this.dataSource = dataSource;
  }

  override save<T extends DeepPartial<Entity>>(
    entities: T[],
    options: SaveOptions & { reload: false },
  ): Promise<T[]>;
  override save<T extends DeepPartial<Entity>>(
    entities: T[],
    options?: SaveOptions,
  ): Promise<(T & Entity)[]>;
  override save<T extends DeepPartial<Entity>>(
    entity: T,
    options: SaveOptions & { reload: false },
  ): Promise<T>;
  override save<T extends DeepPartial<Entity>>(
    entity: T,
    options?: SaveOptions,
  ): Promise<T & Entity>;
  override save<T extends DeepPartial<Entity>>(
    entityOrEntities: T | T[],
    options?: SaveOptions & { reload?: boolean },
  ): Promise<T | T[] | (T & Entity) | (T & Entity)[]> {
    // Invalidate cache
    this.clearQueryResultCache();

    if (Array.isArray(entityOrEntities)) {
      return Promise.all(
        entityOrEntities.map((entity) => super.save(entity, options)),
      ) as Promise<(T & Entity)[]>;
    } else {
      return super.save(entityOrEntities, options) as Promise<T & Entity>;
    }
  }
  
  //... other overrides
  
  private clearQueryResultCache() {
    // If TypeORM is configured to use the query cache
    // clear the entire cache
    if (this.dataSource.queryResultCache) {
      this.dataSource.queryResultCache.clear();
    }
  }

Results

To test the effectiveness of our cache, we added some instrumentation to count how many cache hits we had, but since the cache was local to each request, so it was tough to measure across a browsing session. So instead we enabled the pg_stat_statements extension in our local development PostgreSQL database, and this would give us counts on the number of queries being run. We could reset the stats, click around in our web application, and get an accurate number of queries from the database directly.

Our testing showed that the cache was in fact very effective in reducing repeat queries. By navigating through the same pages with the query cache disabled and then with the cache enabled, we saw a number of queries executed reduced by 30 percent!

Conclusion

This post shows that with a couple hundred LOC, you can make sizable impact to your application performance, without wide refactoring or making risky changes. Setting up guard rails when scoping and adhering to them during implementation will force you to experiment until you find an optimal solution.

We are in a growth phase at WorkOS and the scale of load on our services is doubling every few months. Making precise and impactful optimizations without re-architecting our tech stack allows the bulk of our engineering team to focus on advancing our products. If you are interested in this kind of work, check out our open roles here.