Skip to content

Commit 7d808c1

Browse files
committed
Update query-document mapping from bundles.
1 parent 15b96a9 commit 7d808c1

File tree

6 files changed

+198
-30
lines changed

6 files changed

+198
-30
lines changed

packages/firestore/src/core/bundle.ts

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -218,16 +218,17 @@ export class BundleLoader {
218218
'Bundled documents ends with a document metadata and missing document.'
219219
);
220220

221+
const result = await applyBundleDocuments(this.localStore, this.documents);
222+
221223
for (const q of this.queries) {
222-
await saveNamedQuery(this.localStore, q);
224+
await saveNamedQuery(
225+
this.localStore,
226+
q,
227+
result.queryDocumentMap.get(q.name!)
228+
);
223229
}
224230

225-
const changedDocs = await applyBundleDocuments(
226-
this.localStore,
227-
this.documents
228-
);
229-
230231
this.progress.taskState = 'Success';
231-
return new BundleLoadResult({ ...this.progress }, changedDocs);
232+
return new BundleLoadResult({ ...this.progress }, result.changedDocuments);
232233
}
233234
}

packages/firestore/src/local/local_store.ts

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1249,6 +1249,11 @@ export async function ignoreIfPrimaryLeaseLoss(
12491249
}
12501250
}
12511251

1252+
export interface ApplyBundleDocumentsResult {
1253+
changedDocuments: MaybeDocumentMap;
1254+
queryDocumentMap: Map<string, DocumentKeySet>;
1255+
}
1256+
12521257
/**
12531258
* Applies the documents from a bundle to the "ground-state" (remote)
12541259
* documents.
@@ -1259,11 +1264,12 @@ export async function ignoreIfPrimaryLeaseLoss(
12591264
export function applyBundleDocuments(
12601265
localStore: LocalStore,
12611266
documents: BundledDocuments
1262-
): Promise<MaybeDocumentMap> {
1267+
): Promise<ApplyBundleDocumentsResult> {
12631268
const localStoreImpl = debugCast(localStore, LocalStoreImpl);
12641269
const bundleConverter = new BundleConverter(localStoreImpl.serializer);
12651270
let documentMap = maybeDocumentMap();
12661271
let versionMap = documentVersionMap();
1272+
const queryDocumentMap: Map<string, DocumentKeySet> = new Map();
12671273
for (const bundleDoc of documents) {
12681274
const documentKey = bundleConverter.toDocumentKey(bundleDoc.metadata.name!);
12691275
documentMap = documentMap.insert(
@@ -1274,6 +1280,14 @@ export function applyBundleDocuments(
12741280
documentKey,
12751281
bundleConverter.toSnapshotVersion(bundleDoc.metadata.readTime!)
12761282
);
1283+
if (bundleDoc.metadata.queries) {
1284+
for (const queryName of bundleDoc.metadata.queries) {
1285+
const documentKeys = (
1286+
queryDocumentMap.get(queryName) || documentKeySet()
1287+
).add(documentKey);
1288+
queryDocumentMap.set(queryName, documentKeys);
1289+
}
1290+
}
12771291
}
12781292

12791293
const documentBuffer = localStoreImpl.remoteDocuments.newChangeBuffer({
@@ -1300,6 +1314,12 @@ export function applyBundleDocuments(
13001314
txn,
13011315
changedDocs
13021316
);
1317+
})
1318+
.next(changedDocuments => {
1319+
return PersistencePromise.resolve({
1320+
changedDocuments,
1321+
queryDocumentMap
1322+
});
13031323
});
13041324
}
13051325
);
@@ -1373,7 +1393,8 @@ export function getNamedQuery(
13731393
*/
13741394
export async function saveNamedQuery(
13751395
localStore: LocalStore,
1376-
query: bundleProto.NamedQuery
1396+
query: bundleProto.NamedQuery,
1397+
documents: DocumentKeySet = documentKeySet()
13771398
): Promise<void> {
13781399
// Allocate a target for the named query such that it can be resumed
13791400
// from associated read time if users use it to listen.
@@ -1400,10 +1421,22 @@ export async function saveNamedQuery(
14001421
transaction,
14011422
newTargetData
14021423
);
1424+
localStoreImpl.targetDataByTarget = localStoreImpl.targetDataByTarget.insert(
1425+
newTargetData.targetId,
1426+
newTargetData
1427+
);
14031428
}
1404-
return updateReadTime.next(() =>
1405-
localStoreImpl.bundleCache.saveNamedQuery(transaction, query)
1406-
);
1429+
return updateReadTime
1430+
.next(() =>
1431+
localStoreImpl.bundleCache.saveNamedQuery(transaction, query)
1432+
)
1433+
.next(() =>
1434+
localStoreImpl.targetCache.addMatchingKeys(
1435+
transaction,
1436+
documents,
1437+
allocated.targetId
1438+
)
1439+
);
14071440
}
14081441
);
14091442
}

