diff --git a/Database/MongoDB.hs b/Database/MongoDB.hs index e96607c..be45429 100644 --- a/Database/MongoDB.hs +++ b/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 a connection to a MongoDB server on a non-standard port -connectOnPort :: HostName -> Network.PortID -> IO Connection -connectOnPort host port = do +-- | 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 + 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" + diff --git a/mongoDB.cabal b/mongoDB.cabal index 179fbfe..7ab06b0 100644 --- a/mongoDB.cabal +++ b/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