Merge remote branch 'rrichardson/master'

Database/MongoDB.hs

@@ -28,6 +28,7 @@ module Database.MongoDB
      -- * Connection
      Connection,
      connect, connectOnPort, conClose, disconnect, dropDatabase,
+     connectCluster, setTarget,
      serverInfo, serverShutdown,
      databasesInfo, databaseNames,
      -- * Database
@@ -35,14 +36,18 @@ module Database.MongoDB
      ColCreateOpt(..),
      collectionNames, createCollection, dropCollection,
      renameCollection, runCommand, validateCollection,
+     login, addUser,
      -- * Collection
      Collection, FieldSelector, FullCollection,
      NumToSkip, NumToReturn, Selector,
      QueryOpt(..),
      UpdateFlag(..),
      count, countMatching, delete, insert, insertMany, query, remove, update,
+     save, 
      -- * Convenience collection operations
      find, findOne, quickFind, quickFind',
+     -- * Query Helpers
+     whereClause,
      -- * Cursor
      Cursor,
      allDocs, allDocs', finish, nextDoc,
@@ -54,10 +59,11 @@ module Database.MongoDB
 where
 import Control.Exception
 import Control.Monad
-import Data.Binary
+import Data.Binary()
 import Data.Binary.Get
 import Data.Binary.Put
 import Data.Bits
+import Data.ByteString.Char8 (pack)
 import Data.ByteString.Internal (c2w)
 import qualified Data.ByteString.Lazy as L
 import qualified Data.ByteString.Lazy.UTF8 as L8
@@ -67,6 +73,7 @@ import qualified Data.List as List
 import qualified Data.Map as Map
 import Data.Maybe
 import Data.Typeable
+import Data.Digest.OpenSSL.MD5
 import Database.MongoDB.BSON as BSON
 import Database.MongoDB.Util
 import qualified Network
@@ -76,27 +83,61 @@ import System.IO
 import System.IO.Unsafe
 import System.Random
 
--- | A handle to a database connection
+-- | A list of handles to database connections
-data Connection = Connection { cHandle :: Handle, cRand :: IORef [Int] }
+data Connection = Connection { cHandles :: [Handle]
+                              ,cIndex :: IORef Int
+                              ,cRand :: IORef [Int] }
 
 -- | Establish a connection to a MongoDB server
 connect :: HostName -> IO Connection
 connect = flip connectOnPort $ Network.PortNumber 27017
 
--- | Establish a connection to a MongoDB server on a non-standard port
+-- | Establish connections to a list of MongoDB servers
-connectOnPort :: HostName -> Network.PortID -> IO Connection
+connectCluster :: [HostName] -> IO Connection
-connectOnPort host port = do
+connectCluster [] = throwOpFailure "No hostnames in list"
+connectCluster xs = do 
+  c <- newConnection 
+  connectAll c xs $ Network.PortNumber 27017
+
+connectAll :: Connection -> [HostName] -> Network.PortID -> IO Connection
+connectAll c [] _ = return c
+connectAll c (host:xs) port = do
   h <- Network.connectTo host port
   hSetBuffering h NoBuffering
+  connectAll (c {cHandles = h:(cHandles c)}) xs port
+
+-- | Establish a connection to a MongoDB server on a non-standard port
+connectOnPort :: HostName -> Network.PortID -> IO Connection
+connectOnPort host port = do 
+  c <- newConnection 
+  connectAll c [host] port
+
+newConnection :: IO Connection
+newConnection = do 
   r <- newStdGen
   let ns = randomRs (fromIntegral (minBound :: Int32),
                      fromIntegral (maxBound :: Int32)) r
   nsRef <- newIORef ns
-  return Connection { cHandle = h, cRand = nsRef }
+  nsIdx <- newIORef 0
+  return $ Connection [] nsIdx nsRef
 
+getHandle :: Connection -> IO Handle
+getHandle c = do 
+  i <- readIORef $ cIndex c
+  return $ (cHandles c) !! i
+
+cPut :: Connection -> L.ByteString -> IO ()
+cPut c msg = getHandle c >>= flip L.hPut msg
+ 
 -- | Close database connection
 conClose :: Connection -> IO ()
