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

This commit is contained in:
Rick Richardson 2010-02-02 16:06:53 -05:00
parent 9bc616ccf5
commit 0cf0da8ab6
2 changed files with 87 additions and 16 deletions

View file

@ -28,6 +28,7 @@ module Database.MongoDB
-- * Connection -- * Connection
Connection, Connection,
connect, connectOnPort, conClose, disconnect, dropDatabase, connect, connectOnPort, conClose, disconnect, dropDatabase,
coonnectCluster, setTarget,
serverInfo, serverShutdown, serverInfo, serverShutdown,
databasesInfo, databaseNames, databasesInfo, databaseNames,
-- * Database -- * Database
@ -35,12 +36,14 @@ 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, modify, replace, repsert,
-- * Convenience collection operations -- * Convenience collection operations
find, findOne, quickFind, quickFind', find, findOne, quickFind, quickFind',
-- * Cursor -- * Cursor
@ -54,10 +57,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 +71,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 +81,59 @@ 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 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 -- | Establish a connection to a MongoDB server on a non-standard port
connectOnPort :: HostName -> Network.PortID -> IO Connection connectOnPort :: HostName -> Network.PortID -> IO Connection
connectOnPort host port = do connectOnPort host port = do
h <- Network.connectTo host port c <- newConnection
hSetBuffering h NoBuffering 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
write :: Connection -> L.ByteString -> IO ()
write 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
@ -402,7 +439,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 write c msg
return reqID return reqID
-- | An alias for 'delete'. -- | An alias for 'delete'.
@ -417,7 +454,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 write c msg
return reqID return reqID
-- | Insert a list of documents into /FullCollection/. -- | Insert a list of documents into /FullCollection/.
@ -428,7 +465,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 write 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,
@ -459,7 +496,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
@ -502,9 +539,37 @@ 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 write c msg
return reqID 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 { data Hdr = Hdr {
hMsgLen :: Int32, hMsgLen :: Int32,
-- hReqID :: Int32, -- hReqID :: Int32,
@ -595,7 +660,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
@ -626,7 +691,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
if cid == 0 if cid == 0
then return () then return ()
@ -746,3 +811,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

@ -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