{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE BangPatterns #-}
module Lmdb.Connection where
import Database.LMDB.Raw
import Lmdb.Types
import Lmdb.Internal
import Data.Word
import Foreign.Storable
import Data.Coerce
import Data.Functor
import Data.Bits
import Control.Concurrent (runInBoundThread,isCurrentThreadBound)
import Data.Bool (bool)
import System.IO (withFile,IOMode(ReadMode))
import Foreign.C.Types (CSize(..))
import Foreign.Ptr (Ptr,plusPtr)
import Foreign.Marshal.Alloc (allocaBytes,alloca)
import Control.Monad
import Control.Exception (finally, bracketOnError)
withCursor ::
Transaction e
-> Database k v
-> (Cursor e k v -> IO a)
-> IO a
withCursor (Transaction txn) (Database dbi settings) f = do
cur <- mdb_cursor_open_X txn dbi
a <- f (Cursor cur settings)
mdb_cursor_close_X cur
return a
withMultiCursor ::
Transaction e
-> MultiDatabase k v
-> (MultiCursor e k v -> IO a)
-> IO a
withMultiCursor (Transaction txn) (MultiDatabase dbi settings) f = do
cur <- mdb_cursor_open_X txn dbi
a <- f (MultiCursor cur settings)
mdb_cursor_close_X cur
return a
withAbortableTransaction ::
ModeBool e
=> Environment e
-> (Transaction e -> IO (Maybe a))
-> IO (Maybe a)
withAbortableTransaction e@(Environment env) f = do
let isReadOnly = modeIsReadOnly e
txn <- mdb_txn_begin env Nothing isReadOnly
ma <- f (Transaction txn)
case ma of
Nothing -> do
mdb_txn_abort txn
return Nothing
Just a -> do
mdb_txn_commit txn
return (Just a)
withTransaction ::
ModeBool e
=> Environment e
-> (Transaction e -> IO a)
-> IO a
withTransaction e@(Environment env) f = do
let isReadOnly = modeIsReadOnly e
bool runInBoundThread id isReadOnly $ bracketOnError
(mdb_txn_begin env Nothing isReadOnly)
mdb_txn_abort
$ \txn -> do
a <- f (Transaction txn)
mdb_txn_commit txn
return a
withNestedTransaction ::
Environment 'ReadWrite
-> Transaction 'ReadWrite
-> (Transaction 'ReadWrite -> IO a)
-> IO a
withNestedTransaction e@(Environment env) (Transaction parentTxn) f = do
txn <- mdb_txn_begin env (Just parentTxn) True
a <- f (Transaction txn)
mdb_txn_commit txn
return a
-- This function can be improved to handle custom sorting.
openDatabase ::
ModeBool e
=> Transaction e
-> Maybe String -- ^ Database name
-> DatabaseSettings k v
-> IO (Database k v)
openDatabase t@(Transaction txn) name settings = do
let rwOpts = if modeIsReadOnly t then [] else [MDB_CREATE]
(keySafeFfi, keyExtraCmd, sortOpts) = case settings of
DatabaseSettings keySort _ keyDec _ _ -> case keySort of
SortNative s -> (,,) False (\_ -> return ()) $ case s of
NativeSortLexographic -> []
NativeSortLexographicBackward -> [MDB_REVERSEKEY]
NativeSortInteger -> [MDB_INTEGERKEY]
SortCustom s -> customSortConfig True keyDec t s
opts = rwOpts ++ sortOpts
dbi <- mdb_dbi_open_X keySafeFfi txn name opts
keyExtraCmd dbi
return (Database dbi settings)
customSortConfig :: Bool -> Decoding a -> Transaction e -> CustomSort a -> (Bool, DbiByFfi -> IO (), [MDB_DbFlag])
customSortConfig isKey (Decoding decode) (Transaction txn) s = (,,) True (\dbiOuter -> case dbiOuter of
DbiSafe dbi -> case s of
CustomSortUnsafe f -> setCmp txn dbi f
CustomSortSafe f -> (setCmp txn dbi =<<) $ wrapCmpFn $ \aPtr bPtr -> do
MDB_val aSize aData <- peek aPtr
MDB_val bSize bData <- peek bPtr
a <- decode aSize aData
b <- decode bSize bData
return $ case f a b of
GT -> 1
EQ -> 0
LT -> (-1)
DbiUnsafe _ -> fail "customSortConfig: logical error in sorting. Open an issue if this happens."
) []
where setCmp = if isKey then mdb_set_compare else mdb_set_dupsort
-- This function can be improved to handle custom sorting.
openMultiDatabase ::
ModeBool e
=> Transaction e
-> Maybe String -- ^ Database name
-> MultiDatabaseSettings k v
-> IO (MultiDatabase k v)
openMultiDatabase t@(Transaction txn) name settings = do
let rwOpts = if modeIsReadOnly t then [] else [MDB_CREATE]
(keySafeFfi,keyExtraCmd,keySortOpts) = case settings of
MultiDatabaseSettings keySort _ _ keyDec _ _ -> case keySort of
SortNative s -> (,,) False (\_ -> return ()) $ case s of
NativeSortLexographic -> []
NativeSortLexographicBackward -> [MDB_REVERSEKEY]
NativeSortInteger -> [MDB_INTEGERKEY]
SortCustom s -> customSortConfig True keyDec t s
(valSafeFfi,valExtraCmd,valSortOpts) = case settings of
MultiDatabaseSettings _ valSort _ _ _ valDec -> case valSort of
SortNative s -> (,,) False (\_ -> return ()) $ case s of
NativeSortLexographic -> []
NativeSortLexographicBackward -> [MDB_REVERSEDUP]
NativeSortInteger -> [MDB_INTEGERDUP]
SortCustom s -> customSortConfig False valDec t s
multiOpts = case settings of
MultiDatabaseSettings _ _ _ _ valEnc _ -> case valEnc of
EncodingVariable _ -> []
EncodingMachineWord _ -> [MDB_DUPFIXED]
EncodingFixed _ _ -> [MDB_DUPFIXED]
baseOpts = [MDB_DUPSORT]
safeFfi = keySafeFfi || valSafeFfi
opts = rwOpts ++ keySortOpts ++ valSortOpts ++ multiOpts ++ baseOpts
dbi <- mdb_dbi_open_X safeFfi txn name opts
keyExtraCmd dbi
valExtraCmd dbi
return (MultiDatabase dbi settings)
-- | This should not normally be used.
withDatabase ::
ModeBool e
=> Environment e
-> Transaction e
-> Maybe String
-> DatabaseSettings k v
-> (Database k v -> IO a)
-> IO a
withDatabase env txn mname settings f = do
dbi <- openDatabase txn mname settings
finally (f dbi) (closeDatabase env dbi)
-- | This should not normally be used.
withMultiDatabase ::
ModeBool e
=> Environment e
-> Transaction e
-> Maybe String
-> MultiDatabaseSettings k v
-> (MultiDatabase k v -> IO a)
-> IO a
withMultiDatabase env txn mname settings f = do
dbi <- openMultiDatabase txn mname settings
finally (f dbi) (closeMultiDatabase env dbi)
-- | Internally, this calls @mdb_env_create@ and @mdb_env_open@.
initializeReadOnlyEnvironment ::
Int -- ^ Map size in bytes
-> Int -- ^ Maximum number of readers (recommended: 126)
-> Int -- ^ Maximum number of databases
-> FilePath -- ^ Directory for lmdb data and locks
-> IO (Environment 'ReadOnly)
initializeReadOnlyEnvironment =
initializeEnvironmentInternal [MDB_RDONLY]
withReadOnlyEnvironment ::
Int -- ^ Map size in bytes
-> Int -- ^ Maximum number of readers (recommended: 126)
-> Int -- ^ Maximum number of databases
-> FilePath -- ^ Directory for lmdb data and locks
-> (Environment 'ReadOnly -> IO a) -- ^ Computation requiring an 'Environment'
-> IO a
withReadOnlyEnvironment a b c d f = do
env <- initializeReadOnlyEnvironment a b c d
finally (f env) (closeEnvironment env)
-- | Internally, this calls @mdb_env_create@ and @mdb_env_open@.
initializeReadWriteEnvironment ::
Int -- ^ Map size in bytes
-> Int -- ^ Maximum number of readers (recommended: 126)
-> Int -- ^ Maximum number of databases
-> FilePath -- ^ Directory for lmdb data and locks
-> IO (Environment 'ReadWrite)
initializeReadWriteEnvironment = initializeEnvironmentInternal []
-- | It is not clear whether or not it is actually neccessary
-- to use 'runInBoundThread' here. It is done just as an extra
-- precaution.
initializeEnvironmentInternal ::
[MDB_EnvFlag] -- ^ Flags
-> Int -- ^ Map size in bytes
-> Int -- ^ Maximum number of readers (recommended: 126)
-> Int -- ^ Maximum number of databases
-> FilePath -- ^ Directory for lmdb data and locks
-> IO (Environment e)
initializeEnvironmentInternal flags maxSize maxReaders maxDbs dir =
runInBoundThread $ do
env <- mdb_env_create
mdb_env_set_mapsize env maxSize
mdb_env_set_maxreaders env maxReaders
mdb_env_set_maxdbs env maxDbs
mdb_env_open env dir flags
return (Environment env)
closeDatabase :: Environment e -> Database k v -> IO ()
closeDatabase (Environment env) (Database dbi _) =
mdb_dbi_close_X env dbi
closeMultiDatabase :: Environment e -> MultiDatabase k v -> IO ()
closeMultiDatabase (Environment env) (MultiDatabase dbi _) =
mdb_dbi_close_X env dbi
closeEnvironment :: Environment e -> IO ()
closeEnvironment (Environment env) =
runInBoundThread $ mdb_env_close env
makeSettings ::
Sort s k -- ^ Key sorting function
-> Codec s k -- ^ Key codec
-> Codec sv v -- ^ Value codec
-> DatabaseSettings k v
makeSettings sort (Codec kEnc kDec) (Codec vEnc vDec) =
DatabaseSettings sort kEnc kDec vEnc vDec
makeMultiSettings ::
Sort sk k -- ^ Key sorting function
-> Sort sv v -- ^ Value sorting function
-> Codec sk k -- ^ Key codec
-> Codec sv v -- ^ Value codec
-> MultiDatabaseSettings k v
makeMultiSettings ksort vsort (Codec kEnc kDec) (Codec vEnc vDec) =
MultiDatabaseSettings ksort vsort kEnc kDec vEnc vDec
readonly :: Transaction 'ReadWrite -> Transaction 'ReadOnly
readonly = coerce
readonlyEnvironment :: Environment 'ReadWrite -> Environment 'ReadOnly
readonlyEnvironment = coerce
-- csvsToLmdb ::
-- FilePath -- ^ Geolite IPv4 Blocks
-- -> FilePath -- ^ Geolite City Locations
-- -> FilePath -- ^ Directory for LMDB
-- -> IO ()
-- csvsToLmdb geoBlocksPath geoCitiesPath lmdbDir = do
-- env <- mdb_env_create
-- mdb_env_set_mapsize env (2 ^ 31)
-- mdb_env_set_maxreaders env 1
-- mdb_env_set_maxdbs env 2
-- mdb_env_open env lmdbDir []
-- txn <- mdb_txn_begin env Nothing True
-- dbiBlocks <- mdb_dbi_open' txn (Just "blocks") [MDB_INTEGERKEY,MDB_CREATE]
-- let appendFlag = compileWriteFlags [MDB_APPEND]
-- r <- withFile filename ReadMode $ \h -> runEffect $
-- fmap (SD.convertDecodeError "utf-8") (PT.decode (PT.utf8 . PT.eof) $ PB.fromHandle h)
-- >-> fmap Just blocks
-- >-> Pipes.mapM_ (\block -> do
-- alloca $ \(w64Ptr :: Ptr Word64) -> do
-- poke w64Ptr (fromIntegral w32)
-- let IPv4Range (IPv4 w32) _ = blockNetwork block
-- w8Ptr = (castPtr :: Ptr Word64 -> Ptr Word8) w64Ptr
-- mdb_put' appendFlag txn dbiBlocks
-- (MDB_val (CSize 8) w8Ptr)
-- (MDB_val (CSize 8) w8Ptr)
-- )
-- case r of
-- Nothing -> assertBool "impossible" True
-- Just err -> assertFailure (Decoding.prettyError Text.unpack err)
-- putBlock :: Block -> Put
-- putBlock (Block network geonameId registered represented isAnon isSat postal lat lon accuracy)