Skip to content

Commit b79345e

Browse files
swantzterswantzter
committed
Use DataLoader's maxBatchSize feature rather than hand-rolling one
1 parent 07ee2aa commit b79345e

File tree

1 file changed

+124
-133
lines changed

1 file changed

+124
-133
lines changed

‎src/cache.ts‎

Lines changed: 124 additions & 133 deletions
Original file line numberDiff line numberDiff line change
@@ -1,133 +1,124 @@
1-
import { CollectionReference, FieldPath, QuerySnapshot } from '@google-cloud/firestore'
2-
import { KeyValueCache } from 'apollo-server-caching'
3-
import DataLoader from 'dataloader'
4-
import { FirestoreDataSourceOptions } from './datasource'
5-
import { replacer, reviverFactory } from './helpers'
6-
7-
// https://github.com/graphql/dataloader#batch-function
8-
const orderDocs = <V>(ids: readonly string[]) => (
9-
docs: Array<V | undefined>,
10-
keyFn?: (source: V) => string
11-
) => {
12-
const keyFnDef =
13-
keyFn ??
14-
((source: V & { id?: string }) => {
15-
if (source.id) return source.id
16-
throw new Error(
17-
'Could not find ID for object; if using an alternate key, pass in a key function'
18-
)
19-
})
20-
21-
const checkNotUndefined = (input: V | undefined): input is V => {
22-
return Boolean(input)
23-
}
24-
25-
const idMap: Record<string, V> = docs
26-
.filter(checkNotUndefined)
27-
.reduce((prev: Record<string, V>, cur: V) => {
28-
prev[keyFnDef(cur)] = cur
29-
return prev
30-
}, {})
31-
return ids.map((id) => idMap[id])
32-
}
33-
34-
export interface createCatchingMethodArgs<DType> {
35-
collection: CollectionReference<DType>
36-
cache: KeyValueCache
37-
options: FirestoreDataSourceOptions
38-
}
39-
40-
export interface FindArgs {
41-
ttl?: number
42-
}
43-
44-
export interface CachedMethods<DType> {
45-
findOneById: (id: string, args?: FindArgs) => Promise<DType | undefined>
46-
findManyByIds: (
47-
ids: string[],
48-
args?: FindArgs
49-
) => Promise<Array<DType | undefined>>
50-
deleteFromCacheById: (id: string) => Promise<void>
51-
dataLoader?: DataLoader<string, DType, string>
52-
cache?: KeyValueCache
53-
cachePrefix?: string
54-
primeLoader: (item: DType | DType[], ttl?: number) => void
55-
}
56-
57-
export const createCachingMethods = <DType extends { id: string }>({
58-
collection,
59-
cache,
60-
options
61-
}: createCatchingMethodArgs<DType>): CachedMethods<DType> => {
62-
const loader = new DataLoader<string, DType>(async (ids) => {
63-
const qSnaps: Array<Promise<QuerySnapshot<DType>>> = []
64-
// the 'in' operator only supports up to 10 values at a time
65-
for (let idx = 0; idx < ids.length; idx += 10) {
66-
const idSlice = ids.slice(idx, idx + 10)
67-
options?.logger?.debug(
68-
`FirestoreDataSource/DataLoader: loading for IDs: ${idSlice}`
69-
)
70-
qSnaps.push(collection.where(FieldPath.documentId(), 'in', idSlice).get())
71-
}
72-
const documents = (await Promise.all(qSnaps)).flatMap(qSnap => qSnap.docs).map(dSnap => dSnap.exists ? dSnap.data() : undefined)
73-
options?.logger?.debug(
74-
`FirestoreDataSource/DataLoader: response count: ${documents.length}`
75-
)
76-
77-
return orderDocs<DType>(ids)(documents)
78-
})
79-
80-
const reviver = reviverFactory(collection)
81-
82-
const cachePrefix = `firestore-${collection.path}-`
83-
84-
const methods: CachedMethods<DType> = {
85-
findOneById: async (id, { ttl } = {}) => {
86-
options?.logger?.debug(`FirestoreDataSource: Running query for ID ${id}`)
87-
const key = cachePrefix + id
88-
89-
const cacheDoc = await cache.get(key)
90-
if (cacheDoc) {
91-
return JSON.parse(cacheDoc, reviver) as DType
92-
}
93-
94-
const doc = await loader.load(id)
95-
96-
if (Number.isInteger(ttl)) {
97-
await cache.set(key, JSON.stringify(doc, replacer), { ttl })
98-
}
99-
100-
return doc
101-
},
102-
103-
findManyByIds: async (ids, args = {}) => {
104-
options?.logger?.debug(`FirestoreDataSource: Running query for IDs ${ids}`)
105-
return await Promise.all(ids.map(async (id) => await methods.findOneById(id, args)))
106-
},
107-
108-
deleteFromCacheById: async (id) => {
109-
loader.clear(id)
110-
await cache.delete(cachePrefix + id)
111-
},
112-
/**
113-
* Loads an item or items into DataLoader and optionally the cache (if TTL is specified)
114-
* Use this when running a query outside of the findOneById/findManyByIds methos
115-
* that automatically and transparently do this
116-
*/
117-
primeLoader: async (docs, ttl?: number) => {
118-
docs = Array.isArray(docs) ? docs : [docs]
119-
for (const doc of docs) {
120-
loader.prime(doc.id, doc)
121-
const key = cachePrefix + doc.id
122-
if (!!ttl || !!(await cache.get(key))) {
123-
await cache.set(key, JSON.stringify(doc, replacer), { ttl })
124-
}
125-
}
126-
},
127-
dataLoader: loader,
128-
cache,
129-
cachePrefix
130-
}
131-
132-
return methods
133-
}
1+
import { CollectionReference, FieldPath } from '@google-cloud/firestore'
2+
import { KeyValueCache } from 'apollo-server-caching'
3+
import DataLoader from 'dataloader'
4+
import { FirestoreDataSourceOptions } from './datasource'
5+
import { replacer, reviverFactory } from './helpers'
6+
7+
// https://github.com/graphql/dataloader#batch-function
8+
const orderDocs = <V>(ids: readonly string[]) => (
9+
docs: Array<V | undefined>,
10+
keyFn?: (source: V) => string
11+
) => {
12+
const keyFnDef =
13+
keyFn ??
14+
((source: V & { id?: string }) => {
15+
if (source.id) return source.id
16+
throw new Error(
17+
'Could not find ID for object; if using an alternate key, pass in a key function'
18+
)
19+
})
20+
21+
const checkNotUndefined = (input: V | undefined): input is V => {
22+
return Boolean(input)
23+
}
24+
25+
const idMap: Record<string, V> = docs
26+
.filter(checkNotUndefined)
27+
.reduce((prev: Record<string, V>, cur: V) => {
28+
prev[keyFnDef(cur)] = cur
29+
return prev
30+
}, {})
31+
return ids.map((id) => idMap[id])
32+
}
33+
34+
export interface createCatchingMethodArgs<DType> {
35+
collection: CollectionReference<DType>
36+
cache: KeyValueCache
37+
options: FirestoreDataSourceOptions
38+
}
39+
40+
export interface FindArgs {
41+
ttl?: number
42+
}
43+
44+
export interface CachedMethods<DType> {
45+
findOneById: (id: string, args?: FindArgs) => Promise<DType | undefined>
46+
findManyByIds: (
47+
ids: string[],
48+
args?: FindArgs
49+
) => Promise<Array<DType | undefined>>
50+
deleteFromCacheById: (id: string) => Promise<void>
51+
dataLoader?: DataLoader<string, DType, string>
52+
cache?: KeyValueCache
53+
cachePrefix?: string
54+
primeLoader: (item: DType | DType[], ttl?: number) => void
55+
}
56+
57+
export const createCachingMethods = <DType extends { id: string }>({
58+
collection,
59+
cache,
60+
options
61+
}: createCatchingMethodArgs<DType>): CachedMethods<DType> => {
62+
const loader = new DataLoader<string, DType>(async (ids) => {
63+
options?.logger?.debug(`FirestoreDataSource/DataLoader: loading for IDs: ${ids}`)
64+
const qSnap = await collection.where(FieldPath.documentId(), 'in', ids).get()
65+
const documents = qSnap.docs.map(dSnap => dSnap.exists ? dSnap.data() : undefined)
66+
options?.logger?.debug(`FirestoreDataSource/DataLoader: response count: ${documents.length}`)
67+
68+
return orderDocs<DType>(ids)(documents)
69+
}, { maxBatchSize: 10 })
70+
71+
const reviver = reviverFactory(collection)
72+
73+
const cachePrefix = `firestore-${collection.path}-`
74+
75+
const methods: CachedMethods<DType> = {
76+
findOneById: async (id, { ttl } = {}) => {
77+
options?.logger?.debug(`FirestoreDataSource: Running query for ID ${id}`)
78+
const key = cachePrefix + id
79+
80+
const cacheDoc = await cache.get(key)
81+
if (cacheDoc) {
82+
return JSON.parse(cacheDoc, reviver) as DType
83+
}
84+
85+
const doc = await loader.load(id)
86+
87+
if (Number.isInteger(ttl)) {
88+
await cache.set(key, JSON.stringify(doc, replacer), { ttl })
89+
}
90+
91+
return doc
92+
},
93+
94+
findManyByIds: async (ids, args = {}) => {
95+
options?.logger?.debug(`FirestoreDataSource: Running query for IDs ${ids}`)
96+
return await Promise.all(ids.map(async (id) => await methods.findOneById(id, args)))
97+
},
98+
99+
deleteFromCacheById: async (id) => {
100+
loader.clear(id)
101+
await cache.delete(cachePrefix + id)
102+
},
103+
/**
104+
* Loads an item or items into DataLoader and optionally the cache (if TTL is specified)
105+
* Use this when running a query outside of the findOneById/findManyByIds methos
106+
* that automatically and transparently do this
107+
*/
108+
primeLoader: async (docs, ttl?: number) => {
109+
docs = Array.isArray(docs) ? docs : [docs]
110+
for (const doc of docs) {
111+
loader.prime(doc.id, doc)
112+
const key = cachePrefix + doc.id
113+
if (!!ttl || !!(await cache.get(key))) {
114+
await cache.set(key, JSON.stringify(doc, replacer), { ttl })
115+
}
116+
}
117+
},
118+
dataLoader: loader,
119+
cache,
120+
cachePrefix
121+
}
122+
123+
return methods
124+
}

0 commit comments

Comments
 (0)