-conClose = hClose . cHandle
+conClose c = sequence_ $ map hClose $ cHandles c
+
+setTarget :: Connection -> Int -> IO ()
+setTarget c i = 
+  if i > length (cHandles c)
+    then throwOpFailure "Target index higher than length of list"
+    else writeIORef (cIndex c) i >> return ()
 
 -- | Information about the databases on the server.
 databasesInfo :: Connection -> IO BsonDoc
@@ -107,7 +148,7 @@ databasesInfo c =
 databaseNames :: Connection -> IO [Database]
 databaseNames c = do
     info <- databasesInfo c
-    let (BsonArray dbs) = fromJust $ Map.lookup (s2L "databases") info
+    let (BsonArray dbs) = fromLookup $ Map.lookup (s2L "databases") info
         names = mapMaybe (Map.lookup (s2L "name") . fromBson) dbs
     return $ List.map fromBson (names::[BsonValue])
 
@@ -138,7 +179,7 @@ serverShutdown c =
 collectionNames :: Connection -> Database -> IO [FullCollection]
 collectionNames c db = do
   docs <- quickFind' c (L.append db $ s2L ".system.namespaces") empty
-  let names = flip List.map docs $ fromBson . fromJust . BSON.lookup "name"
+  let names = flip List.map docs $ fromBson . fromLookup . BSON.lookup "name"
   return $ List.filter (L.notElem $ c2w '$') names
 
 data ColCreateOpt = CCOSize Int64  -- ^ Desired initial size for the
