streamly-lmdb 0.1.0 → 0.2.0

raw patch · 6 files changed

+190/−39 lines, 6 filesdep ~streamly-lmdbPVP ok

version bump matches the API change (PVP)

Dependency ranges changed: streamly-lmdb

API changes (from Hackage documentation)

Files

• ChangeLog.md +4/−0
• README.md +5/−5
• src/Streamly/External/LMDB.hs +63/−9
• src/Streamly/External/LMDB/Internal/Foreign.hsc +9/−0
• streamly-lmdb.cabal +3/−3
• test/Streamly/External/LMDB/Tests.hs +106/−22

@@ -1,3 +1,7 @@+## 0.2.0++* Added read options for backward reading and starting from a specific key.+ ## 0.1.0  * Added a write option that disallows overwriting but does not throw an error when attempting to write a key-value pair that already exists in the database.

@@ -20,9 +20,9 @@ module Main where  import Streamly.External.LMDB-    (Limits (mapSize), WriteOptions (chunkSize), defaultLimits,-    defaultWriteOptions, getDatabase, openEnvironment, readLMDB,-    tebibyte, writeLMDB)+    (Limits (mapSize), WriteOptions (writeTransactionSize), defaultLimits,+    defaultReadOptions, defaultWriteOptions, getDatabase, openEnvironment,+    readLMDB, tebibyte, writeLMDB) import qualified Streamly.Prelude as S  main :: IO ()@@ -38,7 +38,7 @@     db <- getDatabase env Nothing      -- Stream key-value pairs into the database.-    let fold' = writeLMDB db defaultWriteOptions { chunkSize = 1 }+    let fold' = writeLMDB db defaultWriteOptions { writeTransactionSize = 1 }     let writeStream = S.fromList [("baz", "a"), ("foo", "b"), ("bar", "c")]     _ <- S.fold fold' writeStream @@ -48,7 +48,7 @@     --     ("bar","c")     --     ("baz","a")     --     ("foo","b")-    let unfold' = readLMDB db+    let unfold' = readLMDB db defaultReadOptions     let readStream = S.unfold unfold' undefined     S.mapM_ print readStream ```

@@ -13,6 +13,8 @@     Mode,     ReadWrite,     ReadOnly,+    ReadDirection(..),+    ReadOptions(..),     OverwriteOptions(..),     WriteOptions(..), @@ -28,6 +30,7 @@     clearDatabase,      -- ** Reading+    defaultReadOptions,     readLMDB,     unsafeReadLMDB, @@ -48,6 +51,7 @@ import Foreign.C (Errno (Errno), eNOTDIR) import Foreign.C.String (CStringLen) import Foreign.Marshal.Utils (with)+import Foreign.Storable (poke) import Streamly.Internal.Data.Fold (Fold (Fold)) import Streamly.Internal.Data.Stream.StreamD (newFinalizedIORef, runIORefFinalizer) import Streamly.Internal.Data.Stream.StreamD.Type (Step (Stop, Yield))@@ -147,7 +151,21 @@     mdb_clear ptxn dbi     mdb_txn_commit ptxn) >>= wait --- | Creates an unfold with which we can stream all key-value pairs from the given database in increasing key order.+-- | Direction of key iteration.+data ReadDirection = Forward | Backward deriving (Show)++data ReadOptions = ReadOptions+    { readDirection :: !ReadDirection+    -- | If 'Nothing', a forward [backward] iteration starts at the beginning [end] of the database.+    -- Otherwise, it starts at the first key that is greater [less] than or equal to the 'Just' key.+    , readStart     :: !(Maybe ByteString) } deriving (Show)++defaultReadOptions :: ReadOptions+defaultReadOptions = ReadOptions+    { readDirection = Forward+    , readStart = Nothing }++-- | Creates an unfold with which we can stream key-value pairs from the given database. -- -- A read transaction is kept open for the duration of the unfold; one should therefore -- bear in mind LMDB's [caveats regarding long-lived transactions](https://git.io/JJZE6).@@ -155,10 +173,10 @@ -- If you don’t want the overhead of intermediate 'ByteString's (on your -- way to your eventual data structures), use 'unsafeReadLMDB' instead. {-# INLINE readLMDB #-}-readLMDB :: (MonadIO m, Mode mode) => Database mode -> Unfold m Void (ByteString, ByteString)-readLMDB db = unsafeReadLMDB db packCStringLen packCStringLen+readLMDB :: (MonadIO m, Mode mode) => Database mode -> ReadOptions -> Unfold m Void (ByteString, ByteString)+readLMDB db ropts = unsafeReadLMDB db ropts packCStringLen packCStringLen --- | Creates an unfold with which we can stream all key-value pairs from the given database in increasing key order.+-- | Creates an unfold with which we can stream key-value pairs from the given database. -- -- A read transaction is kept open for the duration of the unfold; one should therefore -- bear in mind LMDB's [caveats regarding long-lived transactions](https://git.io/JJZE6).@@ -168,20 +186,50 @@ -- transform the 'CStringLen's to your desired data structures is to use 'Data.ByteString.Unsafe.unsafePackCStringLen'. {-# INLINE unsafeReadLMDB #-} unsafeReadLMDB :: (MonadIO m, Mode mode)-               => Database mode -> (CStringLen -> IO k) -> (CStringLen -> IO v) -> Unfold m Void (k, v)-unsafeReadLMDB (Database penv dbi) kmap vmap =-    flip supply mdb_first $ Unfold+               => Database mode+               -> ReadOptions+               -> (CStringLen -> IO k)+               -> (CStringLen -> IO v)+               -> Unfold m Void (k, v)+unsafeReadLMDB (Database penv dbi) ropts kmap vmap =+    let (firstOp, subsequentOp) = case (readDirection ropts, readStart ropts) of+            (Forward, Nothing) -> (mdb_first, mdb_next)+            (Forward, Just _) -> (mdb_set_range, mdb_next)+            (Backward, Nothing) -> (mdb_last, mdb_prev)+            (Backward, Just _) ->  (mdb_set_range, mdb_prev)+    in flip supply firstOp $ Unfold         (\(op, pcurs, pk, pv, ref) -> do-            found <- liftIO $ c_mdb_cursor_get pcurs pk pv op >>= \rc ->+            rc <- liftIO $+                if op == mdb_set_range && subsequentOp == mdb_prev then do+                    -- Reverse MDB_SET_RANGE.+                    kfst' <- peek pk+                    kfst <- packCStringLen (mv_data kfst', fromIntegral $ mv_size kfst')+                    rc <- c_mdb_cursor_get pcurs pk pv op+                    if rc /= 0 && rc == mdb_notfound then+                        c_mdb_cursor_get pcurs pk pv mdb_last+                    else if rc == 0 then do+                        k' <- peek pk+                        k <- unsafePackCStringLen (mv_data k', fromIntegral $ mv_size k')+                        if k /= kfst then+                            c_mdb_cursor_get pcurs pk pv mdb_prev+                        else+                            return rc+                    else+                        return rc+                else+                    c_mdb_cursor_get pcurs pk pv op++            found <- liftIO $                 if rc /= 0 && rc /= mdb_notfound then do                     runIORefFinalizer ref                     throwLMDBErrNum "mdb_cursor_get" rc                 else                     return $ rc /= mdb_notfound+             if found then do                 !k <- liftIO $ (\x -> kmap (mv_data x, fromIntegral $ mv_size x)) =<< peek pk                 !v <- liftIO $ (\x -> vmap (mv_data x, fromIntegral $ mv_size x)) =<< peek pv-                return $ Yield (k, v) (mdb_next, pcurs, pk, pv, ref)+                return $ Yield (k, v) (subsequentOp, pcurs, pk, pv, ref)             else do                 runIORefFinalizer ref                 return Stop)@@ -191,6 +239,12 @@                 pcurs <- liftIO $ mdb_cursor_open ptxn dbi                 pk <- liftIO malloc                 pv <- liftIO malloc++                _ <- case readStart ropts of+                    Nothing -> return ()+                    Just k -> unsafeUseAsCStringLen k $ \(kp, kl) ->+                        poke pk (MDB_val (fromIntegral kl) kp)+                 ref <- liftIO . newFinalizedIORef $ do                     free pv >> free pk                     c_mdb_cursor_close pcurs

@@ -188,8 +188,17 @@ mdb_first :: MDB_cursor_op_t mdb_first = #const MDB_FIRST +mdb_last :: MDB_cursor_op_t+mdb_last = #const MDB_LAST+ mdb_next :: MDB_cursor_op_t mdb_next = #const MDB_NEXT++mdb_prev :: MDB_cursor_op_t+mdb_prev = #const MDB_PREV++mdb_set_range :: MDB_cursor_op_t+mdb_set_range = #const MDB_SET_RANGE  mdb_env_create :: IO (Ptr MDB_env) mdb_env_create = do

@@ -4,10 +4,10 @@ -- -- see: https://github.com/sol/hpack ----- hash: 43ef967b54a793b5ad742a0e6d24b951d3961ac437567bc8756786917e48d678+-- hash: 9ad1d377f90069c99f24bcc557c1ff43a9f075fdfead1071725615fe8f352f50  name:           streamly-lmdb-version:        0.1.0+version:        0.2.0 synopsis:       Stream data to or from LMDB databases using the streamly library. description:    Please see the README on GitHub at <https://github.com/shlok/streamly-lmdb#readme> category:       Database, Streaming, Streamly@@ -70,7 +70,7 @@     , bytestring >=0.10.10.0 && <0.11     , directory >=1.3.6.0 && <1.4     , streamly >=0.7.2 && <0.8-    , streamly-lmdb ==0.1.0+    , streamly-lmdb ==0.2.0     , tasty >=1.2.3 && <1.3     , tasty-quickcheck >=0.10.1.1 && <0.11     , temporary >=1.3 && <1.4

@@ -5,12 +5,13 @@ import Control.Concurrent.Async (asyncBound, wait) import Control.Exception (SomeException, onException, try) import Control.Monad (forM_)-import Data.ByteString (ByteString, pack)+import Data.ByteString (ByteString, pack, unpack) import Data.ByteString.Unsafe (unsafeUseAsCStringLen) import Data.List (find, foldl', nubBy, sort)+import Data.Word (Word8) import Foreign (castPtr, nullPtr, with) import Streamly.Prelude (fromList, toList, unfold)-import Test.QuickCheck (choose)+import Test.QuickCheck (NonEmptyList (..), Gen, choose, elements, frequency) import Test.QuickCheck.Monadic (PropertyM, monadicIO, pick, run) import Test.Tasty (TestTree) import Test.Tasty.QuickCheck (arbitrary, testProperty)@@ -18,8 +19,8 @@ import qualified Data.ByteString as B import qualified Streamly.Prelude as S -import Streamly.External.LMDB (Mode, ReadWrite, OverwriteOptions(..), WriteOptions(..),-    clearDatabase, defaultWriteOptions, readLMDB, unsafeReadLMDB, writeLMDB)+import Streamly.External.LMDB (Mode, ReadWrite, OverwriteOptions (..), ReadDirection (..), ReadOptions (..),+    WriteOptions(..), clearDatabase, defaultReadOptions, defaultWriteOptions, readLMDB, unsafeReadLMDB, writeLMDB) import Streamly.External.LMDB.Internal (Database (..)) import Streamly.External.LMDB.Internal.Foreign (MDB_val (..), combineOptions,     mdb_nooverwrite, mdb_put, mdb_txn_begin, mdb_txn_commit)@@ -30,43 +31,40 @@     , testUnsafeReadLMDB db     , testWriteLMDB db     , testWriteLMDB_2 db-    , testWriteLMDB_3 db ]+    , testWriteLMDB_3 db+    , testBetween ]  -- | Clear the database, write key-value pairs to it in a normal manner, read -- them back using our library, and make sure the result is what we wrote. testReadLMDB :: (Mode mode) => IO (Database mode) -> TestTree testReadLMDB res = testProperty "readLMDB" . monadicIO $ do     db <- run res-    keyValuePairs <- arbitraryKeyValuePairs+    keyValuePairs <- arbitraryKeyValuePairs''     run $ clearDatabase db      run $ writeChunk db False keyValuePairs     let keyValuePairsInDb = sort . removeDuplicateKeys $ keyValuePairs -    let stream = unfold (readLMDB db) undefined-    readPairsAll <- run . toList $ stream-    let allAsExpected = readPairsAll == keyValuePairsInDb--    -- Also make sure the stream’s 'take' functionality works.-    readPairsFirstFew <- run . toList $ S.take 3 stream-    let firstFewAsExpected = readPairsFirstFew == take 3 keyValuePairsInDb+    (readOpts, expectedResults) <- pick $ readOptionsAndResults keyValuePairsInDb+    results <- run . toList $ unfold (readLMDB db readOpts) undefined -    return $ allAsExpected && firstFewAsExpected+    return $ results == expectedResults  -- | Similar to 'testReadLMDB', except that it tests the unsafe function in a different manner. testUnsafeReadLMDB :: (Mode mode) => IO (Database mode) -> TestTree testUnsafeReadLMDB res = testProperty "unsafeReadLMDB" . monadicIO $ do     db <- run res-    keyValuePairs <- arbitraryKeyValuePairs+    keyValuePairs <- arbitraryKeyValuePairs''     run $ clearDatabase db      run $ writeChunk db False keyValuePairs-    let lengthsInDb = map (\(k, v) -> (B.length k, B.length v)) . sort . removeDuplicateKeys $ keyValuePairs+    let keyValuePairsInDb = sort . removeDuplicateKeys $ keyValuePairs -    let stream = unfold (unsafeReadLMDB db (return . snd) (return . snd)) undefined-    readLengthsAll <- run . toList $ stream+    (readOpts, expectedResults) <- pick $ readOptionsAndResults keyValuePairsInDb+    let expectedLengths = map (\(k, v) -> (B.length k, B.length v)) expectedResults+    lengths <- run . toList $ unfold (unsafeReadLMDB db readOpts (return . snd) (return . snd)) undefined -    return $ readLengthsAll == lengthsInDb+    return $ lengths == expectedLengths  -- | Clear the database, write key-value pairs to it using our library with key overwriting allowed, read -- them back using our library (already covered by 'testReadLMDB'), and make sure the result is what we wrote.@@ -85,7 +83,7 @@     run $ asyncBound (S.fold fol' (fromList keyValuePairs)) >>= wait     let keyValuePairsInDb = sort . removeDuplicateKeys $ keyValuePairs -    readPairsAll <- run . toList $ unfold (readLMDB db) undefined+    readPairsAll <- run . toList $ unfold (readLMDB db defaultReadOptions) undefined      return $ keyValuePairsInDb == readPairsAll @@ -110,7 +108,7 @@             Right _ -> return . not $ hasDuplicateKeys keyValuePairs      let keyValuePairsInDb = sort . prefixBeforeDuplicate  $ keyValuePairs-    readPairsAll <- run . toList $ unfold (readLMDB db) undefined+    readPairsAll <- run . toList $ unfold (readLMDB db defaultReadOptions) undefined     let pairsAsExpected = keyValuePairsInDb == readPairsAll      return $ exceptionAsExpected && pairsAsExpected@@ -137,7 +135,7 @@             Right _ -> return . not $ hasDuplicateKeysWithDiffVals keyValuePairs      let keyValuePairsInDb = sort . removeDuplicateKeys . prefixBeforeDuplicateWithDiffVal $ keyValuePairs-    readPairsAll <- run . toList $ unfold (readLMDB db) undefined+    readPairsAll <- run . toList $ unfold (readLMDB db defaultReadOptions) undefined     let pairsAsExpected = keyValuePairsInDb == readPairsAll      return $ exceptionAsExpected && pairsAsExpected@@ -164,6 +162,22 @@     else         return arb +-- A variation that makes more likely keys with same the prefix and a difference of trailing zero bytes.+arbitraryKeyValuePairs'' :: PropertyM IO [(ByteString, ByteString)]+arbitraryKeyValuePairs'' = do+    arb <- arbitraryKeyValuePairs+    if null arb then+        return arb+    else pick $ frequency [(1, return arb), (3, do+        let (k, v) = head arb+        b' <- arbitrary+        v' <- if b' then return v else pack <$> arbitrary+        i <- choose (0, length arb - 1)+        let (arb1, arb2) = splitAt i arb+        let arb3 = map (\i' -> (k `B.append` B.replicate i' 0, v')) [1..(i+1)]+        let arb' = arb1 ++ arb3 ++ arb2+        return arb' )]+ -- | Note that this function retains the last value for each key. removeDuplicateKeys :: (Eq a) => [(a, b)] -> [(a, b)] removeDuplicateKeys = foldl' (\acc (a, b) -> if any ((== a) . fst) acc then acc else (a, b) : acc) [] . reverse@@ -191,6 +205,76 @@     in case fstDup of         Nothing -> xs         Just i -> take i xs++-- Assumes first < second.+between :: [Word8] -> [Word8] -> [Word8] -> Maybe [Word8]+between [] [] _ = error "first = second"+between _ [] _ = error "first > second"+between [] (w:ws) commonPrefixRev+    | w == 0 && null ws = Nothing+    | w == 0 = between [] ws (w:commonPrefixRev)+    | otherwise = Just $ reverse (0:commonPrefixRev)+between (w1:ws1) (w2:ws2) commonPrefixRev+    | w1 == w2 = between ws1 ws2 (w1:commonPrefixRev)+    | w1 > w2 = error "first > second"+    | otherwise = Just $ reverse commonPrefixRev ++ [w1] ++ ws1 ++ [0]++testBetween :: TestTree+testBetween = testProperty "testBetween" $ \ws1 ws2 ->+    (ws1 == ws2) ||+        let (smaller, bigger) = if ws1 < ws2 then (ws1, ws2) else (ws2, ws1)+        in case between smaller bigger [] of+            Nothing -> drop (length ws1) ws2 == replicate (length ws2 - length ws1) 0+            Just betw -> smaller < betw && betw < bigger++betweenBs :: ByteString -> ByteString -> Maybe ByteString+betweenBs bs1 bs2 = between (unpack bs1) (unpack bs2) [] >>= (return . pack)++type PairsInDatabase = [(ByteString, ByteString)]+type ExpectedReadResult = [(ByteString, ByteString)]++-- | Given database pairs, randomly generates read options and corresponding expected results.+readOptionsAndResults :: PairsInDatabase -> Gen (ReadOptions, ExpectedReadResult)+readOptionsAndResults pairsInDb = do+    forw <- arbitrary+    let dir = if forw then Forward else Backward+    let len = length pairsInDb+    readAll <- frequency [ (1, return True), (3, return False) ]+    let ropts = defaultReadOptions { readDirection = dir }+    if readAll then+        return (ropts { readStart = Nothing }, (if forw then id else reverse) pairsInDb)+    else if len == 0 then do+        bs <- arbitrary >>= \(NonEmpty ws) -> return $ pack ws+        return (ropts { readStart = Just bs }, [])+    else do+        idx <- if len < 3 then choose (0, len - 1) else+                frequency [ (1, choose (1, len - 2)), (3, elements [0, len - 1]) ]+        let keyAt i = fst $ pairsInDb !! i+        let nextKey+                | idx + 1 <= len - 1 = betweenBs (keyAt idx) (keyAt $ idx + 1)+                | otherwise = Just $ keyAt (len - 1) `B.append` B.singleton 0+        let prevKey+                | idx == 0 && keyAt idx /= B.singleton 0 = Just $ B.singleton 0 -- Keys are known to be non-empty.+                | idx == 0 = Nothing+                | otherwise = betweenBs (keyAt $ idx - 1) (keyAt idx)+        let forwEq = (ropts { readStart = Just $ keyAt idx }, drop idx pairsInDb)+        let backwEq = (ropts { readStart = Just $ keyAt idx }, reverse $ take (idx + 1) pairsInDb)+        ord <- arbitrary @Ordering -- Proximity to the key at idx (if possible).+        return $ case (ord, dir) of+            (EQ, Forward) -> forwEq+            (EQ, Backward) -> backwEq+            (GT, Forward) -> case nextKey of+                Nothing -> forwEq+                Just nextKey' -> (ropts { readStart = Just nextKey' }, drop (idx + 1) pairsInDb)+            (GT, Backward) -> case nextKey of+                Nothing -> backwEq+                Just nextKey' -> (ropts { readStart = Just nextKey' }, reverse $ take (idx + 1) pairsInDb)+            (LT, Forward) -> case prevKey of+                Nothing -> forwEq+                Just prevKey' -> (ropts { readStart = Just prevKey' }, drop idx pairsInDb)+            (LT, Backward) -> case prevKey of+                Nothing -> backwEq+                Just prevKey' -> (ropts { readStart = Just prevKey' }, reverse $ take idx pairsInDb)  -- Writes the given key-value pairs to the given database. writeChunk :: (Foldable t, Mode mode) => Database mode -> Bool -> t (ByteString, ByteString) -> IO ()