diff --git a/Database/MongoDB.hs b/Database/MongoDB.hs new file mode 100644 --- /dev/null +++ b/Database/MongoDB.hs @@ -0,0 +1,628 @@ +{- + +Copyright (C) 2010 Scott R Parish + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +-} + +module Database.MongoDB + ( + -- * Connection + Connection, + connect, connectOnPort, conClose, disconnect, + dropDatabase, + -- * Database + Database, MongoDBCollectionInvalid, + ColCreateOpt(..), + collectionNames, createCollection, dropCollection, validateCollection, + -- * Collection + Collection, FieldSelector, NumToSkip, NumToReturn, Selector, + QueryOpt(..), + UpdateFlag(..), + count, countMatching, delete, insert, insertMany, query, remove, update, + -- * Convience collection operations + find, findOne, quickFind, quickFind', + -- * Cursor + Cursor, + allDocs, allDocs', finish, nextDoc, + ) +where +import Control.Exception +import Control.Monad +import Data.Binary +import Data.Binary.Get +import Data.Binary.Put +import Data.Bits +import Data.ByteString.Char8 hiding (count, find) +import qualified Data.ByteString.Lazy as L +import qualified Data.ByteString.Lazy.UTF8 as L8 +import Data.Int +import Data.IORef +import qualified Data.List as List +import Data.Maybe +import Data.Typeable +import Database.MongoDB.BSON as BSON +import Database.MongoDB.Util +import qualified Network +import Network.Socket hiding (connect, send, sendTo, recv, recvFrom) +import Prelude hiding (getContents) +import System.IO +import System.IO.Unsafe +import System.Random + +-- | A handle to a database connection +data Connection = Connection { cHandle :: Handle, cRand :: IORef [Int] } + +-- | Estabilish a connection to a MongoDB server +connect :: HostName -> IO Connection +connect = flip connectOnPort $ Network.PortNumber 27017 + +-- | Estabilish 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 + r <- newStdGen + let ns = randomRs (fromIntegral (minBound :: Int32), + fromIntegral (maxBound :: Int32)) r + nsRef <- newIORef ns + return $ Connection { cHandle = h, cRand = nsRef } + +-- | Close database connection +conClose :: Connection -> IO () +conClose = hClose . cHandle + +-- | Alias for 'conClose' +disconnect :: Connection -> IO () +disconnect = conClose + +-- | Drop a database. +dropDatabase :: Connection -> Database -> IO () +dropDatabase c db = do + _ <- dbCmd c db $ toBsonDoc [("dropDatabase", toBson (1::Int))] + return () + +-- | Return a list of collections in /Database/. +collectionNames :: Connection -> Database -> IO [Collection] +collectionNames c db = do + docs <- quickFind' c (db ++ ".system.namespaces") BSON.empty + let names = flip List.map docs $ \doc -> + fromBson $ fromJust $ BSON.lookup "name" doc + return $ List.filter (List.notElem '$') names + +data ColCreateOpt = CCOSize Int64 -- ^ Desired initial size for the + -- collection (in bytes). must be + -- less than or equal to + -- 10000000000. For capped + -- collections this size is the max + -- size of the collection. + | CCOCapped Bool -- ^ If 'True', this is a capped collection. + | CCOMax Int64 -- ^ Maximum number of objects if capped. + deriving (Show, Eq) + +colCreateOptToBson :: ColCreateOpt -> (String, BsonValue) +colCreateOptToBson (CCOSize sz) = ("size", toBson sz) +colCreateOptToBson (CCOCapped b) = ("capped", toBson b) +colCreateOptToBson (CCOMax m) = ("max", toBson m) + +-- | Create a new collection in this database. +-- +-- Normally collection creation is automatic. This function should +-- only be needed if you want to specify 'ColCreateOpt's on creation. +-- 'MongoDBCollectionInvalid' is thrown if the collection already +-- exists. +createCollection :: Connection -> Collection -> [ColCreateOpt] -> IO () +createCollection c col opts = do + let db = dbFromCol col + col' = colMinusDB col + dbcols <- collectionNames c db + case col `List.elem` dbcols of + True -> throwColInvalid $ "Collection already exists: " ++ show col + False -> return () + case ".." `List.elem` (List.group col) of + True -> throwColInvalid $ "Collection can't contain \"..\": " ++ show col + False -> return () + case '$' `List.elem` col && + not ("oplog.$mail" `List.isPrefixOf` col' || + "$cmd" `List.isPrefixOf` col') of + True -> throwColInvalid $ "Collection can't contain '$': " ++ show col + False -> return () + case List.head col == '.' || List.last col == '.' of + True -> throwColInvalid $ + "Collection can't start or end with '.': " ++ show col + False -> return () + let cmd = ("create", toBson col') : List.map colCreateOptToBson opts + _ <- dbCmd c db $ toBsonDoc cmd + return () + +-- | Drop a collection. +dropCollection :: Connection -> Collection -> IO () +dropCollection c col = do + let db = dbFromCol col + col' = colMinusDB col + _ <- dbCmd c db $ toBsonDoc [("drop", toBson col')] + return () + +-- | Return a string of validation info about the collection. +-- +-- Example output (note this probably can/will change with different +-- versions of the server): +-- +-- > validate +-- > details: 0x7fe5cc2c1da4 ofs:e7da4 +-- > firstExtent:0:24100 ns:test.foo.bar +-- > lastExtent:0:24100 ns:test.foo.bar +-- > # extents:1 +-- > datasize?:180 nrecords?:5 lastExtentSize:1024 +-- > padding:1 +-- > first extent: +-- > loc:0:24100 xnext:null xprev:null +-- > nsdiag:test.foo.bar +-- > size:1024 firstRecord:0:241e4 lastRecord:0:24280 +-- > 5 objects found, nobj:5 +-- > 260 bytes data w/headers +-- > 180 bytes data wout/headers +-- > deletedList: 0100100000000000000 +-- > deleted: n: 4 size: 588 +-- > nIndexes:1 +-- > test.foo.bar.$_id_ keys:5 +validateCollection :: Connection -> Collection -> IO String +validateCollection c col = do + let db = dbFromCol col + col' = colMinusDB col + res <- dbCmd c db $ toBsonDoc [("validate", toBson col')] + return $ fromBson $ fromJust $ BSON.lookup "result" res + +dbFromCol :: Collection -> Database +dbFromCol = List.takeWhile (/= '.') + +colMinusDB :: Collection -> Collection +colMinusDB = List.tail . List.dropWhile (/= '.') + +dbCmd :: Connection -> Database -> BsonDoc -> IO BsonDoc +dbCmd c db cmd = do + mres <- findOne c (db ++ ".$cmd") cmd + let res = fromJust mres + case fromBson $ fromJust $ BSON.lookup "ok" res :: Int of + 1 -> return () + _ -> throwOpFailure $ "command \"" ++ show cmd ++ "\" failed: " ++ + (fromBson $ fromJust $ BSON.lookup "errmsg" res) + return res + +-- | An Itertaor over the results of a query. Use 'nextDoc' to get each +-- successive result document, or 'allDocs' or 'allDocs'' to get lazy or +-- strict lists of results. +data Cursor = Cursor { + curCon :: Connection, + curID :: IORef Int64, + curNumToRet :: Int32, + curCol :: Collection, + curDocBytes :: IORef L.ByteString, + curClosed :: IORef Bool + } + +data Opcode + = OP_REPLY -- 1 Reply to a client request. responseTo is set + | OP_MSG -- 1000 generic msg command followed by a string + | OP_UPDATE -- 2001 update document + | OP_INSERT -- 2002 insert new document + | OP_GET_BY_OID -- 2003 is this used? + | OP_QUERY -- 2004 query a collection + | OP_GET_MORE -- 2005 Get more data from a query. See Cursors + | OP_DELETE -- 2006 Delete documents + | OP_KILL_CURSORS -- 2007 Tell database client is done with a cursor + deriving (Show, Eq) + +data MongoDBInternalError = MongoDBInternalError String + deriving (Eq, Show, Read) + +mongoDBInternalError :: TyCon +mongoDBInternalError = mkTyCon "Database.MongoDB.MongoDBInternalError" + +instance Typeable MongoDBInternalError where + typeOf _ = mkTyConApp mongoDBInternalError [] + +instance Exception MongoDBInternalError + +data MongoDBCollectionInvalid = MongoDBCollectionInvalid String + deriving (Eq, Show, Read) + +mongoDBCollectionInvalid :: TyCon +mongoDBCollectionInvalid = mkTyCon "Database.MongoDB.MongoDBcollectionInvalid" + +instance Typeable MongoDBCollectionInvalid where + typeOf _ = mkTyConApp mongoDBCollectionInvalid [] + +instance Exception MongoDBCollectionInvalid + +throwColInvalid :: String -> a +throwColInvalid s = throw $ MongoDBCollectionInvalid s + +data MongoDBOperationFailure = MongoDBOperationFailure String + deriving (Eq, Show, Read) + +mongoDBOperationFailure :: TyCon +mongoDBOperationFailure = mkTyCon "Database.MongoDB.MongoDBoperationFailure" + +instance Typeable MongoDBOperationFailure where + typeOf _ = mkTyConApp mongoDBOperationFailure [] + +instance Exception MongoDBOperationFailure + +throwOpFailure :: String -> a +throwOpFailure s = throw $ MongoDBOperationFailure s + +fromOpcode :: Opcode -> Int32 +fromOpcode OP_REPLY = 1 +fromOpcode OP_MSG = 1000 +fromOpcode OP_UPDATE = 2001 +fromOpcode OP_INSERT = 2002 +fromOpcode OP_GET_BY_OID = 2003 +fromOpcode OP_QUERY = 2004 +fromOpcode OP_GET_MORE = 2005 +fromOpcode OP_DELETE = 2006 +fromOpcode OP_KILL_CURSORS = 2007 + +toOpcode :: Int32 -> Opcode +toOpcode 1 = OP_REPLY +toOpcode 1000 = OP_MSG +toOpcode 2001 = OP_UPDATE +toOpcode 2002 = OP_INSERT +toOpcode 2003 = OP_GET_BY_OID +toOpcode 2004 = OP_QUERY +toOpcode 2005 = OP_GET_MORE +toOpcode 2006 = OP_DELETE +toOpcode 2007 = OP_KILL_CURSORS +toOpcode n = throw $ MongoDBInternalError $ "Got unexpected Opcode: " ++ show n + +-- | The name of a database. +type Database = String + +-- | The full collection name. The full collection name is the +-- concatenation of the database name with the collection name, using +-- a @.@ for the concatenation. For example, for the database @foo@ +-- and the collection @bar@, the full collection name is @foo.bar@. +type Collection = String + +-- | A 'BsonDoc' representing restrictions for a query much like the +-- /where/ part of an SQL query. +type Selector = BsonDoc +-- | A list of field names that limits the fields in the returned +-- documents. The list can contains zero or more elements, each of +-- which is the name of a field that should be returned. An empty list +-- means that no limiting is done and all fields are returned. +type FieldSelector = [L8.ByteString] +type RequestID = Int32 +-- | Sets the number of documents to omit - starting from the first +-- document in the resulting dataset - when returning the result of +-- the query. +type NumToSkip = Int32 +-- | This controls how many documents are returned at a time. The +-- cursor works by requesting /NumToReturn/ documents, which are then +-- immediately all transfered over the network; these are held locally +-- until the those /NumToReturn/ are all consumed and then the network +-- will be hit again for the next /NumToReturn/ documents. +-- +-- If the value @0@ is given, the database will choose the number of +-- documents to return. +-- +-- Otherwise choosing a good value is very dependant on the document size +-- and the way the cursor is being used. +type NumToReturn = Int32 + +-- | Options that control the behavior of a 'query' operation. +data QueryOpt = QO_TailableCursor + | QO_SlaveOK + | QO_OpLogReplay + | QO_NoCursorTimeout + deriving (Show) + +fromQueryOpts :: [QueryOpt] -> Int32 +fromQueryOpts opts = List.foldl (.|.) 0 $ fmap toVal opts + where toVal QO_TailableCursor = 2 + toVal QO_SlaveOK = 4 + toVal QO_OpLogReplay = 8 + toVal QO_NoCursorTimeout = 16 + +-- | Options that effect the behavior of a 'update' operation. +data UpdateFlag = UF_Upsert + | UF_Multiupdate + deriving (Show, Enum) + +fromUpdateFlags :: [UpdateFlag] -> Int32 +fromUpdateFlags flags = List.foldl (.|.) 0 $ + flip fmap flags $ (1 `shiftL`) . fromEnum + +-- | Return the number of documents in /Collection/. +count :: Connection -> Collection -> IO Int64 +count c col = countMatching c col BSON.empty + +-- | Return the number of documents in /Collection/ matching /Selector/ +countMatching :: Connection -> Collection -> Selector -> IO Int64 +countMatching c col sel = do + let db = dbFromCol col + col' = colMinusDB col + res <- dbCmd c db $ toBsonDoc [("count", toBson col'), + ("query", BsonObject sel)] + return $ fromBson $ fromJust $ BSON.lookup "n" res + +-- | Delete documents matching /Selector/ from the given /Collection/. +delete :: Connection -> Collection -> Selector -> IO RequestID +delete c col sel = do + let body = runPut $ do + putI32 0 + putCol col + putI32 0 + put sel + (reqID, msg) <- packMsg c OP_DELETE body + L.hPut (cHandle c) msg + return reqID + +-- | An alias for 'delete'. +remove :: Connection -> Collection -> Selector -> IO RequestID +remove = delete + +-- | Insert a single document into /Collection/. +insert :: Connection -> Collection -> BsonDoc -> IO RequestID +insert c col doc = do + let body = runPut $ do + putI32 0 + putCol col + put doc + (reqID, msg) <- packMsg c OP_INSERT body + L.hPut (cHandle c) msg + return reqID + +-- | Insert a list of documents into /Collection/. +insertMany :: Connection -> Collection -> [BsonDoc] -> IO RequestID +insertMany c col docs = do + let body = runPut $ do + putI32 0 + putCol col + forM_ docs put + (reqID, msg) <- packMsg c OP_INSERT body + L.hPut (cHandle c) msg + return reqID + +-- | Open a cursor to find documents. If you need full functionality, +-- see 'query' +find :: Connection -> Collection -> Selector -> IO Cursor +find c col sel = query c col [] 0 0 sel [] + +-- | Query, but only return the first result, if any. +findOne :: Connection -> Collection -> Selector -> IO (Maybe BsonDoc) +findOne c col sel = do + cur <- query c col [] 0 (-1) sel [] + el <- nextDoc cur + finish cur + return el + +-- | Perform a query and return the result as a lazy list. Be sure to +-- understand the comments about using the lazy list given for +-- 'allDocs'. +quickFind :: Connection -> Collection -> Selector -> IO [BsonDoc] +quickFind c col sel = find c col sel >>= allDocs + +-- | Perform a query and return the result as a strict list. +quickFind' :: Connection -> Collection -> Selector -> IO [BsonDoc] +quickFind' c col sel = find c col sel >>= allDocs' + +-- | Open a cursor to find documents in /Collection/ that match +-- /Selector/. See the documentation for each argument's type for +-- information about how it effects the query. +query :: Connection -> Collection -> [QueryOpt] -> NumToSkip -> NumToReturn -> + Selector -> FieldSelector -> IO Cursor +query c col opts nskip ret sel fsel = do + let h = cHandle c + + let body = runPut $ do + putI32 $ fromQueryOpts opts + putCol col + putI32 nskip + putI32 ret + put sel + case fsel of + [] -> putNothing + _ -> put $ toBsonDoc $ List.zip fsel $ repeat $ BsonInt32 1 + (reqID, msg) <- packMsg c OP_QUERY 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 () + 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 documents with /BsonDoc/ in /Collection/ that match /Selector/. +update :: Connection -> Collection -> + [UpdateFlag] -> Selector -> BsonDoc -> IO RequestID +update c col flags sel obj = do + let body = runPut $ do + putI32 0 + putCol col + putI32 $ fromUpdateFlags flags + put sel + put obj + (reqID, msg) <- packMsg c OP_UPDATE body + L.hPut (cHandle c) msg + return reqID + +data Hdr = Hdr { + hMsgLen :: Int32, + -- hReqID :: Int32, + hRespTo :: Int32, + hOp :: Opcode + } deriving (Show) + +data Reply = Reply { + rRespFlags :: Int32, + rCursorID :: Int64 + -- rStartFrom :: Int32, + -- rNumReturned :: Int32 + } deriving (Show) + +getHeader :: Handle -> IO Hdr +getHeader h = do + hdrBytes <- L.hGet h 16 + return $ flip runGet hdrBytes $ do + msgLen <- getI32 + skip 4 -- reqID <- getI32 + respTo <- getI32 + op <- getI32 + return $ Hdr msgLen respTo $ toOpcode op + +getReply :: Handle -> IO Reply +getReply h = do + replyBytes <- L.hGet h 20 + return $ flip runGet replyBytes $ do + respFlags <- getI32 + cursorID <- getI64 + skip 4 -- startFrom <- getI32 + skip 4 -- numReturned <- getI32 + return $ (Reply respFlags cursorID) + + +-- | Return one document or Nothing if there are no more. +-- Automatically closes the curosr when last document is read +nextDoc :: Cursor -> IO (Maybe BsonDoc) +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 + +-- | Return a lazy list of all (of the rest) of the documents in the +-- cursor. This works much like hGetContents--it will lazily read the +-- cursor data out of the database as the list is used. The cursor is +-- automatically closed when the list has been fully read. +-- +-- If you manually finish the cursor before consuming off this list +-- you won't get all the original documents in the cursor. +-- +-- If you don't consume to the end of the list, you must manually +-- close the cursor or you will leak the cursor, which may also leak +-- on the database side. +allDocs :: Cursor -> IO [BsonDoc] +allDocs cur = unsafeInterleaveIO $ do + doc <- nextDoc cur + case doc of + Nothing -> return [] + Just d -> allDocs cur >>= return . (d :) + +-- | Returns a strict list of all (of the rest) of the documents in +-- the cursor. This means that all of the documents will immediately +-- be read out of the database and loaded into memory. +allDocs' :: Cursor -> IO [BsonDoc] +allDocs' cur = do + doc <- nextDoc cur + case doc of + Nothing -> return [] + Just d -> allDocs' cur >>= return . (d :) + +getFirstDoc :: L.ByteString -> (BsonDoc, L.ByteString) +getFirstDoc docBytes = flip runGet docBytes $ do + doc <- get + docBytes' <- getRemainingLazyByteString + return (doc, docBytes') + +getMore :: Cursor -> IO (Maybe BsonDoc) +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 () + case rCursorID reply of + 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 + +-- | Manually close a cursor -- usually not needed if you use +-- 'allDocs', 'allDocs'', or 'nextDoc'. +finish :: Cursor -> IO () +finish cur = do + let h = cHandle $ curCon cur + cid <- readIORef $ curID cur + let body = runPut $ do + putI32 0 + putI32 1 + putI64 cid + (_reqID, msg) <- packMsg (curCon cur) OP_KILL_CURSORS body + L.hPut h msg + writeIORef (curClosed cur) True + return () + +putCol :: Collection -> Put +putCol col = putByteString (pack col) >> putNull + +packMsg :: Connection -> Opcode -> L.ByteString -> IO (RequestID, L.ByteString) +packMsg c op body = do + reqID <- randNum c + let msg = runPut $ do + putI32 $ fromIntegral $ L.length body + 16 + putI32 reqID + putI32 0 + putI32 $ fromOpcode op + putLazyByteString body + return (reqID, msg) + +randNum :: Connection -> IO Int32 +randNum Connection { cRand = nsRef } = atomicModifyIORef nsRef $ \ns -> + (List.tail ns, + fromIntegral $ List.head ns) diff --git a/Database/MongoDB/BSON.hs b/Database/MongoDB/BSON.hs new file mode 100644 --- /dev/null +++ b/Database/MongoDB/BSON.hs @@ -0,0 +1,475 @@ +{- + +Copyright (C) 2010 Scott R Parish + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +-} + +module Database.MongoDB.BSON + ( + -- * Types + BsonValue(..), + BsonDoc(..), + toBsonDoc, + BinarySubType(..), + -- * BsonDoc Operations + empty, lookup, + -- * Conversion + fromBson, toBson + ) +where +import Prelude hiding (lookup) + +import Control.Monad +import Data.Binary +import Data.Binary.Get +import Data.Binary.IEEE754 +import Data.Binary.Put +import Data.ByteString.Char8 as C8 hiding (empty) +import qualified Data.ByteString.Lazy as L +import qualified Data.ByteString.Lazy.UTF8 as L8 +import qualified Data.ByteString.UTF8 as S8 +import Data.Convertible +import Data.Int +import qualified Data.Map as Map +import qualified Data.List as List +import Data.Time.Clock.POSIX +import Data.Typeable +import Database.MongoDB.Util + +-- | BsonValue is the type that can be used as a key in a 'BsonDoc'. +data BsonValue + = BsonDouble Double + | BsonString L8.ByteString + | BsonObject BsonDoc + | BsonArray [BsonValue] + | BsonUndefined + | BsonBinary BinarySubType L.ByteString + | BsonObjectId L.ByteString + | BsonBool !Bool + | BsonDate POSIXTime + | BsonNull + | BsonRegex L8.ByteString String + | BsonSymbol L8.ByteString + | BsonInt32 Int32 + | BsonInt64 Int64 + | BsonMinKey + | BsonMaxKey + deriving (Show, Eq, Ord) + +instance Typeable BsonValue where + typeOf _ = mkTypeName "BsonValue" + +-- | BSON Document: this is the top-level (but recursive) type that +-- all MongoDB collections work in terms of. It is a mapping between +-- strings ('Data.ByteString.Lazu.UTF8.ByteString') and 'BsonValue's. +-- It can be constructed either from a 'Map' (eg @'BsonDoc' myMap@) or +-- from a associative list (eg @'toBsonDoc' myAL@). +newtype BsonDoc = BsonDoc { + bdFromBsonDoc :: Map.Map L8.ByteString BsonValue + } + deriving (Eq, Ord, Show) + +class BsonDocOps a where + -- | Construct a BsonDoc from an associative list + toBsonDoc :: [(a, BsonValue)] -> BsonDoc + -- | Unwrap BsonDoc to be a Map + fromBsonDoc :: BsonDoc -> Map.Map a BsonValue + -- | Return the BsonValue for given key, if any. + lookup :: a -> BsonDoc -> Maybe BsonValue + +-- | An empty BsonDoc +empty :: BsonDoc +empty = BsonDoc Map.empty + +instance BsonDocOps L8.ByteString where + toBsonDoc = BsonDoc . Map.fromList + fromBsonDoc = bdFromBsonDoc + lookup k = Map.lookup k . fromBsonDoc + +instance BsonDocOps String where + toBsonDoc = BsonDoc . Map.mapKeys L8.fromString .Map.fromList + fromBsonDoc = Map.mapKeys L8.toString . bdFromBsonDoc + lookup k = Map.lookup (L8.fromString k) . fromBsonDoc + +data DataType = + Data_min_key | -- -1 + Data_number | -- 1 + Data_string | -- 2 + Data_object | -- 3 + Data_array | -- 4 + Data_binary | -- 5 + Data_undefined | -- 6 + Data_oid | -- 7 + Data_boolean | -- 8 + Data_date | -- 9 + Data_null | -- 10 + Data_regex | -- 11 + Data_ref | -- 12 + Data_code | -- 13 + Data_symbol | -- 14 + Data_code_w_scope | -- 15 + Data_int | -- 16 + Data_timestamp | -- 17 + Data_long | -- 18 + Data_max_key -- 127 + deriving (Show, Read, Enum, Eq, Ord) + +toDataType :: Int -> DataType +toDataType (-1) = Data_min_key +toDataType 127 = Data_max_key +toDataType d = toEnum d + +fromDataType :: DataType -> Int +fromDataType Data_min_key = (-1) +fromDataType Data_max_key = 127 +fromDataType d = fromEnum d + +data BinarySubType = + BSTUNDEFINED_1 | + BSTFunction | -- 1 + BSTByteArray | -- 2 + BSTUUID | -- 3 + BSTUNDEFINED_2 | + BSTMD5 | -- 5 + BSTUserDefined + deriving (Show, Read, Enum, Eq, Ord) + +toBinarySubType :: Int -> BinarySubType +toBinarySubType 0x80 = BSTUserDefined +toBinarySubType d = toEnum d + +fromBinarySubType :: BinarySubType -> Int +fromBinarySubType BSTUserDefined = 0x80 +fromBinarySubType d = fromEnum d + +instance Binary BsonDoc where + get = liftM snd getDoc + put = putObj + +getVal :: DataType -> Get (Integer, BsonValue) +getVal Data_number = getFloat64le >>= return . (,) 8 . BsonDouble +getVal Data_string = do + sLen1 <- getI32 + (_sLen2, s) <- getS + return (fromIntegral $ 4 + sLen1, BsonString s) +getVal Data_object = getDoc >>= \(len, obj) -> return (len, BsonObject obj) +getVal Data_array = do + (len, arr) <- getRawObj + let arr2 = Map.fold (:) [] arr -- reverse and remove key + return (len, BsonArray arr2) +getVal Data_binary = do + skip 4 + st <- getI8 + len2 <- getI32 + bs <- getLazyByteString $ fromIntegral len2 + return (4 + 1 + 4 + fromIntegral len2, BsonBinary (toBinarySubType st) bs) +getVal Data_undefined = return (1, BsonUndefined) +getVal Data_oid = getLazyByteString 12 >>= return . (,) 12 . BsonObjectId +getVal Data_boolean = + getI8 >>= return . (,) (1::Integer) . BsonBool . (/= (0::Int)) +getVal Data_date = + getI64 >>= return . (,) 8 . BsonDate . flip (/) 1000 . realToFrac +getVal Data_null = return (1, BsonNull) +getVal Data_regex = fail "Data_code not yet supported" -- TODO +getVal Data_ref = fail "Data_ref is deprecated" +getVal Data_code = fail "Data_code not yet supported" -- TODO +getVal Data_symbol = do + sLen1 <- getI32 + (_sLen2, s) <- getS + return (fromIntegral $ 4 + sLen1, BsonString s) +getVal Data_code_w_scope = fail "Data_code_w_scope not yet supported" -- TODO +getVal Data_int = getI32 >>= return . (,) 4 . BsonInt32 . fromIntegral +getVal Data_timestamp = fail "Data_timestamp not yet supported" -- TODO + +getVal Data_long = getI64 >>= return . (,) 8 . BsonInt64 +getVal Data_min_key = return (0, BsonMinKey) +getVal Data_max_key = return (0, BsonMaxKey) + +getInnerObj :: Int32 -> Get (Map.Map L8.ByteString BsonValue) + -> Get (Map.Map L8.ByteString BsonValue) +getInnerObj 1 obj = obj +getInnerObj bytesLeft obj = do + typ <- getDataType + (keySz, key) <- getS + (valSz, val) <- getVal typ + getInnerObj (bytesLeft - 1 - fromIntegral keySz - fromIntegral valSz) $ + liftM (Map.insert key val) obj + +getRawObj :: Get (Integer, Map.Map L8.ByteString BsonValue) +getRawObj = do + bytes <- getI32 + obj <- getInnerObj (bytes - 4) $ return Map.empty + getNull + return (fromIntegral bytes, obj) + +getDoc :: Get (Integer, BsonDoc) +getDoc = getRawObj >>= \(len, obj) -> return (len, BsonDoc obj) + +getDataType :: Get DataType +getDataType = liftM toDataType getI8 + +putType :: BsonValue -> Put +putType BsonDouble{} = putDataType Data_number +putType BsonString{} = putDataType Data_string +putType BsonObject{} = putDataType Data_object +putType BsonArray{} = putDataType Data_array +putType BsonBinary{} = putDataType Data_binary +putType BsonUndefined = putDataType Data_undefined +putType BsonObjectId{} = putDataType Data_oid +putType BsonBool{} = putDataType Data_boolean +putType BsonDate{} = putDataType Data_date +putType BsonNull = putDataType Data_null +putType BsonRegex{} = putDataType Data_regex +-- putType = putDataType Data_ref +-- putType = putDataType Data_code +putType BsonSymbol{} = putDataType Data_symbol +-- putType = putDataType Data_code_w_scope +putType BsonInt32 {} = putDataType Data_int +putType BsonInt64 {} = putDataType Data_long +-- putType = putDataType Data_timestamp +putType BsonMinKey = putDataType Data_min_key +putType BsonMaxKey = putDataType Data_max_key + +putVal :: BsonValue -> Put +putVal (BsonDouble d) = putFloat64le d +putVal (BsonString s) = putI32 (fromIntegral $ 1 + L8.length s) >> putS s +putVal (BsonObject o) = putObj o +putVal (BsonArray es) = putOutterObj bs + where bs = runPut $ forM_ (List.zip [(0::Int) .. ] es) $ \(i, e) -> + putType e >> (putS $ L8.fromString $ show i) >> putVal e +putVal (BsonBinary t bs)= do putI32 $ fromIntegral $ 4 + L.length bs + putI8 $ fromBinarySubType t + putI32 $ fromIntegral $ L.length bs + putLazyByteString bs +putVal BsonUndefined = putNothing +putVal (BsonObjectId o) = putLazyByteString o +putVal (BsonBool False) = putI8 (0::Int) +putVal (BsonBool True) = putI8 (1::Int) +putVal (BsonDate pt) = putI64 $ round $ 1000 * (realToFrac pt :: Double) +putVal BsonNull = putNothing +putVal (BsonRegex r opt)= do putS r + putByteString $ pack $ List.sort opt + putNull +putVal (BsonSymbol s) = putI32 (fromIntegral $ 1 + L8.length s) >> putS s +putVal (BsonInt32 i) = putI32 i +putVal (BsonInt64 i) = putI64 i +putVal BsonMinKey = putNothing +putVal BsonMaxKey = putNothing + +putObj :: BsonDoc -> Put +putObj obj = putOutterObj bs + where bs = runPut $ forM_ (Map.toList (fromBsonDoc obj)) $ \(k, v) -> + putType v >> putS k >> putVal v + +putOutterObj :: L.ByteString -> Put +putOutterObj bytes = do + -- the length prefix and null term are included in the length + putI32 $ fromIntegral $ 4 + 1 + L.length bytes + putLazyByteString bytes + putNull + +putDataType :: DataType -> Put +putDataType = putI8 . fromDataType + +class BsonConv a b where + -- | Convert a BsonValue into a native Haskell type. + fromBson :: Convertible a b => a -> b + -- | Convert a native Haskell type into a BsonValue. + toBson :: Convertible b a => b -> a + +instance BsonConv BsonValue a where + fromBson = convert + toBson = convert + +instance BsonConv (Maybe BsonValue) (Maybe a) where + fromBson = convert + toBson = convert + +unsupportedError :: (Typeable a, Convertible BsonValue a) => + BsonValue -> ConvertResult a +unsupportedError = convError "Unsupported conversion" + +instance Convertible Double BsonValue where + safeConvert = return . BsonDouble + +instance Convertible Float BsonValue where + safeConvert = return . BsonDouble . realToFrac + +instance Convertible String BsonValue where + safeConvert = return . BsonString . L8.fromString + +instance Convertible L8.ByteString BsonValue where + safeConvert = return . BsonString + +instance Convertible S8.ByteString BsonValue where + safeConvert = return . BsonString . L.fromChunks . return + +instance Convertible [Double] BsonValue where + safeConvert ds = BsonArray `liftM` mapM safeConvert ds + +instance Convertible [Float] BsonValue where + safeConvert fs = BsonArray `liftM` mapM safeConvert fs + +instance Convertible [String] BsonValue where + safeConvert ss = BsonArray `liftM` mapM safeConvert ss + +instance Convertible [L8.ByteString] BsonValue where + safeConvert bs = BsonArray `liftM` mapM safeConvert bs + +instance Convertible [S8.ByteString] BsonValue where + safeConvert bs = BsonArray `liftM` mapM safeConvert bs + +instance Convertible [Bool] BsonValue where + safeConvert bs = BsonArray `liftM` mapM safeConvert bs + +instance Convertible [POSIXTime] BsonValue where + safeConvert ts = BsonArray `liftM` mapM safeConvert ts + +instance Convertible [Int] BsonValue where + safeConvert is = BsonArray `liftM` mapM safeConvert is + +instance Convertible [Integer] BsonValue where + safeConvert is = BsonArray `liftM` mapM safeConvert is + +instance Convertible [Int32] BsonValue where + safeConvert is = BsonArray `liftM` mapM safeConvert is + +instance Convertible [Int64] BsonValue where + safeConvert is = BsonArray `liftM` mapM safeConvert is + +instance Convertible POSIXTime BsonValue where + safeConvert = return . BsonDate + +instance Convertible Bool BsonValue where + safeConvert = return . BsonBool + +instance Convertible Int BsonValue where + safeConvert i = if i >= (fromIntegral (minBound::Int32)) && + i <= (fromIntegral (maxBound::Int32)) + then return $ BsonInt32 $ fromIntegral i + else return $ BsonInt64 $ fromIntegral i + +instance Convertible Integer BsonValue where + safeConvert i = if i >= (fromIntegral (minBound::Int32)) && + i <= (fromIntegral (maxBound::Int32)) + then return $ BsonInt32 $ fromIntegral i + else return $ BsonInt64 $ fromIntegral i + +instance Convertible Int32 BsonValue where + safeConvert i = return $ BsonInt32 i + +instance Convertible Int64 BsonValue where + safeConvert i = return $ BsonInt64 i + +instance (Convertible a BsonValue) => + Convertible (Maybe a) BsonValue where + safeConvert Nothing = return BsonNull + safeConvert (Just a) = safeConvert a + +instance Convertible BsonValue Double where + safeConvert (BsonDouble d) = return d + safeConvert (BsonInt32 i) = safeConvert i + safeConvert (BsonInt64 i) = safeConvert i + safeConvert v = unsupportedError v + +instance Convertible BsonValue Float where + safeConvert (BsonDouble d) = safeConvert d + safeConvert (BsonInt32 i) = safeConvert i + safeConvert (BsonInt64 i) = safeConvert i + safeConvert v = unsupportedError v + +instance Convertible BsonValue String where + safeConvert (BsonString bs) = return $ L8.toString bs + safeConvert v = unsupportedError v + +instance Convertible BsonValue L8.ByteString where + safeConvert (BsonString bs) = return bs + safeConvert v = unsupportedError v + +instance Convertible BsonValue S8.ByteString where + safeConvert (BsonString bs) = return $ C8.concat $ L.toChunks bs + safeConvert v = unsupportedError v + +instance Convertible BsonValue [Double] where + safeConvert (BsonArray a) = mapM safeConvert a + safeConvert v = unsupportedError v + +instance Convertible BsonValue [Float] where + safeConvert (BsonArray a) = mapM safeConvert a + safeConvert v = unsupportedError v + +instance Convertible BsonValue [String] where + safeConvert (BsonArray a) = mapM safeConvert a + safeConvert v = unsupportedError v + +instance Convertible BsonValue [Bool] where + safeConvert (BsonArray a) = mapM safeConvert a + safeConvert v = unsupportedError v + +instance Convertible BsonValue [POSIXTime] where + safeConvert (BsonArray a) = mapM safeConvert a + safeConvert v = unsupportedError v + +instance Convertible BsonValue [Int32] where + safeConvert (BsonArray a) = mapM safeConvert a + safeConvert v = unsupportedError v + +instance Convertible BsonValue [Int64] where + safeConvert (BsonArray a) = mapM safeConvert a + safeConvert v = unsupportedError v + +instance Convertible BsonValue Bool where + safeConvert (BsonBool b) = return b + safeConvert v = unsupportedError v + +instance Convertible BsonValue POSIXTime where + safeConvert (BsonDate t) = return t + safeConvert v = unsupportedError v + +instance Convertible BsonValue Int where + safeConvert (BsonDouble d) = safeConvert d + safeConvert (BsonInt32 d) = safeConvert d + safeConvert (BsonInt64 d) = safeConvert d + safeConvert v = unsupportedError v + +instance Convertible BsonValue Integer where + safeConvert (BsonDouble d) = safeConvert d + safeConvert (BsonInt32 d) = safeConvert d + safeConvert (BsonInt64 d) = safeConvert d + safeConvert v = unsupportedError v + +instance Convertible BsonValue Int32 where + safeConvert (BsonDouble d) = safeConvert d + safeConvert (BsonInt32 d) = return d + safeConvert (BsonInt64 d) = safeConvert d + safeConvert v = unsupportedError v + +instance Convertible BsonValue Int64 where + safeConvert (BsonDouble d) = safeConvert d + safeConvert (BsonInt32 d) = safeConvert d + safeConvert (BsonInt64 d) = return d + safeConvert v = unsupportedError v + +instance (Convertible BsonValue a) => + Convertible (Maybe BsonValue) (Maybe a) where + safeConvert Nothing = return Nothing + safeConvert (Just a) = liftM Just $ safeConvert a diff --git a/Database/MongoDB/Util.hs b/Database/MongoDB/Util.hs new file mode 100644 --- /dev/null +++ b/Database/MongoDB/Util.hs @@ -0,0 +1,77 @@ +{- + +Copyright (C) 2010 Scott R Parish + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +-} + +module Database.MongoDB.Util + ( + putI8, putI32, putI64, putNothing, putNull, putS, + getI8, getI32, getI64, getC, getS, getNull, + ) +where +import Control.Exception (assert) +import Control.Monad +import Data.Binary +import Data.Binary.Get +import Data.Binary.Put +import Data.ByteString.Char8 +import qualified Data.ByteString.Lazy as L +import qualified Data.ByteString.Lazy.UTF8 as L8 +import Data.Char (chr) +import Data.Int + +getC :: Get Char +getC = liftM chr getI8 + +getI8 :: (Integral a) => Get a +getI8 = liftM fromIntegral getWord8 + +getI32 :: Get Int32 +getI32 = liftM fromIntegral getWord32le + +getI64 :: Get Int64 +getI64 = liftM fromIntegral getWord64le + +getS :: Get (Integer, L8.ByteString) +getS = getLazyByteStringNul >>= \s -> return (fromIntegral $ L.length s + 1, s) + +getNull :: Get () +getNull = do {c <- getC; assert (c == '\0') $ return ()} + +putI8 :: (Integral i) => i -> Put +putI8 = putWord8 . fromIntegral + +putI32 :: Int32 -> Put +putI32 = putWord32le . fromIntegral + +putI64 :: Int64 -> Put +putI64 = putWord64le . fromIntegral + +putNothing :: Put +putNothing = putByteString $ pack "" + +putNull :: Put +putNull = putI8 (0::Int) + +putS :: L8.ByteString -> Put +putS s = putLazyByteString s >> putNull diff --git a/Setup.lhs b/Setup.lhs new file mode 100644 --- /dev/null +++ b/Setup.lhs @@ -0,0 +1,3 @@ +#!/usr/bin/env runhaskell +> import Distribution.Simple +> main = defaultMain diff --git a/mongoDB.cabal b/mongoDB.cabal new file mode 100644 --- /dev/null +++ b/mongoDB.cabal @@ -0,0 +1,31 @@ +Name: mongoDB +Version: 0.1 +License: MIT +Maintainer: Scott Parish +Author: Scott Parish +Copyright: Copyright (c) 2010-2010 Scott Parish +Category: Database +Synopsis: A driver for MongoDB +Description: This driver lets you connect to MongoDB, do inserts, + queries, updates, etc. Also has many convience functions + inspired by HDBC such as more easily converting between + the BsonValue types and native Haskell types. +Stability: alpha +Build-Depends: base < 5, + binary, + bytestring, + containers, + convertible, + data-binary-ieee754, + network, + random, + time, + utf8-string +Build-Type: Simple +Exposed-modules: Database.MongoDB, + Database.MongoDB.BSON +Other-modules: Database.MongoDB.Util +ghc-options: -Wall +extensions: FlexibleContexts, FlexibleInstances, MultiParamTypeClasses, + TypeSynonymInstances +cabal-version: >= 1.4