From abf2686be6bd1ec1bde0e11a2cda23c1e0cf6861 Mon Sep 17 00:00:00 2001
From: Enok <416828041@qq.com>
Date: Sun, 15 Mar 2020 15:52:57 +0800
Subject: [PATCH] count support

---
 README.md            |  3 +++
 src/command/count.rs | 29 +++++++++++++++++++++++++++++
 src/command/mod.rs   |  2 ++
 src/lib.rs           |  2 ++
 test.ts              |  7 +++++++
 ts/collection.ts     | 17 +++++++++++++++++
 ts/types.ts          |  1 +
 7 files changed, 61 insertions(+)
 create mode 100644 src/command/count.rs

diff --git a/README.md b/README.md
index 5908cb81..a695318e 100644
--- a/README.md
+++ b/README.md
@@ -52,6 +52,9 @@ const user1 = await users.findOne({ _id: insertId });
 // find
 const users = await users.find({ username: { $ne: null } });
 
+// count
+const count = await users.count({ username: { $ne: null } });
+
 // aggregation
 const docs = await users.aggregation([
   { $match: { username: "many" } },
diff --git a/src/command/count.rs b/src/command/count.rs
new file mode 100644
index 00000000..50130e13
--- /dev/null
+++ b/src/command/count.rs
@@ -0,0 +1,29 @@
+use crate::*;
+use serde_json::Value;
+use util::maybe_json_to_document;
+
+#[derive(Deserialize, Debug)]
+#[serde(rename_all = "camelCase")]
+struct CountArgs {
+    db_name: String,
+    collection_name: String,
+    filter: Option<Value>,
+}
+
+pub fn count(command: Command) -> CoreOp {
+    let fut = async move {
+        let client = command.get_client();
+        let data = command.data;
+        let args: CountArgs = serde_json::from_slice(data.unwrap().as_ref()).unwrap();
+        let db_name = args.db_name;
+        let collection_name = args.collection_name;
+        let filter = maybe_json_to_document(args.filter);
+
+        let database = client.database(&db_name);
+        let collection = database.collection(&collection_name);
+
+        let count = collection.count_documents(filter, None).unwrap();
+        Ok(util::async_result(&command.args, count))
+    };
+    CoreOp::Async(fut.boxed())
+}
diff --git a/src/command/mod.rs b/src/command/mod.rs
index dc679a19..de6824a4 100644
--- a/src/command/mod.rs
+++ b/src/command/mod.rs
@@ -1,5 +1,6 @@
 mod aggregation;
 mod connect;
+mod count;
 mod delete;
 mod find;
 mod indexes;
@@ -10,6 +11,7 @@ mod update;
 
 pub use aggregation::aggregate;
 pub use connect::{connect_with_options, connect_with_uri};
+pub use count::count;
 pub use delete::delete;
 pub use find::find;
 pub use indexes::create_indexes;
diff --git a/src/lib.rs b/src/lib.rs
index ec227e4a..237942d1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -40,6 +40,7 @@ pub enum CommandType {
     Delete,
     Aggregate,
     Update,
+    Count,
     CreateIndexes,
 }
 
@@ -103,6 +104,7 @@ fn op_command(data: &[u8], zero_copy: Option<ZeroCopyBuf>) -> CoreOp {
         CommandType::Update => command::update,
         CommandType::Aggregate => command::aggregate,
         CommandType::CreateIndexes => command::create_indexes,
+        CommandType::Count => command::count,
     };
 
     executor(Command::new(args, zero_copy))
diff --git a/test.ts b/test.ts
index 1b9244e7..b950efdd 100644
--- a/test.ts
+++ b/test.ts
@@ -117,6 +117,13 @@ test(async function testFind() {
   assertEquals(notFound, []);
 });
 
+test(async function testCount() {
+  const db = getClient().database("test");
+  const users = db.collection("mongo_test_users");
+  const count = await users.count({ username: "many" });
+  assertEquals(count, 2);
+});
+
 test(async function testAggregation() {
   const db = getClient().database("test");
   const users = db.collection("mongo_test_users");
diff --git a/ts/collection.ts b/ts/collection.ts
index 6aeccdec..6daf03af 100644
--- a/ts/collection.ts
+++ b/ts/collection.ts
@@ -28,6 +28,23 @@ export class Collection {
     return doc;
   }
 
+  public async count(filter?: Object): Promise<number> {
+    const count = await dispatchAsync(
+      {
+        command_type: CommandType.Count,
+        client_id: this.client.clientId
+      },
+      encode(
+        JSON.stringify({
+          dbName: this.dbName,
+          collectionName: this.collectionName,
+          filter
+        })
+      )
+    );
+    return count as number;
+  }
+
   public async findOne(filter?: Object): Promise<any> {
     return this._find(filter, { findOne: true });
   }
diff --git a/ts/types.ts b/ts/types.ts
index 9b2b629b..783f6b7d 100644
--- a/ts/types.ts
+++ b/ts/types.ts
@@ -11,6 +11,7 @@ export enum CommandType {
   Delete = "Delete",
   Update = "Update",
   Aggregate = "Aggregate",
+  Count = "Count",
   CreateIndexes = "CreateIndexes"
 }