diff --git a/Database/MongoDB.hs b/Database/MongoDB.hs --- a/Database/MongoDB.hs +++ b/Database/MongoDB.hs @@ -17,9 +17,9 @@ > run = do > clearTeams > insertTeams -> printDocs "All Teams" =<< allTeams -> printDocs "National League Teams" =<< nationalLeagueTeams -> printDocs "New York Teams" =<< newYorkTeams +> allTeams >>= printDocs "All Teams" +> nationalLeagueTeams >>= printDocs "National League Teams" +> newYorkTeams >>= printDocs "New York Teams" > > clearTeams = delete (select [] "team") > diff --git a/Database/MongoDB/Connection.hs b/Database/MongoDB/Connection.hs --- a/Database/MongoDB/Connection.hs +++ b/Database/MongoDB/Connection.hs @@ -4,20 +4,21 @@ module Database.MongoDB.Connection ( -- * Util - IOE, runIOE, + Secs, IOE, runIOE, -- * Connection Pipe, close, isClosed, -- * Server - Host(..), PortID(..), defaultPort, host, showHostPort, readHostPort, readHostPortM, connect, + Host(..), PortID(..), defaultPort, host, showHostPort, readHostPort, readHostPortM, + globalConnectTimeout, connect, connect', -- * Replica Set - ReplicaSetName, openReplicaSet, ReplicaSet, primary, secondaryOk, closeReplicaSet + ReplicaSetName, openReplicaSet, openReplicaSet', + ReplicaSet, primary, secondaryOk, closeReplicaSet, replSetName ) where import Prelude hiding (lookup) -import Database.MongoDB.Internal.Protocol (Pipe, writeMessage, readMessage) -import System.IO.Pipeline (IOE, IOStream(..), newPipeline, close, isClosed) +import Database.MongoDB.Internal.Protocol (Pipe, newPipe) +import System.IO.Pipeline (IOE, close, isClosed) import System.IO.Error as E (try) -import System.IO (hClose) import Network (HostName, PortID(..), connectTo) import Text.ParserCombinators.Parsec as T (parse, many1, letter, digit, char, eof, spaces, try, (<|>)) import Control.Monad.Identity (runIdentity) @@ -30,6 +31,9 @@ import Database.MongoDB.Query (access, slaveOk, Failure(ConnectionFailure), Command, runCommand) import Database.MongoDB.Internal.Util (untilSuccess, liftIOE, runIOE, updateAssocs, shuffle) import Data.List as L (lookup, intersect, partition, (\\), delete) +import Data.IORef (IORef, newIORef, readIORef) +import System.Timeout (timeout) +import System.IO.Unsafe (unsafePerformIO) adminCommand :: Command -> Pipe -> IOE Document -- ^ Run command against admin database on server connected to pipe. Fail if connection fails. @@ -80,33 +84,55 @@ -- ^ Read string \"hostname:port\" as @Host hostname (PortNumber port)@ or \"hostname\" as @host hostname@ (default port). Error if string does not match either syntax. readHostPort = runIdentity . readHostPortM +type Secs = Double + +globalConnectTimeout :: IORef Secs +-- ^ 'connect' (and 'openReplicaSet') fails if it can't connect within this many seconds (default is 6 seconds). Use 'connect\'' (and 'openReplicaSet\'') if you want to ignore this global and specify your own timeout. Note, this timeout only applies to initial connection establishment, not when reading/writing to the connection. +globalConnectTimeout = unsafePerformIO (newIORef 6) +{-# NOINLINE globalConnectTimeout #-} + connect :: Host -> IOE Pipe --- ^ Connect to Host returning pipelined TCP connection. Throw IOError if problem connecting. -connect (Host hostname port) = do - handle <- ErrorT . E.try $ connectTo hostname port - lift $ newPipeline $ IOStream (writeMessage handle) (readMessage handle) (hClose handle) +-- ^ Connect to Host returning pipelined TCP connection. Throw IOError if connection refused or no response within 'globalConnectTimeout'. +connect h = lift (readIORef globalConnectTimeout) >>= flip connect' h +connect' :: Secs -> Host -> IOE Pipe +-- ^ Connect to Host returning pipelined TCP connection. Throw IOError if connection refused or no response within given number of seconds. +connect' timeoutSecs (Host hostname port) = do + handle <- ErrorT . E.try $ do + mh <- timeout (round $ timeoutSecs * 1000000) (connectTo hostname port) + maybe (ioError $ userError "connect timed out") return mh + lift $ newPipe handle + -- * Replica Set type ReplicaSetName = UString -- | Maintains a connection (created on demand) to each server in the named replica set -data ReplicaSet = ReplicaSet ReplicaSetName (MVar [(Host, Maybe Pipe)]) +data ReplicaSet = ReplicaSet ReplicaSetName (MVar [(Host, Maybe Pipe)]) Secs +replSetName :: ReplicaSet -> UString +-- ^ name of connected replica set +replSetName (ReplicaSet rsName _ _) = rsName + openReplicaSet :: (ReplicaSetName, [Host]) -> IOE ReplicaSet --- ^ Open connections (on demand) to servers in replica set. Supplied hosts is seed list. At least one of them must be a live member of the named replica set, otherwise fail. -openReplicaSet (rsName, seedList) = do - rs <- ReplicaSet rsName <$> newMVar (map (, Nothing) seedList) +-- ^ Open connections (on demand) to servers in replica set. Supplied hosts is seed list. At least one of them must be a live member of the named replica set, otherwise fail. The value of 'globalConnectTimeout' at the time of this call is the timeout used for future member connect attempts. To use your own value call 'openReplicaSet\'' instead. +openReplicaSet rsSeed = lift (readIORef globalConnectTimeout) >>= flip openReplicaSet' rsSeed + +openReplicaSet' :: Secs -> (ReplicaSetName, [Host]) -> IOE ReplicaSet +-- ^ Open connections (on demand) to servers in replica set. Supplied hosts is seed list. At least one of them must be a live member of the named replica set, otherwise fail. Supplied seconds timeout is used for connect attempts to members. +openReplicaSet' timeoutSecs (rsName, seedList) = do + vMembers <- newMVar (map (, Nothing) seedList) + let rs = ReplicaSet rsName vMembers timeoutSecs _ <- updateMembers rs return rs closeReplicaSet :: ReplicaSet -> IO () -- ^ Close all connections to replica set -closeReplicaSet (ReplicaSet _ vMembers) = withMVar vMembers $ mapM_ (maybe (return ()) close . snd) +closeReplicaSet (ReplicaSet _ vMembers _) = withMVar vMembers $ mapM_ (maybe (return ()) close . snd) primary :: ReplicaSet -> IOE Pipe -- ^ Return connection to current primary of replica set. Fail if no primary available. -primary rs@(ReplicaSet rsName _) = do +primary rs@(ReplicaSet rsName _ _) = do mHost <- statedPrimary <$> updateMembers rs case mHost of Just host' -> connection rs Nothing host' @@ -133,7 +159,7 @@ updateMembers :: ReplicaSet -> IOE ReplicaInfo -- ^ Fetch replica info from any server and update members accordingly -updateMembers rs@(ReplicaSet _ vMembers) = do +updateMembers rs@(ReplicaSet _ vMembers _) = do (host', info) <- untilSuccess (fetchReplicaInfo rs) =<< readMVar vMembers modifyMVar vMembers $ \members -> do let ((members', old), new) = intersection (map readHostPort $ at "hosts" info) members @@ -147,7 +173,7 @@ fetchReplicaInfo :: ReplicaSet -> (Host, Maybe Pipe) -> IOE ReplicaInfo -- Connect to host and fetch replica info from host creating new connection if missing or closed (previously failed). Fail if not member of named replica set. -fetchReplicaInfo rs@(ReplicaSet rsName _) (host', mPipe) = do +fetchReplicaInfo rs@(ReplicaSet rsName _ _) (host', mPipe) = do pipe <- connection rs mPipe host' info <- adminCommand ["isMaster" =: (1 :: Int)] pipe case D.lookup "setName" info of @@ -157,11 +183,11 @@ connection :: ReplicaSet -> Maybe Pipe -> Host -> IOE Pipe -- ^ Return new or existing connection to member of replica set. If pipe is already known for host it is given, but we still test if it is open. -connection (ReplicaSet _ vMembers) mPipe host' = +connection (ReplicaSet _ vMembers timeoutSecs) mPipe host' = maybe conn (\p -> lift (isClosed p) >>= \bad -> if bad then conn else return p) mPipe where conn = modifyMVar vMembers $ \members -> do - let new = connect host' >>= \pipe -> return (updateAssocs host' (Just pipe) members, pipe) + let new = connect' timeoutSecs host' >>= \pipe -> return (updateAssocs host' (Just pipe) members, pipe) case L.lookup host' members of Just (Just pipe) -> lift (isClosed pipe) >>= \bad -> if bad then new else return (members, pipe) _ -> new diff --git a/Database/MongoDB/Internal/Protocol.hs b/Database/MongoDB/Internal/Protocol.hs --- a/Database/MongoDB/Internal/Protocol.hs +++ b/Database/MongoDB/Internal/Protocol.hs @@ -7,11 +7,9 @@ module Database.MongoDB.Internal.Protocol ( FullCollection, -- * Pipe - Pipe, send, call, - -- * Message - writeMessage, readMessage, + Pipe, newPipe, send, call, -- ** Notice - Notice(..), UpdateOption(..), DeleteOption(..), CursorId, + Notice(..), InsertOption(..), UpdateOption(..), DeleteOption(..), CursorId, -- ** Request Request(..), QueryOption(..), -- ** Reply @@ -24,9 +22,9 @@ import Control.Applicative ((<$>)) import Control.Arrow ((***)) import Data.ByteString.Lazy as B (length, hPut) -import System.IO.Pipeline (IOE, Pipeline) +import System.IO.Pipeline (IOE, Pipeline, newPipeline, IOStream(..)) import qualified System.IO.Pipeline as P (send, call) -import System.IO (Handle) +import System.IO (Handle, hClose) import Data.Bson (Document, UString) import Data.Bson.Binary import Data.Binary.Put @@ -47,6 +45,10 @@ type Pipe = Pipeline Response Message -- ^ Thread-safe TCP connection with pipelined requests +newPipe :: Handle -> IO Pipe +-- ^ Create pipe over handle +newPipe handle = newPipeline $ IOStream (writeMessage handle) (readMessage handle) (hClose handle) + send :: Pipe -> [Notice] -> IOE () -- ^ Send notices as a contiguous batch to server with no reply. Throw IOError if connection fails. send pipe notices = P.send pipe (notices, Nothing) @@ -135,6 +137,7 @@ data Notice = Insert { iFullCollection :: FullCollection, + iOptions :: [InsertOption], iDocuments :: [Document]} | Update { uFullCollection :: FullCollection, @@ -149,6 +152,9 @@ kCursorIds :: [CursorId]} deriving (Show, Eq) +data InsertOption = KeepGoing -- ^ If set, the database will not stop processing a bulk insert if one fails (eg due to duplicate IDs). This makes bulk insert behave similarly to a series of single inserts, except lastError will be set if any insert fails, not just the last one. (new in 1.9.1) + deriving (Show, Eq) + data UpdateOption = Upsert -- ^ If set, the database will insert the supplied object into the collection if no matching document is found | MultiUpdate -- ^ If set, the database will update all matching objects in the collection. Otherwise only updates first matching doc @@ -170,24 +176,33 @@ putNotice :: Notice -> RequestId -> Put putNotice notice requestId = do putHeader (nOpcode notice) requestId - putInt32 0 case notice of Insert{..} -> do + putInt32 (iBits iOptions) putCString iFullCollection mapM_ putDocument iDocuments Update{..} -> do + putInt32 0 putCString uFullCollection putInt32 (uBits uOptions) putDocument uSelector putDocument uUpdater Delete{..} -> do + putInt32 0 putCString dFullCollection putInt32 (dBits dOptions) putDocument dSelector KillCursors{..} -> do + putInt32 0 putInt32 $ toEnum (X.length kCursorIds) mapM_ putInt64 kCursorIds +iBit :: InsertOption -> Int32 +iBit KeepGoing = bit 0 + +iBits :: [InsertOption] -> Int32 +iBits = bitOr . map iBit + uBit :: UpdateOption -> Int32 uBit Upsert = bit 0 uBit MultiUpdate = bit 1 @@ -203,7 +218,7 @@ -- ** Request --- | A request is a message that is sent with a 'Reply' returned +-- | A request is a message that is sent with a 'Reply' expected in return data Request = Query { qOptions :: [QueryOption], @@ -224,6 +239,8 @@ | NoCursorTimeout -- ^ The server normally times out idle cursors after 10 minutes to prevent a memory leak in case a client forgets to close a cursor. Set this option to allow a cursor to live forever until it is closed. | AwaitData -- ^ Use with TailableCursor. If we are at the end of the data, block for a while rather than returning no data. After a timeout period, we do return as normal. -- | Exhaust -- ^ Stream the data down full blast in multiple "more" packages, on the assumption that the client will fully read all data queried. Faster when you are pulling a lot of data and know you want to pull it all down. Note: the client is not allowed to not read all the data unless it closes the connection. +-- Exhaust commented out because not compatible with current `Pipeline` implementation + | Partial -- ^ Get partial results from a _mongos_ if some shards are down, instead of throwing an error. deriving (Show, Eq) -- *** Binary format @@ -255,6 +272,7 @@ qBit NoCursorTimeout = bit 4 qBit AwaitData = bit 5 --qBit Exhaust = bit 6 +qBit Partial = bit 7 qBits :: [QueryOption] -> Int32 qBits = bitOr . map qBit diff --git a/Database/MongoDB/Query.hs b/Database/MongoDB/Query.hs --- a/Database/MongoDB/Query.hs +++ b/Database/MongoDB/Query.hs @@ -4,7 +4,7 @@ module Database.MongoDB.Query ( -- * Monad - Action, access, Failure(..), + Action, access, Failure(..), ErrorCode, AccessMode(..), GetLastError, master, slaveOk, accessMode, MonadDB(..), -- * Database @@ -18,17 +18,17 @@ Select(select), -- * Write -- ** Insert - insert, insert_, insertMany, insertMany_, + insert, insert_, insertMany, insertMany_, insertAll, insertAll_, -- ** Update save, replace, repsert, Modifier, modify, -- ** Delete delete, deleteOne, -- * Read -- ** Query - Query(..), QueryOption(NoCursorTimeout), Projector, Limit, Order, BatchSize, + Query(..), QueryOption(NoCursorTimeout, TailableCursor, AwaitData, Partial), Projector, Limit, Order, BatchSize, explain, find, findOne, fetch, count, distinct, -- *** Cursor - Cursor, next, nextN, rest, closeCursor, isCursorClosed, + Cursor, nextBatch, next, nextN, rest, closeCursor, isCursorClosed, -- ** Group Group(..), GroupKey(..), group, -- ** MapReduce @@ -41,7 +41,7 @@ import Prelude as X hiding (lookup) import Data.UString as U (UString, dropWhile, any, tail) import Data.Bson (Document, at, lookup, look, Field(..), (=:), (=?), Label, Value(String,Doc), Javascript, genObjectId) -import Database.MongoDB.Internal.Protocol (Pipe, Notice(..), Request(GetMore), Reply(..), QueryOption(..), ResponseFlag(..), UpdateOption(..), DeleteOption(..), CursorId, FullCollection, Username, Password, pwKey) +import Database.MongoDB.Internal.Protocol (Pipe, Notice(..), Request(GetMore), Reply(..), QueryOption(..), ResponseFlag(..), InsertOption(..), UpdateOption(..), DeleteOption(..), CursorId, FullCollection, Username, Password, pwKey) import qualified Database.MongoDB.Internal.Protocol as P (send, call, Request(Query)) import Database.MongoDB.Internal.Util (MonadIO', loop, liftIOE, true1, (<.>)) import Control.Monad.MVar @@ -85,19 +85,19 @@ -- | Type of reads and writes to perform data AccessMode = - ReadStaleOk -- Read-only action, reading stale data from a slave is OK. - | UnconfirmedWrites -- Read-write action, slave not OK, every write is fire & forget. - | ConfirmWrites GetLastError -- Read-write action, slave not OK, every write is confirmed with getLastError. + ReadStaleOk -- ^ Read-only action, reading stale data from a slave is OK. + | UnconfirmedWrites -- ^ Read-write action, slave not OK, every write is fire & forget. + | ConfirmWrites GetLastError -- ^ Read-write action, slave not OK, every write is confirmed with getLastError. type GetLastError = Document --- ^ Parameters for getLastError command. For example ["w" =: 2] tells the server to wait for the write to reach at least two servers in replica set before acknowledging. See "http://www.mongodb.org/display/DOCS/Last+Error+Commands" for more options. +-- ^ Parameters for getLastError command. For example @[\"w\" =: 2]@ tells the server to wait for the write to reach at least two servers in replica set before acknowledging. See for more options. master :: AccessMode --- ^ @'ConfirmWrites' []@ +-- ^ Same as 'ConfirmWrites' [] master = ConfirmWrites [] slaveOk :: AccessMode --- ^ @'ReadStaleOk'@ +-- ^ Same as 'ReadStaleOk' slaveOk = ReadStaleOk accessMode :: (Monad m) => AccessMode -> Action m a -> Action m a @@ -208,7 +208,7 @@ -- ^ Selects documents in collection that match selector type Selector = Document --- ^ Filter for a query, analogous to the where clause in SQL. @[]@ matches all documents in collection. @[x =: a, y =: b]@ is analogous to @where x = a and y = b@ in SQL. See for full selector syntax. +-- ^ Filter for a query, analogous to the where clause in SQL. @[]@ matches all documents in collection. @[\"x\" =: a, \"y\" =: b]@ is analogous to @where x = a and y = b@ in SQL. See for full selector syntax. whereJS :: Selector -> Javascript -> Selector -- ^ Add Javascript predicate to selector, in which case a document must match both selector and predicate @@ -253,17 +253,29 @@ insert_ col doc = insert col doc >> return () insertMany :: (MonadIO m) => Collection -> [Document] -> Action m [Value] --- ^ Insert documents into collection and return their \"_id\" values, which are created automatically if not supplied -insertMany col docs = do - db <- thisDatabase - docs' <- liftIO $ mapM assignId docs - write (Insert (db <.> col) docs') - mapM (look "_id") docs' +-- ^ Insert documents into collection and return their \"_id\" values, which are created automatically if not supplied. If a document fails to be inserted (eg. due to duplicate key) then remaining docs are aborted, and LastError is set. +insertMany = insert' [] insertMany_ :: (MonadIO m) => Collection -> [Document] -> Action m () -- ^ Same as 'insertMany' except don't return _ids insertMany_ col docs = insertMany col docs >> return () +insertAll :: (MonadIO m) => Collection -> [Document] -> Action m [Value] +-- ^ Insert documents into collection and return their \"_id\" values, which are created automatically if not supplied. If a document fails to be inserted (eg. due to duplicate key) then remaining docs are still inserted. LastError is set if any doc fails, not just last one. +insertAll = insert' [KeepGoing] + +insertAll_ :: (MonadIO m) => Collection -> [Document] -> Action m () +-- ^ Same as 'insertAll' except don't return _ids +insertAll_ col docs = insertAll col docs >> return () + +insert' :: (MonadIO m) => [InsertOption] -> Collection -> [Document] -> Action m [Value] +-- ^ Insert documents into collection and return their \"_id\" values, which are created automatically if not supplied +insert' opts col docs = do + db <- thisDatabase + docs' <- liftIO $ mapM assignId docs + write (Insert (db <.> col) opts docs') + mapM (look "_id") docs' + assignId :: Document -> IO Document -- ^ Assign a unique value to _id field if missing assignId doc = if X.any (("_id" ==) . label) doc @@ -342,13 +354,13 @@ } deriving (Show, Eq) type Projector = Document --- ^ Fields to return, analogous to the select clause in SQL. @[]@ means return whole document (analogous to * in SQL). @[x =: 1, y =: 1]@ means return only @x@ and @y@ fields of each document. @[x =: 0]@ means return all fields except @x@. +-- ^ Fields to return, analogous to the select clause in SQL. @[]@ means return whole document (analogous to * in SQL). @[\"x\" =: 1, \"y\" =: 1]@ means return only @x@ and @y@ fields of each document. @[\"x\" =: 0]@ means return all fields except @x@. type Limit = Word32 -- ^ Maximum number of documents to return, i.e. cursor will close after iterating over this number of documents. 0 means no limit. type Order = Document --- ^ Fields to sort by. Each one is associated with 1 or -1. Eg. @[x =: 1, y =: -1]@ means sort by @x@ ascending then @y@ descending +-- ^ Fields to sort by. Each one is associated with 1 or -1. Eg. @[\"x\" =: 1, \"y\" =: -1]@ means sort by @x@ ascending then @y@ descending type BatchSize = Word32 -- ^ The number of document to return in each batch response from the server. 0 means use Mongo default. @@ -460,6 +472,17 @@ addMVarFinalizer var (closeCursor cursor) return cursor +nextBatch :: (MonadMVar m) => Cursor -> Action m [Document] +-- ^ Return next batch of documents in query result, which will be empty if finished. +nextBatch (Cursor fcol batchSize var) = modifyMVar var $ \dBatch -> do + -- Pre-fetch next batch promise from server and return current batch. + Batch limit cid docs <- fulfill dBatch + dBatch' <- if cid /= 0 then nextBatch' limit cid else return $ return (Batch 0 0 []) + return (dBatch', docs) + where + nextBatch' limit cid = request [] (GetMore fcol batchSize' cid, remLimit) + where (batchSize', remLimit) = batchSizeRemainingLimit batchSize limit + next :: (MonadMVar m) => Cursor -> Action m (Maybe Document) -- ^ Return next document in query result, or Nothing if finished. next (Cursor fcol batchSize var) = modifyMVar var nextState where @@ -470,13 +493,13 @@ case docs of doc : docs' -> do dBatch' <- if null docs' && cid /= 0 - then nextBatch limit cid + then nextBatch' limit cid else return $ return (Batch limit cid docs') return (dBatch', Just doc) [] -> if cid == 0 then return (return $ Batch 0 0 [], Nothing) -- finished else error $ "server returned empty batch but says more results on server" - nextBatch limit cid = request [] (GetMore fcol batchSize' cid, remLimit) + nextBatch' limit cid = request [] (GetMore fcol batchSize' cid, remLimit) where (batchSize', remLimit) = batchSizeRemainingLimit batchSize limit nextN :: (MonadMVar m, Functor m) => Int -> Cursor -> Action m [Document] diff --git a/mongoDB.cabal b/mongoDB.cabal --- a/mongoDB.cabal +++ b/mongoDB.cabal @@ -1,5 +1,5 @@ name: mongoDB -version: 1.0.0 +version: 1.0.1 build-type: Simple license: OtherLicense license-file: LICENSE @@ -58,9 +58,9 @@ include-dirs: hs-source-dirs: . other-modules: -ghc-prof-options: +ghc-prof-options: -auto-all ghc-shared-options: -ghc-options: -Wall -O2 +ghc-options: -Wall hugs-options: nhc98-options: jhc-options: