Merge remote branch 'rrichardson/master'

This commit is contained in:
Scott R. Parish 2010-02-06 14:40:14 -06:00
commit 2f452e989c
4 changed files with 128 additions and 27 deletions

View file

@ -28,6 +28,7 @@ module Database.MongoDB
-- * Connection -- * Connection
Connection, Connection,
connect, connectOnPort, conClose, disconnect, dropDatabase, connect, connectOnPort, conClose, disconnect, dropDatabase,
connectCluster, setTarget,
serverInfo, serverShutdown, serverInfo, serverShutdown,
databasesInfo, databaseNames, databasesInfo, databaseNames,
-- * Database -- * Database
@ -35,14 +36,18 @@ module Database.MongoDB
ColCreateOpt(..), ColCreateOpt(..),
collectionNames, createCollection, dropCollection, collectionNames, createCollection, dropCollection,
renameCollection, runCommand, validateCollection, renameCollection, runCommand, validateCollection,
login, addUser,
-- * Collection -- * Collection
Collection, FieldSelector, FullCollection, Collection, FieldSelector, FullCollection,
NumToSkip, NumToReturn, Selector, NumToSkip, NumToReturn, Selector,
QueryOpt(..), QueryOpt(..),
UpdateFlag(..), UpdateFlag(..),
count, countMatching, delete, insert, insertMany, query, remove, update, count, countMatching, delete, insert, insertMany, query, remove, update,
save,
-- * Convenience collection operations -- * Convenience collection operations
find, findOne, quickFind, quickFind', find, findOne, quickFind, quickFind',
-- * Query Helpers
whereClause,
-- * Cursor -- * Cursor
Cursor, Cursor,
allDocs, allDocs', finish, nextDoc, allDocs, allDocs', finish, nextDoc,
@ -54,10 +59,11 @@ module Database.MongoDB
where where
import Control.Exception import Control.Exception
import Control.Monad import Control.Monad
import Data.Binary import Data.Binary()
import Data.Binary.Get import Data.Binary.Get
import Data.Binary.Put import Data.Binary.Put
import Data.Bits import Data.Bits
import Data.ByteString.Char8 (pack)
import Data.ByteString.Internal (c2w) import Data.ByteString.Internal (c2w)
import qualified Data.ByteString.Lazy as L import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Lazy.UTF8 as L8 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 qualified Data.Map as Map
import Data.Maybe import Data.Maybe
import Data.Typeable import Data.Typeable
import Data.Digest.OpenSSL.MD5
import Database.MongoDB.BSON as BSON import Database.MongoDB.BSON as BSON
import Database.MongoDB.Util import Database.MongoDB.Util
import qualified Network import qualified Network
@ -76,27 +83,61 @@ import System.IO
import System.IO.Unsafe import System.IO.Unsafe
import System.Random 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 -- | Establish a connection to a MongoDB server
connect :: HostName -> IO Connection connect :: HostName -> IO Connection
connect = flip connectOnPort $ Network.PortNumber 27017 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 h <- Network.connectTo host port
hSetBuffering h NoBuffering 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 r <- newStdGen
let ns = randomRs (fromIntegral (minBound :: Int32), let ns = randomRs (fromIntegral (minBound :: Int32),
fromIntegral (maxBound :: Int32)) r fromIntegral (maxBound :: Int32)) r
nsRef <- newIORef ns 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 -- | Close database connection
conClose :: Connection -> IO () 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. -- | Information about the databases on the server.
databasesInfo :: Connection -> IO BsonDoc databasesInfo :: Connection -> IO BsonDoc
@ -107,7 +148,7 @@ databasesInfo c =
databaseNames :: Connection -> IO [Database] databaseNames :: Connection -> IO [Database]
databaseNames c = do databaseNames c = do
info <- databasesInfo c 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 names = mapMaybe (Map.lookup (s2L "name") . fromBson) dbs
return $ List.map fromBson (names::[BsonValue]) return $ List.map fromBson (names::[BsonValue])
@ -138,7 +179,7 @@ serverShutdown c =
collectionNames :: Connection -> Database -> IO [FullCollection] collectionNames :: Connection -> Database -> IO [FullCollection]
collectionNames c db = do collectionNames c db = do
docs <- quickFind' c (L.append db $ s2L ".system.namespaces") empty 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 return $ List.filter (L.notElem $ c2w '$') names
data ColCreateOpt = CCOSize Int64 -- ^ Desired initial size for the data ColCreateOpt = CCOSize Int64 -- ^ Desired initial size for the
@ -216,7 +257,7 @@ validateCollection :: Connection -> FullCollection -> IO String
validateCollection c col = do validateCollection c col = do
let (db, col') = splitFullCol col let (db, col') = splitFullCol col
res <- runCommand c db $ toBsonDoc [("validate", toBson 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 :: FullCollection -> (Database, Collection)
splitFullCol col = (L.takeWhile (c2w '.' /=) col, splitFullCol col = (L.takeWhile (c2w '.' /=) col,
@ -228,10 +269,10 @@ splitFullCol col = (L.takeWhile (c2w '.' /=) col,
runCommand :: Connection -> Database -> BsonDoc -> IO BsonDoc runCommand :: Connection -> Database -> BsonDoc -> IO BsonDoc
runCommand c db cmd = do runCommand c db cmd = do
mres <- findOne c (L.append db $ s2L ".$cmd") cmd 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: " ++ throwOpFailure $ "command \"" ++ show cmd ++ "\" failed: " ++
fromBson (fromJust $ BSON.lookup "errmsg" res) fromBson (fromLookup $ BSON.lookup "errmsg" res)
return res return res
-- | An Iterator over the results of a query. Use 'nextDoc' to get each -- | 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 let (db, col') = splitFullCol col
res <- runCommand c db $ toBsonDoc [("count", toBson col'), res <- runCommand c db $ toBsonDoc [("count", toBson col'),
("query", toBson sel)] ("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 documents matching /Selector/ from the given /FullCollection/.
delete :: Connection -> FullCollection -> Selector -> IO RequestID delete :: Connection -> FullCollection -> Selector -> IO RequestID
@ -402,7 +443,7 @@ delete c col sel = do
putI32 0 putI32 0
putBsonDoc sel putBsonDoc sel
(reqID, msg) <- packMsg c OPDelete body (reqID, msg) <- packMsg c OPDelete body
L.hPut (cHandle c) msg cPut c msg
return reqID return reqID
-- | An alias for 'delete'. -- | An alias for 'delete'.
@ -417,7 +458,7 @@ insert c col doc = do
putCol col putCol col
putBsonDoc doc putBsonDoc doc
(reqID, msg) <- packMsg c OPInsert body (reqID, msg) <- packMsg c OPInsert body
L.hPut (cHandle c) msg cPut c msg
return reqID return reqID
-- | Insert a list of documents into /FullCollection/. -- | Insert a list of documents into /FullCollection/.
@ -428,7 +469,7 @@ insertMany c col docs = do
putCol col putCol col
forM_ docs putBsonDoc forM_ docs putBsonDoc
(reqID, msg) <- packMsg c OPInsert body (reqID, msg) <- packMsg c OPInsert body
L.hPut (cHandle c) msg cPut c msg
return reqID return reqID
-- | Open a cursor to find documents. If you need full functionality, -- | 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] -> query :: Connection -> FullCollection -> [QueryOpt] ->
NumToSkip -> NumToReturn -> Selector -> FieldSelector -> IO Cursor NumToSkip -> NumToReturn -> Selector -> FieldSelector -> IO Cursor
query c col opts nskip ret sel fsel = do query c col opts nskip ret sel fsel = do
let h = cHandle c h <- getHandle c
let body = runPut $ do let body = runPut $ do
putI32 $ fromQueryOpts opts putI32 $ fromQueryOpts opts
@ -501,9 +542,52 @@ update c col flags sel obj = do
putBsonDoc sel putBsonDoc sel
putBsonDoc obj putBsonDoc obj
(reqID, msg) <- packMsg c OPUpdate body (reqID, msg) <- packMsg c OPUpdate body
L.hPut (cHandle c) msg cPut c msg
return reqID 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 { data Hdr = Hdr {
hMsgLen :: Int32, hMsgLen :: Int32,
-- hReqID :: Int32, -- hReqID :: Int32,
@ -594,7 +678,7 @@ getFirstDoc docBytes = flip runGet docBytes $ do
getMore :: Cursor -> IO (Maybe BsonDoc) getMore :: Cursor -> IO (Maybe BsonDoc)
getMore cur = do getMore cur = do
let h = cHandle $ curCon cur h <- getHandle $ curCon cur
cid <- readIORef $ curID cur cid <- readIORef $ curID cur
let body = runPut $ do let body = runPut $ do
@ -625,7 +709,7 @@ getMore cur = do
-- 'allDocs', 'allDocs'', or 'nextDoc'. -- 'allDocs', 'allDocs'', or 'nextDoc'.
finish :: Cursor -> IO () finish :: Cursor -> IO ()
finish cur = do finish cur = do
let h = cHandle $ curCon cur h <- getHandle $ curCon cur
cid <- readIORef $ curID cur cid <- readIORef $ curID cur
unless (cid == 0) $ do unless (cid == 0) $ do
let body = runPut $ do let body = runPut $ do
@ -743,3 +827,8 @@ validateCollectionName col = do
when (L.head col == c2w '.' || L.last col == c2w '.') $ when (L.head col == c2w '.' || L.last col == c2w '.') $
throwColInvalid $ "Collection can't start or end with '.': " ++ show col throwColInvalid $ "Collection can't start or end with '.': " ++ show col
return (db, col') return (db, col')
fromLookup :: (Maybe a) -> a
fromLookup (Just m) = m
fromLookup Nothing = throwColInvalid "cannot find key"

View file

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

View file

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

View file

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