-
Notifications
You must be signed in to change notification settings - Fork 1.8k
/
estimated_document_count.ts
83 lines (70 loc) · 2.42 KB
/
estimated_document_count.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
import type { Document } from '../bson';
import type { Collection } from '../collection';
import type { MongoServerError } from '../error';
import type { Server } from '../sdam/server';
import type { ClientSession } from '../sessions';
import { Callback, maxWireVersion } from '../utils';
import { CommandOperation, CommandOperationOptions } from './command';
import { Aspect, defineAspects } from './operation';
/** @public */
export interface EstimatedDocumentCountOptions extends CommandOperationOptions {
/**
* The maximum amount of time to allow the operation to run.
*
* This option is sent only if the caller explicitly provides a value. The default is to not send a value.
*/
maxTimeMS?: number;
}
/** @internal */
export class EstimatedDocumentCountOperation extends CommandOperation<number> {
override options: EstimatedDocumentCountOptions;
collectionName: string;
constructor(collection: Collection, options: EstimatedDocumentCountOptions = {}) {
super(collection, options);
this.options = options;
this.collectionName = collection.collectionName;
}
override execute(
server: Server,
session: ClientSession | undefined,
callback: Callback<number>
): void {
if (maxWireVersion(server) < 12) {
return this.executeLegacy(server, session, callback);
}
const pipeline = [{ $collStats: { count: {} } }, { $group: { _id: 1, n: { $sum: '$count' } } }];
const cmd: Document = { aggregate: this.collectionName, pipeline, cursor: {} };
if (typeof this.options.maxTimeMS === 'number') {
cmd.maxTimeMS = this.options.maxTimeMS;
}
super.executeCommand(server, session, cmd, (err, response) => {
if (err && (err as MongoServerError).code !== 26) {
callback(err);
return;
}
callback(undefined, response?.cursor?.firstBatch[0]?.n || 0);
});
}
executeLegacy(
server: Server,
session: ClientSession | undefined,
callback: Callback<number>
): void {
const cmd: Document = { count: this.collectionName };
if (typeof this.options.maxTimeMS === 'number') {
cmd.maxTimeMS = this.options.maxTimeMS;
}
super.executeCommand(server, session, cmd, (err, response) => {
if (err) {
callback(err);
return;
}
callback(undefined, response.n || 0);
});
}
}
defineAspects(EstimatedDocumentCountOperation, [
Aspect.READ_OPERATION,
Aspect.RETRYABLE,
Aspect.CURSOR_CREATING
]);