query now returns a cursor, nextDoc iterates over the cursor

This commit is contained in:
Scott R. Parish 2010-01-16 15:54:39 -06:00
parent 8805ea4a9e
commit f5a946c0e0

View file

@ -2,6 +2,7 @@ module Database.MongoDB
( (
connect, connectOnPort, connect, connectOnPort,
delete, insert, insertMany, query, remove, update, delete, insert, insertMany, query, remove, update,
nextDoc,
Collection, FieldSelector, NumToSkip, NumToReturn, RequestID, Selector, Collection, FieldSelector, NumToSkip, NumToReturn, RequestID, Selector,
Opcode(..), Opcode(..),
QueryOpt(..), QueryOpt(..),
@ -43,6 +44,15 @@ connectOnPort host port = do
nsRef <- newIORef ns nsRef <- newIORef ns
return $ Connection { cHandle = h, cRand = nsRef } return $ Connection { cHandle = h, cRand = nsRef }
data Cursor = Cursor {
curCon :: Connection,
curID :: IORef Int64,
curNumToRet :: Int32,
curCol :: Collection,
curDocBytes :: IORef L.ByteString,
curClosed :: IORef Bool
}
data Opcode data Opcode
= OP_REPLY -- 1 Reply to a client request. responseTo is set = OP_REPLY -- 1 Reply to a client request. responseTo is set
| OP_MSG -- 1000 generic msg command followed by a string | OP_MSG -- 1000 generic msg command followed by a string
@ -135,8 +145,10 @@ insertMany c col docs = do
return reqID return reqID
query :: Connection -> Collection -> [QueryOpt] -> NumToSkip -> NumToReturn -> query :: Connection -> Collection -> [QueryOpt] -> NumToSkip -> NumToReturn ->
Selector -> Maybe FieldSelector -> IO [BSONObject] Selector -> Maybe FieldSelector -> IO Cursor
query c col opts skip ret sel fsel = do query c col opts skip ret sel fsel = do
let h = cHandle c
let body = runPut $ do let body = runPut $ do
putI32 $ fromQueryOpts opts putI32 $ fromQueryOpts opts
putCol col putCol col
@ -147,8 +159,24 @@ query c col opts skip ret sel fsel = do
Nothing -> putNothing Nothing -> putNothing
Just fsel -> put fsel Just fsel -> put fsel
(reqID, msg) <- packMsg c OP_QUERY body (reqID, msg) <- packMsg c OP_QUERY body
L.hPut (cHandle c) msg L.hPut h msg
getReply c reqID
hdr <- getHeader h
assert (OP_REPLY == hOp hdr) $ return ()
assert (hRespTo hdr == reqID) $ return ()
reply <- getReply h
assert (rRespFlags reply == 0) $ return ()
docBytes <- (L.hGet h $ fromIntegral $ hMsgLen hdr - 16 - 20) >>= newIORef
closed <- newIORef False
cid <- newIORef $ rCursorID reply
return $ Cursor {
curCon = c,
curID = cid,
curNumToRet = ret,
curCol = col,
curDocBytes = docBytes,
curClosed = closed
}
update :: Connection -> Collection -> update :: Connection -> Collection ->
[UpdateFlag] -> Selector -> BSONObject -> IO RequestID [UpdateFlag] -> Selector -> BSONObject -> IO RequestID
@ -178,29 +206,76 @@ data Reply = Reply {
rNumReturned :: Int32 rNumReturned :: Int32
} deriving (Show) } deriving (Show)
getReply :: Connection -> RequestID -> IO [BSONObject] getHeader h = do
getReply c reqID = do
let h = cHandle c
hdrBytes <- L.hGet h 16 hdrBytes <- L.hGet h 16
let hdr = flip runGet hdrBytes $ do return $ flip runGet hdrBytes $ do
msgLen <- getI32 msgLen <- getI32
reqID <- getI32 reqID <- getI32
respTo <- getI32 respTo <- getI32
op <- getI32 op <- getI32
return $ Hdr msgLen reqID respTo $ toOpcode op return $ Hdr msgLen reqID respTo $ toOpcode op
assert (OP_REPLY == hOp hdr) $ return ()
assert (hRespTo hdr == reqID) $ return () getReply h = do
replyBytes <- L.hGet h 20 replyBytes <- L.hGet h 20
let reply = flip runGet replyBytes $ do return $ flip runGet replyBytes $ do
respFlags <- getI32 respFlags <- getI32
cursorID <- getI64 cursorID <- getI64
startFrom <- getI32 startFrom <- getI32
numReturned <- getI32 numReturned <- getI32
return $ (Reply respFlags cursorID startFrom numReturned) return $ (Reply respFlags cursorID startFrom numReturned)
nextDoc :: Cursor -> IO (Maybe BSONObject)
nextDoc cur = do
closed <- readIORef $ curClosed cur
case closed of
True -> return Nothing
False -> do
docBytes <- readIORef $ curDocBytes cur
cid <- readIORef $ curID cur
case L.length docBytes of
0 -> if cid == 0
then writeIORef (curClosed cur) True >> return Nothing
else getMore cur
_ -> do
let (doc, docBytes') = getFirstDoc docBytes
writeIORef (curDocBytes cur) docBytes'
return $ Just doc
getFirstDoc docBytes = flip runGet docBytes $ do
doc <- get
docBytes' <- getRemainingLazyByteString
return (doc, docBytes')
getMore :: Cursor -> IO (Maybe BSONObject)
getMore cur = do
let h = cHandle $ curCon cur
cid <- readIORef $ curID cur
let body = runPut $ do
putI32 0
putCol $ curCol cur
putI32 $ curNumToRet cur
putI64 cid
(reqID, msg) <- packMsg (curCon cur) OP_GET_MORE body
L.hPut h msg
hdr <- getHeader h
assert (OP_REPLY == hOp hdr) $ return ()
assert (hRespTo hdr == reqID) $ return ()
reply <- getReply h
assert (rRespFlags reply == 0) $ return () assert (rRespFlags reply == 0) $ return ()
docBytes <- L.hGet h $ fromIntegral $ hMsgLen hdr - 16 - 20 cid <- readIORef (curID cur)
return $ flip runGet docBytes $ do case rCursorID reply of
forM [1 .. rNumReturned reply] $ \_ -> get 0 -> writeIORef (curID cur) 0
ncid -> assert (ncid == cid) $ return ()
docBytes <- (L.hGet h $ fromIntegral $ hMsgLen hdr - 16 - 20)
case L.length docBytes of
0 -> writeIORef (curClosed cur) True >> return Nothing
_ -> do
let (doc, docBytes') = getFirstDoc docBytes
writeIORef (curDocBytes cur) docBytes'
return $ Just doc
putCol col = putByteString (pack col) >> putNull putCol col = putByteString (pack col) >> putNull