added login/user mgmt, added multiple handles in the connection

Database/MongoDB.hs

@@ -28,6 +28,7 @@ module Database.MongoDB
      -- * Connection
      Connection,
      connect, connectOnPort, conClose, disconnect, dropDatabase,
+     coonnectCluster, setTarget,
      serverInfo, serverShutdown,
      databasesInfo, databaseNames,
      -- * Database
@@ -35,12 +36,14 @@ 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, modify, replace, repsert, 
      -- * Convenience collection operations
      find, findOne, quickFind, quickFind',
      -- * Cursor
@@ -54,10 +57,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 +71,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 +81,59 @@ import System.IO
 import System.IO.Unsafe
 import System.Random
 
--- | A handle to a database connection
-data Connection = Connection { cHandle :: Handle, cRand :: IORef [Int] }
+-- | A list of handles to database connections
+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 connections to a list of MongoDB servers
+connectCluster :: [HostName] -> IO Connection
+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 [] p = 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 
-  h <- Network.connectTo host port
-  hSetBuffering h NoBuffering
+  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
+
+write :: Connection -> L.ByteString -> IO ()
+write 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
@@ -402,7 +439,7 @@ delete c col sel = do
                      putI32 0
                      putBsonDoc sel
   (reqID, msg) <- packMsg c OPDelete body
-  L.hPut (cHandle c) msg
+  write c msg
   return reqID
 
 -- | An alias for 'delete'.
@@ -417,7 +454,7 @@ insert c col doc = do
                      putCol col
                      putBsonDoc doc
   (reqID, msg) <- packMsg c OPInsert body
-  L.hPut (cHandle c) msg
+  write c msg
   return reqID
 
 -- | Insert a list of documents into /FullCollection/.
@@ -428,7 +465,7 @@ insertMany c col docs = do
                putCol col
                forM_ docs putBsonDoc
   (reqID, msg) <- packMsg c OPInsert body
-  L.hPut (cHandle c) msg
+  write c msg
   return reqID
 
 -- | Open a cursor to find documents. If you need full functionality,
@@ -459,7 +496,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
@@ -502,9 +539,37 @@ update c col flags sel obj = do
                putBsonDoc sel
                putBsonDoc obj
   (reqID, msg) <- packMsg c OPUpdate body
-  L.hPut (cHandle c) msg
+  write c msg
   return reqID
 
+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
+
+addUser :: Connection -> String -> String -> String -> IO BsonDoc
+addUser c db user pass = do
+  let userDoc = toBsonDoc [(s2L"user", toBson user)]
+      fdb = s2L (db ++ ".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'
+
+
+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
+  
+
 data Hdr = Hdr {
       hMsgLen :: Int32,
       -- hReqID :: Int32,
@@ -595,7 +660,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
@@ -626,7 +691,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
   if cid == 0
     then return ()
@@ -746,3 +811,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"  
+

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