add query method w caching

Author: andrejpk

src/cache.ts

@@ -48,9 +48,11 @@ export interface CachedMethods<DType> {
     args?: FindArgs
   ) => Promise<(DType | undefined)[]>;
   deleteFromCacheById: (id: string) => Promise<void>;
+  dataLoader?: DataLoader<string, DType, string>;
+  primeCache: (item: DType, ttl?: number) => void;
 }
 
-export const createCachingMethods = <DType>({
+export const createCachingMethods = <DType extends { id: string }>({
   container,
   cache,
   options,
@@ -64,10 +66,9 @@ export const createCachingMethods = <DType>({
     //   id,
     // }));
     // const response = await container.items.bulk(operations);
+    const idList = ids.map((id) => `'${id}'`).join(",");
     const querySpec = {
-      query: `select * from c where c.id in (${ids
-        .map((id) => `'${id}'`)
-        .join(",")})`,
+      query: `select * from c where c.id in (${idList})`,
       //   query: "select * from c where c.id in (@ids)",
       //   parameters: [{ name: "@ids", value: ids }],
     };
@@ -115,6 +116,19 @@ export const createCachingMethods = <DType>({
       loader.clear(id);
       await cache.delete(cachePrefix + id);
     },
+    /**
+     * Loads an item into DataLoader and optionally the cache (if TTL is specified)
+     * Use this when running a query outside of the findOneById/findManyByIds methos
+     * that automatically and transparently do this
+     */
+    primeCache: (doc: DType, ttl?: number) => {
+      const key = doc.id;
+      loader.prime(key, doc);
+      if (ttl) {
+        cache.set(key, EJSON.stringify(doc), { ttl });
+      }
+    },
+    dataLoader: loader,
   };
 
   return methods;

src/datasource.ts

@@ -1,11 +1,12 @@
 import { DataSource } from "apollo-datasource";
 import { ApolloError } from "apollo-server-errors";
 import { InMemoryLRUCache, KeyValueCache } from "apollo-server-caching";
-import { Container } from "@azure/cosmos";
+import { Container, SqlQuerySpec } from "@azure/cosmos";
 import { Logger } from "./helpers";
 
 import { isCosmosDbContainer } from "./helpers";
-import { createCachingMethods, CachedMethods } from "./cache";
+import { createCachingMethods, CachedMethods, FindArgs } from "./cache";
+import DataLoader from "dataloader";
 
 export interface CosmosDataSourceOptions {
   logger?: Logger;
@@ -15,7 +16,7 @@ const placeholderHandler = () => {
   throw new Error("DataSource not initialized");
 };
 
-export class CosmosDataSource<TData, TContext = any>
+export class CosmosDataSource<TData extends { id: string }, TContext = any>
   extends DataSource<TContext>
   implements CachedMethods<TData> {
   container: Container;
@@ -26,6 +27,20 @@ export class CosmosDataSource<TData, TContext = any>
   findOneById: CachedMethods<TData>["findOneById"] = placeholderHandler;
   findManyByIds: CachedMethods<TData>["findManyByIds"] = placeholderHandler;
   deleteFromCacheById: CachedMethods<TData>["deleteFromCacheById"] = placeholderHandler;
+  dataLoader: CachedMethods<TData>["dataLoader"];
+  primeCache: CachedMethods<TData>["primeCache"] = placeholderHandler;
+
+  // base SQL operations aren't cached so directly declared here
+  async findManyByQuery(query: SqlQuerySpec | string, { ttl }: FindArgs) {
+    const results = await this.container.items.query<TData>(query).fetchAll();
+    // prime these into the dataloader and maybe the cache
+    if (this.dataLoader && results.resources) {
+      results.resources.forEach((r) => {
+        this.primeCache(r, ttl);
+      });
+    }
+    return results.resources;
+  }
 
   constructor(container: Container, options: CosmosDataSourceOptions = {}) {
     super();