@@ -216,7 +257,7 @@ validateCollection :: Connection -> FullCollection -> IO String
 validateCollection c col = do
   let (db, col') = splitFullCol col
   res <- runCommand c db $ toBsonDoc [("validate", toBson col')]
-  return $ fromBson $ fromJust $ BSON.lookup "result" res
+  return $ fromBson $ fromLookup $ BSON.lookup "result" res
 
 splitFullCol :: FullCollection -> (Database, Collection)
 splitFullCol col = (L.takeWhile (c2w '.' /=) col,
@@ -228,10 +269,10 @@ splitFullCol col = (L.takeWhile (c2w '.' /=) col,
 runCommand :: Connection -> Database -> BsonDoc -> IO BsonDoc
 runCommand c db cmd = do
   mres <- findOne c (L.append db $ s2L ".$cmd") cmd
-  let res = fromJust mres
+  let res = fromLookup mres
-  when (1 /= (fromBson $ fromJust $ BSON.lookup "ok" res :: Int)) $
+  when (1 /= (fromBson $ fromLookup $ BSON.lookup "ok" res :: Int)) $
        throwOpFailure $ "command \"" ++ show cmd ++ "\" failed: " ++
-                      fromBson (fromJust $ BSON.lookup "errmsg" res)
+                      fromBson (fromLookup $ BSON.lookup "errmsg" res)
   return res
 
 -- | An Iterator over the results of a query. Use 'nextDoc' to get each
@@ -391,7 +432,7 @@ countMatching c col sel = do
   let (db, col') = splitFullCol col
   res <- runCommand c db $ toBsonDoc [("count", toBson col'),
                                       ("query", toBson sel)]
-  return $ fromBson $ fromJust $ BSON.lookup "n" res
+  return $ fromBson $ fromLookup $ BSON.lookup "n" res
 
 -- | Delete documents matching /Selector/ from the given /FullCollection/.
 delete :: Connection -> FullCollection -> Selector -> IO RequestID
@@ -402,7 +443,7 @@ delete c col sel = do
                      putI32 0
                      putBsonDoc sel
   (reqID, msg) <- packMsg c OPDelete body
-  L.hPut (cHandle c) msg
+  cPut c msg
   return reqID
 
 -- | An alias for 'delete'.
@@ -417,7 +458,7 @@ insert c col doc = do
                      putCol col
                      putBsonDoc doc
   (reqID, msg) <- packMsg c OPInsert body
-  L.hPut (cHandle c) msg
+  cPut c msg
   return reqID
 
 -- | Insert a list of documents into /FullCollection/.
@@ -428,7 +469,7 @@ insertMany c col docs = do
                putCol col
                forM_ docs putBsonDoc
   (reqID, msg) <- packMsg c OPInsert body
-  L.hPut (cHandle c) msg
+  cPut c msg
   return reqID
 
 -- | Open a cursor to find documents. If you need full functionality,
@@ -458,7 +499,7 @@ quickFind' c col sel = find c col sel >>= allDocs'
 query :: Connection -> FullCollection -> [QueryOpt] ->
          NumToSkip -> NumToReturn -> Selector -> FieldSelector -> IO Cursor
 query c col opts nskip ret sel fsel = do
-  let h = cHandle c
+  h <- getHandle c
 
   let body = runPut $ do
                putI32 $ fromQueryOpts opts
@@ -501,9 +542,52 @@ update c col flags sel obj = do
                putBsonDoc sel
                putBsonDoc obj
   (reqID, msg) <- packMsg c OPUpdate body
-  L.hPut (cHandle c) msg
+  cPut c msg
   return reqID
 
+-- | log into the mongodb /Database/ attached to the /Connection/
+login :: Connection -> Database -> String -> String -> IO BsonDoc 
+login c db user pass = do
+  doc <- runCommand c db (toBsonDoc [("getnonce", toBson (1 :: Int))])
+  let nonce = fromBson $ fromLookup $ BSON.lookup "nonce" doc :: String
+      digest = md5sum $ pack $ nonce ++ user ++ 
+                 ( md5sum $ pack (user ++ ":mongo:" ++ pass)) 
+      request = toBsonDoc [("authenticate", toBson (1 :: Int)), 
+                           ("user", toBson user), 
+                           ("nonce", toBson nonce), 
+                           ("key", toBson digest)]
+      in runCommand c db request
+
+-- | create a new user in the current /Database/
+addUser :: Connection -> Database -> String -> String -> IO BsonDoc
+addUser c db user pass = do
+  let userDoc = toBsonDoc [(s2L"user", toBson user)]
+      fdb = L.append db (s2L ".system.users")  
+  doc <- liftM (maybe userDoc id) (findOne c fdb userDoc)
+  let doc' = Map.insert (s2L "pwd") 
+               (toBson ( md5sum $ pack (user ++ ":mongo:" ++ pass))) doc
+  _ <- save c fdb doc'
+  return doc'
+
+-- | Conveniently stores the /BsonDoc/ to the /FullCollection/ 
+-- | if there is an _id present in the /BsonDoc/ then it already has 
+-- | a place in the DB, so we update it using the _id, otherwise
+-- | we insert it
+save :: Connection -> FullCollection -> BsonDoc -> IO RequestID
+save c fc doc =
+  case Map.lookup (s2L "_id") doc of 
+    Nothing -> insert c fc doc
+    Just obj -> update c fc [UFUpsert] (toBsonDoc [("_id", obj)]) doc
+
+-- | Use this in the place of the query portion of a select type query
+-- | This uses javascript and a scope supplied by a /BsonDoc/ to evaluate 
+-- | documents in the database for retrieval. 
+-- | Example: 
+-- | > findOne conn mycoll $ whereClause "this.name == (name1 + name2)" 
+-- | >     (toBsonDoc [("name1", toBson "mar"), ("name2", toBson "tha")])
+whereClause :: String -> BsonDoc -> BsonDoc
+whereClause qry scope = toBsonDoc [("$where", (BsonCodeWScope (s2L qry) scope))]
+
 data Hdr = Hdr {
       hMsgLen :: Int32,
       -- hReqID :: Int32,
@@ -594,7 +678,7 @@ getFirstDoc docBytes = flip runGet docBytes $ do
 
 getMore :: Cursor -> IO (Maybe BsonDoc)
 getMore cur = do
-  let h = cHandle $ curCon cur
+  h <- getHandle $ curCon cur
 
   cid <- readIORef $ curID cur
   let body = runPut $ do
@@ -625,7 +709,7 @@ getMore cur = do
 -- 'allDocs', 'allDocs'', or 'nextDoc'.
 finish :: Cursor -> IO ()
 finish cur = do
-  let h = cHandle $ curCon cur
+  h <- getHandle $ curCon cur
   cid <- readIORef $ curID cur
   unless (cid == 0) $ do
       let body = runPut $ do
@@ -743,3 +827,8 @@ validateCollectionName col = do
   when (L.head col == c2w '.' || L.last col == c2w '.') $
        throwColInvalid $ "Collection can't start or end with '.': " ++ show col
   return (db, col')
+
+fromLookup :: (Maybe a) -> a
+fromLookup (Just m) = m
+fromLookup Nothing = throwColInvalid "cannot find key"  
+

Database/MongoDB/BSON.hs

@@ -73,6 +73,7 @@ data BsonValue
     | BsonSymbol L8.ByteString
     | BsonInt32 Int32
     | BsonInt64 Int64
+    | BsonCodeWScope L8.ByteString BsonDoc
     | BsonMinKey
     | BsonMaxKey
     deriving (Show, Eq, Ord)
@@ -195,7 +196,11 @@ getVal DataSymbol = do
   sLen1 <- getI32
   (_sLen2, s) <- getS
   return (fromIntegral $ 4 + sLen1, BsonString s)
-getVal DataCodeWScope = fail "DataCodeWScope not yet supported" -- TODO
+getVal DataCodeWScope = do
+  sLen1 <- getI32
+  (_, qry) <- getS
+  (_, scope) <- getDoc
+  return (fromIntegral sLen1, BsonCodeWScope qry scope)
 getVal DataInt = liftM ((,) 4 . BsonInt32 . fromIntegral) getI32
 getVal DataTimestamp = fail "DataTimestamp not yet supported" -- TODO
 
@@ -240,7 +245,7 @@ putType BsonRegex{}    = putDataType DataRegex
 -- putType = putDataType DataRef
 -- putType = putDataType DataCode
 putType BsonSymbol{}   = putDataType DataSymbol
--- putType = putDataType DataCodeWScope
+putType BsonCodeWScope{} = putDataType DataCodeWScope
 putType BsonInt32 {}   = putDataType DataInt
 putType BsonInt64 {}   = putDataType DataLong
 -- putType = putDataType DataTimestamp
@@ -249,7 +254,7 @@ putType BsonMaxKey     = putDataType DataMaxKey
 
 putVal :: BsonValue -> Put
 putVal (BsonDouble d)   = putFloat64le d
-putVal (BsonString s)   = putI32 (fromIntegral $ 1 + L8.length s) >> putS s
+putVal (BsonString s)   = putStrSz s
 putVal (BsonObject o)   = putObj o
 putVal (BsonArray es)   = putOutterObj bs
     where bs = runPut $ forM_ (List.zip [(0::Int) .. ] es) $ \(i, e) ->
@@ -270,6 +275,9 @@ putVal (BsonRegex r opt)= do putS r
 putVal (BsonSymbol s)   = putI32 (fromIntegral $ 1 + L8.length s) >> putS s
 putVal (BsonInt32 i)    = putI32 i
 putVal (BsonInt64 i)    = putI64 i
+putVal (BsonCodeWScope q s) =
+  let bytes = runPut (putStrSz q >> putObj s)
+    in (putI32 $ (+4) $ fromIntegral $ L.length bytes) >> putLazyByteString bytes
 putVal BsonMinKey       = putNothing
 putVal BsonMaxKey       = putNothing
 

Database/MongoDB/Util.hs

@@ -26,7 +26,7 @@ WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 module Database.MongoDB.Util
     (
      putI8, putI32, putI64, putNothing, putNull, putS,
-     getI8, getI32, getI64, getC, getS, getNull,
+     getI8, getI32, getI64, getC, getS, getNull, putStrSz,
     )
 where
 import Control.Exception (assert)
@@ -75,3 +75,6 @@ putNull = putI8 (0::Int)
 
 putS :: L8.ByteString -> Put
 putS s = putLazyByteString s >> putNull
+
+putStrSz :: L.ByteString -> Put
+putStrSz s = putI32 (fromIntegral $ 1 + L8.length s) >> putS s

mongoDB.cabal

@@ -21,7 +21,8 @@ Build-Depends:       base < 5,
 		     network,
 		     random,
 		     time,
-		     utf8-string
+		     utf8-string,
+         nano-md5
 Build-Type:          Simple
 Exposed-modules:     Database.MongoDB,
 		     Database.MongoDB.BSON