{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE RecordWildCards #-}
module Database.LevelDB.Higher
(
Key, Value, Item, KeySpace, KeySpaceId
, get, put, delete
, runBatch, putB, deleteB
, scan, ScanQuery(..), queryItems, queryList, queryBegins, queryCount
, withKeySpace, withOptions, withSnapshot
, forkLevelDB
, MonadLevelDB(..), LevelDBT, LevelDB
, mapLevelDBT
, runLevelDB, runLevelDB', runCreateLevelDB
, runResourceT
, Options(..), ReadOptions(..), WriteOptions(..), RWOptions
, WriteBatch, def
, MonadThrow, MonadResourceBase
) where
import Control.Monad.Reader
import Control.Monad.Writer
import Data.Word (Word32)
#if !MIN_VERSION_base(4,8,0)
import Control.Applicative (Applicative)
#endif
import Control.Monad.Base (MonadBase(..))
import Control.Concurrent.MVar.Lifted
import Control.Concurrent (ThreadId)
import qualified Data.ByteString as BS
import Data.ByteString (ByteString)
import Data.Serialize (encode, decode)
import Data.Default (def)
import qualified Database.LevelDB as LDB
import Database.LevelDB
hiding (put, get, delete, write, withSnapshot)
import Control.Monad.Trans.Resource
import Control.Monad.Trans.Control
import Control.Monad.Catch (MonadCatch (..)
, MonadMask (..))
#if MIN_VERSION_mtl(2,2,1)
import qualified Control.Monad.Except as Except
#else
import qualified Control.Monad.Trans.Error as Error
#endif
import qualified Control.Monad.Trans.Cont as Cont
import qualified Control.Monad.Trans.Identity as Identity
import qualified Control.Monad.Trans.List as List
import qualified Control.Monad.Trans.Maybe as Maybe
import qualified Control.Monad.Trans.State as State
import qualified Control.Monad.Trans.Writer as Writer
import qualified Control.Monad.Trans.RWS as RWS
import qualified Control.Monad.Trans.RWS.Strict as Strict
import qualified Control.Monad.Trans.State.Strict as Strict
import qualified Control.Monad.Trans.Writer.Strict as Strict
type Key = ByteString
type Value = ByteString
type KeySpace = ByteString
type KeySpaceId = ByteString
type Item = (Key, Value)
type RWOptions = (ReadOptions, WriteOptions)
data DBContext = DBC { dbcDb :: DB
, dbcKsId :: KeySpaceId
, dbcSyncMV :: MVar Word32
, dbcRWOptions :: RWOptions
}
instance Show DBContext where
show = (<>) "KeySpaceID: " . show . dbcKsId
newtype LevelDBT m a
= LevelDBT { unLevelDBT :: ReaderT DBContext (ResourceT m) a }
deriving ( Functor, Applicative, Monad, MonadIO, MonadThrow)
instance (MonadBase b m) => MonadBase b (LevelDBT m) where
liftBase = lift . liftBase
instance MonadTrans LevelDBT where
lift = LevelDBT . lift . lift
instance (MonadResourceBase m) => MonadResource (LevelDBT m) where
liftResourceT = LevelDBT . liftResourceT
instance MonadCatch m => MonadCatch (LevelDBT m) where
catch (LevelDBT m) c =
LevelDBT . ReaderT $ \r -> runReaderT m r `catch` \e -> runReaderT (unLevelDBT (c e)) r
instance MonadMask m => MonadMask (LevelDBT m) where
mask a = LevelDBT . ReaderT $ \e -> mask $ \u -> runReaderT (unLevelDBT (a $ q u)) e
where
q :: (ResourceT m a -> ResourceT m a) -> LevelDBT m a -> LevelDBT m a
q u (LevelDBT (ReaderT b)) =
LevelDBT $ ReaderT (u . b)
uninterruptibleMask a =
LevelDBT . ReaderT $ \e -> uninterruptibleMask $ \u -> runReaderT (unLevelDBT (a $ q u)) e
where
q :: (ResourceT m a -> ResourceT m a) -> LevelDBT m a -> LevelDBT m a
q u (LevelDBT (ReaderT b)) =
LevelDBT $ ReaderT (u . b)
#if MIN_VERSION_monad_control(1,0,0)
instance MonadTransControl LevelDBT where
type StT LevelDBT a = StT ResourceT (StT (ReaderT DBContext) a)
liftWith f =
LevelDBT $ liftWith $ \run ->
liftWith $ \run' ->
f $ run' . run . unLevelDBT
restoreT = LevelDBT . restoreT . restoreT
instance (MonadBaseControl b m) => MonadBaseControl b (LevelDBT m) where
type StM (LevelDBT m) a = ComposeSt LevelDBT m a
liftBaseWith = defaultLiftBaseWith
restoreM = defaultRestoreM
#else
instance MonadTransControl LevelDBT where
newtype StT LevelDBT a = StLevelDBT
{unStLevelDBT :: StT ResourceT (StT (ReaderT DBContext) a) }
liftWith f =
LevelDBT $ liftWith $ \run ->
liftWith $ \run' ->
f $ liftM StLevelDBT . run' . run . unLevelDBT
restoreT = LevelDBT . restoreT . restoreT . liftM unStLevelDBT
instance (MonadBaseControl b m) => MonadBaseControl b (LevelDBT m) where
newtype StM (LevelDBT m) a = StMT {unStMT :: ComposeSt LevelDBT m a}
liftBaseWith = defaultLiftBaseWith StMT
restoreM = defaultRestoreM unStMT
#endif
class ( Monad m
, MonadThrow m
, MonadIO m
, Applicative m
, MonadResource m
, MonadBase IO m )
=> MonadLevelDB m where
withDBContext :: (DBContext -> DBContext) -> m a -> m a
liftLevelDB :: LevelDBT IO a -> m a
instance (MonadResourceBase m) => MonadLevelDB (LevelDBT m) where
liftLevelDB = mapLevelDBT liftIO
withDBContext = localLDB
#define INST(M,T, F) \
instance (M, MonadLevelDB m) => MonadLevelDB (T m) \
where \
liftLevelDB = lift . liftLevelDB ; \
withDBContext f = F (withDBContext f) ; \
INST(Monad m,ReaderT r, mapReaderT)
INST(Monad m,Maybe.MaybeT, Maybe.mapMaybeT)
INST(Monad m,Identity.IdentityT, Identity.mapIdentityT)
INST(Monad m,List.ListT, List.mapListT)
INST(Monad m,Cont.ContT r, Cont.mapContT)
INST(Monad m,State.StateT s, State.mapStateT )
INST(Monad m,Strict.StateT s, Strict.mapStateT )
INST(Monoid w, Writer.WriterT w, Writer.mapWriterT)
INST(Monoid w, Strict.WriterT w, Strict.mapWriterT)
INST(Monoid w, RWS.RWST r w s, RWS.mapRWST)
INST(Monoid w, Strict.RWST r w s, Strict.mapRWST)
#if MIN_VERSION_mtl(2,2,1)
INST(Monad m, Except.ExceptT e, Except.mapExceptT)
#else
INST(Error.Error e, Error.ErrorT e, Error.mapErrorT)
#endif
#undef INST
type LevelDB a = LevelDBT IO a
runLevelDB :: (MonadResourceBase m)
=> FilePath
-> Options
-> RWOptions
-> KeySpace
-> LevelDBT m a
-> m a
runLevelDB path dbopt rwopt ks ma = runResourceT $ runLevelDB' path dbopt rwopt ks ma
runLevelDB' :: (MonadResourceBase m)
=> FilePath
-> Options
-> RWOptions
-> KeySpace
-> LevelDBT m a
-> ResourceT m a
runLevelDB' path dbopt rwopt ks ma = do
db <- openDB
mv <- newMVar 0
ksId <- withSystemContext db mv $ getKeySpaceId ks
runReaderT (unLevelDBT ma) (DBC db ksId mv rwopt)
where
openDB = LDB.open path dbopt
withSystemContext db mv sctx =
runReaderT (unLevelDBT sctx) $ DBC db systemKeySpaceId mv rwopt
runCreateLevelDB :: (MonadResourceBase m)
=> FilePath
-> KeySpace
-> LevelDBT m a
-> m a
runCreateLevelDB path = runLevelDB path def{createIfMissing=True} def
forkLevelDB :: (MonadLevelDB m)
=> LevelDB ()
-> m ThreadId
forkLevelDB ma = liftLevelDB $ LevelDBT $
mapReaderT resourceForkIO (unLevelDBT ma)
withKeySpace :: (MonadLevelDB m) => KeySpace -> m a -> m a
withKeySpace ks ma = do
ksId <- getKeySpaceId ks
withDBContext (\dbc -> dbc { dbcKsId = ksId}) ma
withOptions :: (MonadLevelDB m) => RWOptions -> m a -> m a
withOptions opts =
withDBContext (\dbc -> dbc { dbcRWOptions = opts })
withSnapshot :: (MonadLevelDB m) => m a -> m a
withSnapshot ma = do
(db, _, _) <- getDB
LDB.withSnapshot db $ \ss ->
withDBContext (\dbc -> dbc {dbcRWOptions = setSnap dbc ss}) ma
where
setSnap dbc ss =
let (ropts, wopts) = dbcRWOptions dbc in
(ropts {useSnapshot = Just ss}, wopts)
put :: (MonadLevelDB m) => Key -> Value -> m ()
put k v = do
(db, ksId, (_, wopt)) <- getDB
let packed = ksId <> k
LDB.put db wopt packed v
get :: (MonadLevelDB m) => Key -> m (Maybe Value)
get k = do
(db, ksId, (ropt, _)) <- getDB
let packed = ksId <> k
LDB.get db ropt packed
delete :: (MonadLevelDB m) => Key -> m ()
delete k = do
(db, ksId, (_, wopt)) <- getDB
let packed = ksId <> k
LDB.delete db wopt packed
runBatch :: (MonadLevelDB m)
=> WriterT WriteBatch m ()
-> m ()
runBatch wb = do
(db, _, (_, wopt)) <- getDB
(_, ops) <- runWriterT wb
LDB.write db wopt ops
putB :: (MonadLevelDB m) => Key -> Value -> WriterT WriteBatch m ()
putB k v = do
(_, ksId, _) <- getDB
tell [Put (ksId <> k) v]
return ()
deleteB :: (MonadLevelDB m) => Key -> WriterT WriteBatch m ()
deleteB k = do
(_, ksId, _) <- getDB
tell [Del (ksId <> k)]
return ()
scan :: (MonadLevelDB m)
=> Key
-> ScanQuery a b
-> m b
scan k ScanQuery{..} = do
(db, ksId, (ropt,_)) <- getDB
withIterator db ropt $ doScan (ksId <> k) ksId
where
doScan prefix ksId iter = do
iterSeek iter prefix
applyIterate scanInit
where
readItem = do
nk <- iterKey iter
nv <- iterValue iter
return $
if sameKsId nk then (fmap (BS.drop 4) nk, nv)
else (Nothing, Nothing)
applyIterate acc = do
item <- readItem
case item of
(Just nk, Just nv) ->
if scanWhile k (nk, nv) acc then do
iterNext iter
items <- applyIterate acc
return $ if scanFilter (nk, nv) then
scanFold (scanMap (nk, nv)) items
else items
else return acc
_ -> return acc
sameKsId Nothing = False
sameKsId (Just nk) = BS.take 4 nk == ksId
data ScanQuery a b = ScanQuery {
scanInit :: b
, scanWhile :: Key -> Item -> b -> Bool
, scanMap :: Item -> a
, scanFilter :: Item -> Bool
, scanFold :: a -> b -> b
}
queryBegins :: ScanQuery a b
queryBegins = ScanQuery
{ scanWhile = \ prefix (nk, _) _ ->
BS.length nk >= BS.length prefix
&& BS.take (BS.length prefix) nk == prefix
, scanInit = error "No scanInit provided."
, scanMap = error "No scanMap provided."
, scanFilter = const True
, scanFold = error "No scanFold provided."
}
queryItems :: ScanQuery Item [Item]
queryItems = queryBegins { scanInit = []
, scanMap = id
, scanFold = (:)
}
queryList :: ScanQuery a [a]
queryList = queryBegins { scanInit = []
, scanFilter = const True
, scanFold = (:)
}
queryCount :: (Num a) => ScanQuery a a
queryCount = queryBegins { scanInit = 0
, scanMap = const 1
, scanFold = (+) }
mapLevelDBT :: (m a -> n b) -> LevelDBT m a -> LevelDBT n b
mapLevelDBT f ma = LevelDBT $
mapReaderT (transResourceT f) $ unLevelDBT ma
getDB :: (MonadLevelDB m) => m (DB, KeySpaceId, RWOptions)
getDB = liftLevelDB $ asksLDB (\dbc ->
(dbcDb dbc, dbcKsId dbc, dbcRWOptions dbc))
asksLDB :: (MonadResourceBase m) => (DBContext -> a) -> LevelDBT m a
asksLDB = LevelDBT . asks
localLDB :: (MonadResourceBase m)
=> (DBContext -> DBContext)
-> LevelDBT m a -> LevelDBT m a
localLDB f ma = LevelDBT $ local f (unLevelDBT ma)
defaultKeySpaceId :: KeySpaceId
defaultKeySpaceId = "\0\0\0\0"
systemKeySpaceId :: KeySpaceId
systemKeySpaceId = "\0\0\0\1"
systemKeySpace :: KeySpace
systemKeySpace = "system"
getKeySpaceId :: (MonadLevelDB m) => KeySpace -> m KeySpaceId
getKeySpaceId ks
| ks == "" = return defaultKeySpaceId
| ks == systemKeySpace = return systemKeySpaceId
| otherwise = liftLevelDB $ withKeySpace systemKeySpace $ do
findKS <- get $ "keyspace:" <> ks
case findKS of
(Just foundId) -> return foundId
Nothing -> do
nextId <- incr "max-keyspace-id"
put ("keyspace:" <> ks) nextId
return nextId
where
incr k = do
mv <- takeMVarDBC
curId <- case mv of
0 -> initKeySpaceIdMV k >> takeMVarDBC
n -> return n
let nextId = curId + 1
put k $ encode nextId
putMVarDBC nextId
return $ encode curId
initKeySpaceIdMV k = do
findMaxId <- get k
case findMaxId of
(Just found) -> putMVarDBC $ decodeKsId found
Nothing -> putMVarDBC 2
putMVarDBC v = asksLDB dbcSyncMV >>= flip putMVar v
takeMVarDBC = asksLDB dbcSyncMV >>= takeMVar
decodeKsId bs =
case decode bs of
Left e -> error $
"Error decoding Key Space ID: " <> show bs <> "\n" <> e
Right i -> i :: Word32