refactor: move CountDocument logic into collection API (mongodb#4148)

Co-authored-by: Bailey Pearson <[email protected]>

Author: nbbeeken

src/collection.ts

@@ -25,7 +25,6 @@ import type {
 import type { AggregateOptions } from './operations/aggregate';
 import { BulkWriteOperation } from './operations/bulk_write';
 import { CountOperation, type CountOptions } from './operations/count';
-import { CountDocumentsOperation, type CountDocumentsOptions } from './operations/count_documents';
 import {
   DeleteManyOperation,
   DeleteOneOperation,
@@ -101,6 +100,14 @@ export interface ModifyResult<TSchema = Document> {
   ok: 0 | 1;
 }
 
+/** @public */
+export interface CountDocumentsOptions extends AggregateOptions {
+  /** The number of documents to skip. */
+  skip?: number;
+  /** The maximum amount of documents to consider. */
+  limit?: number;
+}
+
 /** @public */
 export interface CollectionOptions extends BSONSerializeOptions, WriteConcernOptions {
   /** Specify a read concern for the collection. (only MongoDB 3.2 or higher supported) */
@@ -764,10 +771,23 @@ export class Collection<TSchema extends Document = Document> {
     filter: Filter<TSchema> = {},
     options: CountDocumentsOptions = {}
   ): Promise<number> {
-    return await executeOperation(
-      this.client,
-      new CountDocumentsOperation(this as TODO_NODE_3286, filter, resolveOptions(this, options))
-    );
+    const pipeline = [];
+    pipeline.push({ $match: filter });
+
+    if (typeof options.skip === 'number') {
+      pipeline.push({ $skip: options.skip });
+    }
+
+    if (typeof options.limit === 'number') {
+      pipeline.push({ $limit: options.limit });
+    }
+
+    pipeline.push({ $group: { _id: 1, n: { $sum: 1 } } });
+
+    const cursor = this.aggregate<{ n: number }>(pipeline, options);
+    const doc = await cursor.next();
+    await cursor.close();
+    return doc?.n ?? 0;
   }
 
   /**

src/index.ts

@@ -305,7 +305,12 @@ export type {
   MongoDBResponse,
   MongoDBResponseConstructor
 } from './cmap/wire_protocol/responses';
-export type { CollectionOptions, CollectionPrivate, ModifyResult } from './collection';
+export type {
+  CollectionOptions,
+  CollectionPrivate,
+  CountDocumentsOptions,
+  ModifyResult
+} from './collection';
 export type {
   COMMAND_FAILED,
   COMMAND_STARTED,
@@ -463,7 +468,6 @@ export type {
   OperationParent
 } from './operations/command';
 export type { CountOptions } from './operations/count';
-export type { CountDocumentsOptions } from './operations/count_documents';
 export type {
   ClusteredCollectionOptions,
   CreateCollectionOptions,

src/operations/count_documents.ts

@@ -53,11 +53,6 @@ describe('abstract operation', function () {
         subclassType: mongodb.CountOperation,
         correctCommandName: 'count'
       },
-      {
-        subclassCreator: () => new mongodb.CountDocumentsOperation(collection, { a: 1 }, {}),
-        subclassType: mongodb.CountDocumentsOperation,
-        correctCommandName: 'aggregate'
-      },
       {
         subclassCreator: () => new mongodb.CreateCollectionOperation(db, 'name'),
         subclassType: mongodb.CreateCollectionOperation,
@@ -322,11 +317,7 @@ describe('abstract operation', function () {
           it(`operation.commandName equals key in command document`, async function () {
             const subclassInstance = subclassCreator();
             const yieldDoc =
-              subclassType.name === 'ProfilingLevelOperation'
-                ? { ok: 1, was: 1 }
-                : subclassType.name === 'CountDocumentsOperation'
-                ? { shift: () => ({ n: 1 }) }
-                : { ok: 1 };
+              subclassType.name === 'ProfilingLevelOperation' ? { ok: 1, was: 1 } : { ok: 1 };
             const cmdCallerStub = sinon.stub(Server.prototype, 'command').resolves(yieldDoc);
             if (sameServerOnlyOperationSubclasses.includes(subclassType.name.toString())) {
               await subclassInstance.execute(constructorServer, client.session);

test/integration/crud/abstract_operation.test.ts

@@ -156,6 +156,87 @@ describe('CRUD API', function () {
     });
   });
 
+  describe('countDocuments()', () => {
+    let client: MongoClient;
+    let events;
+    let collection: Collection<{ _id: number }>;
+
+    beforeEach(async function () {
+      client = this.configuration.newClient({ monitorCommands: true });
+      events = [];
+      client.on('commandSucceeded', commandSucceeded =>
+        commandSucceeded.commandName === 'aggregate' ? events.push(commandSucceeded) : null
+      );
+      client.on('commandFailed', commandFailed =>
+        commandFailed.commandName === 'aggregate' ? events.push(commandFailed) : null
+      );
+
+      collection = client.db('countDocuments').collection('countDocuments');
+      await collection.drop().catch(() => null);
+      await collection.insertMany([{ _id: 1 }, { _id: 2 }]);
+    });
+
+    afterEach(async () => {
+      await collection.drop().catch(() => null);
+      await client.close();
+    });
+
+    describe('when the aggregation operation succeeds', () => {
+      it('the cursor for countDocuments is closed', async function () {
+        const spy = sinon.spy(Collection.prototype, 'aggregate');
+        const result = await collection.countDocuments({});
+        expect(result).to.deep.equal(2);
+        expect(events[0]).to.be.instanceOf(CommandSucceededEvent);
+        expect(spy.returnValues[0]).to.have.property('closed', true);
+        expect(spy.returnValues[0]).to.have.nested.property('session.hasEnded', true);
+      });
+    });
+
+    describe('when the aggregation operation fails', () => {
+      beforeEach(async function () {
+        if (semver.lt(this.configuration.version, '4.2.0')) {
+          if (this.currentTest) {
+            this.currentTest.skipReason = `Cannot run fail points on server version: ${this.configuration.version}`;
+          }
+          this.skip();
+        }
+
+        const failPoint: FailPoint = {
+          configureFailPoint: 'failCommand',
+          mode: 'alwaysOn',
+          data: {
+            failCommands: ['aggregate'],
+            // 1 == InternalError, but this value not important to the test
+            errorCode: 1
+          }
+        };
+        await client.db().admin().command(failPoint);
+      });
+
+      afterEach(async function () {
+        if (semver.lt(this.configuration.version, '4.2.0')) {
+          return;
+        }
+
+        const failPoint: FailPoint = {
+          configureFailPoint: 'failCommand',
+          mode: 'off',
+          data: { failCommands: ['aggregate'] }
+        };
+        await client.db().admin().command(failPoint);
+      });
+
+      it('the cursor for countDocuments is closed', async function () {
+        const spy = sinon.spy(Collection.prototype, 'aggregate');
+        const error = await collection.countDocuments({}).catch(error => error);
+        expect(error).to.be.instanceOf(MongoServerError);
+        expect(events.at(0)).to.be.instanceOf(CommandFailedEvent);
+        expect(spy.returnValues.at(0)).to.have.property('closed', true);
+        expect(spy.returnValues.at(0)).to.have.nested.property('session.hasEnded', true);
+      });
+    });
+  });
+
   context('when creating a cursor with find', () => {
     let collection;
 

test/integration/crud/crud_api.test.ts

@@ -157,7 +157,6 @@ export * from '../src/operations/bulk_write';
 export * from '../src/operations/collections';
 export * from '../src/operations/command';
 export * from '../src/operations/count';
-export * from '../src/operations/count_documents';
 export * from '../src/operations/create_collection';
 export * from '../src/operations/delete';
 export * from '../src/operations/distinct';