packages/firestore/src/protos/firestore/bundle.proto

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,9 @@ message BundledDocumentMetadata {
7979

8080
// Whether the document exists.
8181
bool exists = 3;
82+
83+
// The names of the queries in this bundle that this document matches to.
84+
repeated string queries = 4;
8285
}
8386

8487
// Metadata describing the bundle file/stream.

packages/firestore/src/protos/firestore_bundle_proto.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,9 @@ export interface BundledDocumentMetadata {
5454

5555
/** BundledDocumentMetadata exists */
5656
exists?: boolean | null;
57+
58+
/** The names of the queries in this bundle that this document matches to. */
59+
queries?: string[] | null;
5760
}
5861

5962
/** Properties of a BundleMetadata. */

packages/firestore/test/unit/local/local_store.test.ts

Lines changed: 120 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ import { IndexFreeQueryEngine } from '../../../src/local/index_free_query_engine
3535
import { IndexedDbPersistence } from '../../../src/local/indexeddb_persistence';
3636
import {
3737
applyBundleDocuments,
38+
ApplyBundleDocumentsResult,
3839
getNamedQuery,
3940
hasNewerBundle,
4041
LocalStore,
@@ -48,6 +49,7 @@ import { LocalViewChanges } from '../../../src/local/local_view_changes';
4849
import { Persistence } from '../../../src/local/persistence';
4950
import { SimpleQueryEngine } from '../../../src/local/simple_query_engine';
5051
import {
52+
DocumentKeySet,
5153
documentKeySet,
5254
MaybeDocumentMap
5355
} from '../../../src/model/collections';
@@ -91,6 +93,7 @@ import {
9193
query,
9294
setMutation,
9395
TestBundledDocuments,
96+
TestNamedQuery,
9497
TestSnapshotVersion,
9598
transformMutation,
9699
unknownDoc,
@@ -113,6 +116,7 @@ export interface LocalStoreComponents {
113116
class LocalStoreTester {
114117
private promiseChain: Promise<void> = Promise.resolve();
115118
private lastChanges: MaybeDocumentMap | null = null;
119+
private lastBundleQueryDocumentMap: Map<string, DocumentKeySet> | null = null;
116120
private lastTargetId: TargetId | null = null;
117121
private batches: MutationBatch[] = [];
118122

@@ -125,6 +129,7 @@ class LocalStoreTester {
125129
private prepareNextStep(): void {
126130
this.promiseChain = this.promiseChain.then(() => {
127131
this.lastChanges = null;
132+
this.lastBundleQueryDocumentMap = null;
128133
this.lastTargetId = null;
129134
this.queryEngine.resetCounts();
130135
});
@@ -137,7 +142,7 @@ class LocalStoreTester {
137142
| RemoteEvent
138143
| LocalViewChanges
139144
| TestBundledDocuments
140-
| bundleProto.NamedQuery
145+
| TestNamedQuery
141146
): LocalStoreTester {
142147
if (op instanceof Mutation) {
143148
return this.afterMutations([op]);
@@ -188,17 +193,22 @@ class LocalStoreTester {
188193

189194
this.promiseChain = this.promiseChain
190195
.then(() => applyBundleDocuments(this.localStore, documents))
191-
.then((result: MaybeDocumentMap) => {
192-
this.lastChanges = result;
196+
.then((result: ApplyBundleDocumentsResult) => {
197+
this.lastChanges = result.changedDocuments;
198+
this.lastBundleQueryDocumentMap = result.queryDocumentMap;
193199
});
194200
return this;
195201
}
196202

197-
afterNamedQuery(namedQuery: bundleProto.NamedQuery): LocalStoreTester {
203+
afterNamedQuery(testQuery: TestNamedQuery): LocalStoreTester {
198204
this.prepareNextStep();
199205

200206
this.promiseChain = this.promiseChain.then(() =>
201-
saveNamedQuery(this.localStore, namedQuery)
207+
saveNamedQuery(
208+
this.localStore,
209+
testQuery.namedQuery,
210+
testQuery.matchingDocuments
211+
)
202212
);
203213
return this;
204214
}
@@ -372,6 +382,25 @@ class LocalStoreTester {
372382
return this;
373383
}
374384

385+
toReturnQueryDocumentMap(
386+
expected: Map<string, DocumentKeySet>
387+
): LocalStoreTester {
388+
this.promiseChain = this.promiseChain.then(() => {
389+
debugAssert(
390+
!!this.lastBundleQueryDocumentMap,
391+
'Expecting query document map to be returned from applying bundled documents'
392+
);
393+
for (const k of Array.from(this.lastBundleQueryDocumentMap.keys())) {
394+
expect(expected.has(k)).to.be.true;
395+
expect(
396+
this.lastBundleQueryDocumentMap.get(k)!.isEqual(expected.get(k)!)
397+
).to.be.true;
398+
}
399+
this.lastBundleQueryDocumentMap = null;
400+
});
401+
return this;
402+
}
403+
375404
toReturnRemoved(...keyStrings: string[]): LocalStoreTester {
376405
this.promiseChain = this.promiseChain.then(() => {
377406
debugAssert(
@@ -432,6 +461,29 @@ class LocalStoreTester {
432461
return this;
433462
}
434463

464+
toHaveQueryDocumentMapping(
465+
persistence: Persistence,
466+
targetId: TargetId,
467+
expectedKeys: DocumentKeySet
468+
): LocalStoreTester {
469+
this.promiseChain = this.promiseChain.then(() => {
470+
return persistence.runTransaction(
471+
'toHaveQueryDocumentMapping',
472+
'readonly',
473+
transaction => {
474+
return persistence
475+
.getTargetCache()
476+
.getMatchingKeysForTargetId(transaction, targetId)
477+
.next(matchedKeys => {
478+
expect(matchedKeys.isEqual(expectedKeys)).to.be.true;
479+
});
480+
}
481+
);
482+
});
483+
484+
return this;
485+
}
486+
435487
toHaveNewerBundle(
436488
metadata: bundleProto.BundleMetadata,
437489
expected: boolean
@@ -1706,6 +1758,69 @@ function genericLocalStoreTests(
17061758
.finish();
17071759
});
17081760

1761+
it('handles loading named queries allocates targets and updates target document mapping', async () => {
1762+
const expectedQueryDocumentMap = new Map([
1763+
['query-1', documentKeySet(key('foo1/bar'))],
1764+
['query-2', documentKeySet(key('foo2/bar'))]
1765+
]);
1766+
const version1 = SnapshotVersion.fromTimestamp(Timestamp.fromMillis(10000));
1767+
const version2 = SnapshotVersion.fromTimestamp(Timestamp.fromMillis(20000));
1768+
1769+
return expectLocalStore()
1770+
.after(
1771+
bundledDocuments(
1772+
[doc('foo1/bar', 1, { sum: 1337 }), doc('foo2/bar', 2, { sum: 42 })],
1773+
[['query-1'], ['query-2']]
1774+
)
1775+
)
1776+
.toReturnChanged(
1777+
doc('foo1/bar', 1, { sum: 1337 }),
1778+
doc('foo2/bar', 2, { sum: 42 })
1779+
)
1780+
.toReturnQueryDocumentMap(expectedQueryDocumentMap)
1781+
.toContain(doc('foo1/bar', 1, { sum: 1337 }))
1782+
.toContain(doc('foo2/bar', 2, { sum: 42 }))
1783+
.after(
1784+
namedQuery(
1785+
'query-1',
1786+
query('foo1'),
1787+
/* limitType */ 'FIRST',
1788+
version1,
1789+
expectedQueryDocumentMap.get('query-1')
1790+
)
1791+
)
1792+
.toHaveNamedQuery({
1793+
name: 'query-1',
1794+
query: query('foo1'),
1795+
readTime: version1
1796+
})
1797+
.toHaveQueryDocumentMapping(
1798+
persistence,
1799+
/*targetId*/ 2,
1800+
/*expectedKeys*/ documentKeySet(key('foo1/bar'))
1801+
)
1802+
.after(
1803+
namedQuery(
1804+
'query-2',
1805+
query('foo2'),
1806+
/* limitType */ 'FIRST',
1807+
version2,
1808+
expectedQueryDocumentMap.get('query-2')
1809+
)
1810+
)
1811+
.toHaveNamedQuery({
1812+
name: 'query-2',
1813+
query: query('foo2'),
1814+
readTime: version2
1815+
})
1816+
.toHaveQueryDocumentMapping(
1817+
persistence,
1818+
/*targetId*/ 4,
1819+
/*expectedKeys*/ documentKeySet(key('foo2/bar'))
1820+
)
1821+
.finish();
1822+
});
1823+
17091824
it('handles saving and loading limit to last queries', async () => {
17101825
const now = Timestamp.now();
17111826
return expectLocalStore()

packages/firestore/test/util/helpers.ts

Lines changed: 26 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -445,14 +445,16 @@ export class TestBundledDocuments {
445445
}
446446

447447
export function bundledDocuments(
448-
documents: MaybeDocument[]
448+
documents: MaybeDocument[],
449+
queryNames?: string[][]
449450
): TestBundledDocuments {
450-
const result = documents.map(d => {
451+
const result = documents.map((d, index) => {
451452
return {
452453
metadata: {
453454
name: toName(JSON_SERIALIZER, d.key),
454455
readTime: toVersion(JSON_SERIALIZER, d.version),
455-
exists: d instanceof Document
456+
exists: d instanceof Document,
457+
queries: queryNames ? queryNames[index] : undefined
456458
},
457459
document:
458460
d instanceof Document ? toDocument(JSON_SERIALIZER, d) : undefined
@@ -462,21 +464,32 @@ export function bundledDocuments(
462464
return new TestBundledDocuments(result);
463465
}
464466

467+
export class TestNamedQuery {
468+
constructor(
469+
public namedQuery: bundleProto.NamedQuery,
470+
public matchingDocuments: DocumentKeySet
471+
) {}
472+
}
473+
465474
export function namedQuery(
466475
name: string,
467476
query: Query,
468477
limitType: bundleProto.LimitType,
469-
readTime: SnapshotVersion
470-
): bundleProto.NamedQuery {
478+
readTime: SnapshotVersion,
479+
matchingDocuments: DocumentKeySet = documentKeySet()
480+
): TestNamedQuery {
471481
return {
472-
name,
473-
readTime: toTimestamp(JSON_SERIALIZER, readTime.toTimestamp()),
474-
bundledQuery: {
475-
parent: toQueryTarget(JSON_SERIALIZER, queryToTarget(query)).parent,
476-
limitType,
477-
structuredQuery: toQueryTarget(JSON_SERIALIZER, queryToTarget(query))
478-
.structuredQuery
479-
}
482+
namedQuery: {
483+
name,
484+
readTime: toTimestamp(JSON_SERIALIZER, readTime.toTimestamp()),
485+
bundledQuery: {
486+
parent: toQueryTarget(JSON_SERIALIZER, queryToTarget(query)).parent,
487+
limitType,
488+
structuredQuery: toQueryTarget(JSON_SERIALIZER, queryToTarget(query))
489+
.structuredQuery
490+
}
491+
},
492+
matchingDocuments
480493
};
481494
}
482495

0 commit comments

Comments
 (0)