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,
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"

View file

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