[CORL-149] Persisted Queries (#2445)

* feat: enable persisted queries on the client

* fix: use `id` inside websocket message

* feat: initial server support for PQ

* feat: deeper server support

* feat: abstracted persisted query replacing logic

src/core/server/models/queries/cache.ts

@@ -0,0 +1,112 @@
+import DataLoader from "dataloader";
+import LRU from "lru-cache";
+import { Db } from "mongodb";
+
+import { loadPersistedQueries } from "coral-server/graph/tenant/persisted";
+import logger from "coral-server/logger";
+
+import { getQueries, PersistedQuery, primeQueries } from "./queries";
+
+interface PersistedQueryCacheOptions {
+  mongo: Db;
+}
+
+/**
+ * PersistedQueryCache abstracts the persisted query management.
+ */
+export class PersistedQueryCache {
+  private mongo: Db;
+  private queries: Map<string, PersistedQuery>;
+  private cache: LRU<string, PersistedQuery>;
+  private loader: DataLoader<string, PersistedQuery | null>;
+
+  constructor(options: PersistedQueryCacheOptions) {
+    const queries = loadPersistedQueries();
+
+    this.mongo = options.mongo;
+    this.loader = new DataLoader(
+      (ids: string[]) => getQueries(this.mongo, ids),
+      {
+        // Turn off caching as we're using the LRU cache here instead.
+        cache: false,
+      }
+    );
+    this.queries = new Map();
+    this.cache = new LRU({
+      // We'll only retain the amount of queries we have right now so we could
+      // possibly hold the previous version in memory if need be. Ideally, we'll
+      // always have the keys we need in memory.
+      max: queries.length,
+      dispose: (id, query) => {
+        logger.warn(
+          { queryID: id, queryVersion: query.version },
+          "cache full, dropping query from cache"
+        );
+      },
+    });
+
+    // Insert all the queries into the local query cache.
+    for (const query of queries) {
+      this.queries.set(query.id, query);
+    }
+  }
+
+  public get size() {
+    return this.queries.size + this.cache.length;
+  }
+
+  /**
+   * prime will load the local queries into the database so every time that the
+   * server starts, the queries will be available to other instances.
+   */
+  public async prime() {
+    if (this.queries.size === 0) {
+      return;
+    }
+
+    const queries: PersistedQuery[] = [];
+    for (const query of this.queries.values()) {
+      queries.push(query);
+    }
+
+    logger.debug({ queries: queries.length }, "priming queries");
+
+    await primeQueries(this.mongo, queries);
+  }
+
+  /**
+   * get will retrieve a given PersistedQuery by ID.
+   *
+   * @param id the ID of the persisted query to load
+   */
+  public async get(id: string) {
+    // Try to get the query from the local query cache.
+    let query: PersistedQuery | null | undefined = this.queries.get(id);
+    if (query) {
+      return query;
+    }
+
+    // Try to get the query from the remote cache.
+    query = this.cache.get(id);
+    if (query) {
+      return query;
+    }
+
+    // Try to get the query from the loader.
+    query = await this.loader.load(id);
+    if (query) {
+      logger.warn(
+        { queryID: id, queryVersion: query.version },
+        "query did not exist in cache, retrieved from MongoDB"
+      );
+
+      // Cache this query in the memory cache.
+      this.cache.set(query.id, query);
+      return query;
+    }
+
+    logger.warn({ queryID: id }, "query did not exist in cache or MongoDB");
+
+    return null;
+  }
+}

src/core/server/models/queries/index.ts

@@ -0,0 +1,2 @@
+export * from "./cache";
+export * from "./queries";

src/core/server/models/queries/queries.ts

@@ -0,0 +1,48 @@
+import { Db } from "mongodb";
+
+import { createIndexFactory } from "../helpers/indexing";
+
+function collection(mongo: Db) {
+  return mongo.collection<Readonly<PersistedQuery>>("queries");
+}
+
+export interface PersistedQuery {
+  id: string;
+  operation: string;
+  operationName: string;
+  query: string;
+  bundle: string;
+  version: string;
+}
+
+export async function createQueriesIndexes(mongo: Db) {
+  const createIndex = createIndexFactory(collection(mongo));
+
+  // UNIQUE { id }
+  await createIndex({ id: 1 }, { unique: true });
+}
+
+export async function primeQueries(mongo: Db, queries: PersistedQuery[]) {
+  // Setup persisting these queries.
+  const bulk = collection(mongo).initializeUnorderedBulkOp({});
+
+  // Upsert each query.
+  for (const query of queries) {
+    const { id } = query;
+
+    // Add to the bulk operation for MongoDB.
+    bulk
+      .find({ id })
+      .upsert()
+      .replaceOne(query);
+  }
+
+  // Execute the bulk operations.
+  await bulk.execute();
+}
+
+export async function getQueries(mongo: Db, ids: string[]) {
+  const cursor = await collection(mongo).find({ id: { $in: ids } });
+  const queries = await cursor.toArray();
+  return ids.map(id => queries.find(query => query.id === id) || null);
+}