diff --git a/Database/MongoDB.hs b/Database/MongoDB.hs index 907f1cc..74ed602 100644 --- a/Database/MongoDB.hs +++ b/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 -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 [] _ = 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 - when (1 /= (fromBson $ fromJust $ BSON.lookup "ok" res :: Int)) $ + let res = fromLookup mres + 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" + diff --git a/Database/MongoDB/BSON.hs b/Database/MongoDB/BSON.hs index b8bfb5a..9eb2451 100644 --- a/Database/MongoDB/BSON.hs +++ b/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 diff --git a/Database/MongoDB/Util.hs b/Database/MongoDB/Util.hs index a67c712..a187cab 100644 --- a/Database/MongoDB/Util.hs +++ b/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 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