From dc99780c03b4387cdb3d701cde162d4a5fbe0201 Mon Sep 17 00:00:00 2001
From: "Scott R. Parish" <sparish@peak6.com>
Date: Mon, 18 Jan 2010 22:13:32 -0600
Subject: [PATCH] count, countMatching

---
 Database/MongoDB.hs | 17 +++++++++++++++--
 TODO                |  2 --
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/Database/MongoDB.hs b/Database/MongoDB.hs
index 3089fb1..b3fa732 100644
--- a/Database/MongoDB.hs
+++ b/Database/MongoDB.hs
@@ -37,7 +37,7 @@ module Database.MongoDB
      Collection, FieldSelector, NumToSkip, NumToReturn, Selector,
      QueryOpt(..),
      UpdateFlag(..),
-     delete, insert, insertMany, query, remove, update,
+     count, countMatching, delete, insert, insertMany, query, remove, update,
      -- * Convience collection operations
      find, findOne, quickFind, quickFind',
      -- * Cursor
@@ -51,7 +51,7 @@ import Data.Binary
 import Data.Binary.Get
 import Data.Binary.Put
 import Data.Bits
-import Data.ByteString.Char8 hiding (find)
+import Data.ByteString.Char8 hiding (count, find)
 import qualified Data.ByteString.Lazy as L
 import qualified Data.ByteString.Lazy.UTF8 as L8
 import Data.Int
@@ -327,6 +327,19 @@ fromUpdateFlags :: [UpdateFlag] -> Int32
 fromUpdateFlags flags = List.foldl (.|.) 0 $
                         flip fmap flags $ (1 `shiftL`) . fromEnum
 
+-- | Return the number of documents in /Collection/.
+count :: Connection -> Collection -> IO Int64
+count c col = countMatching c col BSON.empty
+
+-- | Return the number of documents in /Collection/ matching /Selector/
+countMatching :: Connection -> Collection -> Selector -> IO Int64
+countMatching c col sel = do
+  let db = dbFromCol col
+      col' = colMinusDB col
+  res <- dbCmd c db $ toBsonDoc [("count", toBson col'),
+                                 ("query", BsonObject sel)]
+  return $ fromBson $ fromJust $ BSON.lookup "n" res
+
 -- | Delete documents matching /Selector/ from the given /Collection/.
 delete :: Connection -> Collection -> Selector -> IO RequestID
 delete c col sel = do
diff --git a/TODO b/TODO
index 4179c5d..ddf158f 100644
--- a/TODO
+++ b/TODO
@@ -42,8 +42,6 @@ MongoDB
     * modify
     * replace
     * repsert
-    * getCount()
-    * getCount(query)
   - index operations
     * createIndex
     * dropIndex