diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,8 @@ +# Changelog for persistent-postgresql-streaming + +# 0.1.0.0 + + * Exposed `Database.Persist.Postgresql.Streaming` + * Added `selectStream` + * Exposed `Database.Persist.Postgresql.Streaming.Internal` + * Added `rawSelectStream` diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Supercede Ltd. (c) 2021 + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Isaac van Bakel nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,17 @@ +# persistent-postgresql-streaming + +This library allows for memory-constant streaming of [`persistent`](https://hackage.haskell.org/package/persistent) entities from PostgreSQL databases. + +This code makes use of the PostgreSQL-only [cursors](https://www.postgresql.org/docs/current/plpgsql-cursors.html), which allow for batched access to the result set of a query at speeds comparable to loading all the results into memory at once. + +See the [main project README](https://github.com/SupercedeTech/persistent-postgresql-streaming/blob/main/README.md) for more. + +## Streaming Persistent queries + +The main function of this library is `selectStream`, which can be used in place of [`selectSource`](https://hackage.haskell.org/package/persistent-2.13.1.2/docs/Database-Persist-Class.html#v:selectSource). `selectSource` runs in a `ConduitT` - consuming the conduit will pull results from the database in constant memory. + +## FAQ + +### Why is `selectStream` so slow? + +Have you configured PostgreSQL correctly? See the section in the [README](https://github.com/SupercedeTech/persistent-postgresql-streaming/blob/main/README.md) about it. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/persistent-postgresql-streaming.cabal b/persistent-postgresql-streaming.cabal new file mode 100644 --- /dev/null +++ b/persistent-postgresql-streaming.cabal @@ -0,0 +1,40 @@ +name: persistent-postgresql-streaming +version: 0.1.0.0 +synopsis: Memory-constant streaming of Persistent entities from PostgreSQL +description: + This library provides a set of APIs for performing queries on Persistent + entities in constant memory, streaming the results using @conduit@. + . + The library relies on PostgreSQL-specific features to avoid loading all the + results of a query into memory at once. This allows for accessing tables of + millions of entities from Haskell without a memory blow-up. +homepage: https://github.com/SupercedeTech/persistent-postgresql-streaming#readme +license: BSD3 +license-file: LICENSE +author: Isaac van Bakel +maintainer: support@supercede.com +copyright: (c) 2021 Supercede Ltd. +category: Conduit, Database +build-type: Simple +extra-source-files: CHANGELOG.md README.md +cabal-version: >=1.10 + +library + hs-source-dirs: src + exposed-modules: Database.Persist.Postgresql.Streaming + , Database.Persist.Postgresql.Streaming.Internal + build-depends: base >= 4.7 && < 5 + , conduit >= 1.3 && < 1.4 + , monad-logger >= 0.3.28.4 && < 0.4 + , mtl >= 2 && < 2.3 + , persistent >= 2.13.2 && < 2.14 + , persistent-postgresql >= 2.13.2 && < 2.14 + , postgresql-simple >= 0.6.1 && < 0.7 + , resourcet >= 0.3 && < 2 + , text >= 1 && < 1.3 + , transformers >= 0.5 && < 0.6 + default-language: Haskell2010 + +source-repository head + type: git + location: https://github.com/SupercedeTech/persistent-postgresql-streaming diff --git a/src/Database/Persist/Postgresql/Streaming.hs b/src/Database/Persist/Postgresql/Streaming.hs new file mode 100644 --- /dev/null +++ b/src/Database/Persist/Postgresql/Streaming.hs @@ -0,0 +1,69 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ViewPatterns #-} + +module Database.Persist.Postgresql.Streaming + ( selectStream + ) where + +import Control.Monad.Reader.Class (MonadReader(ask)) +import Control.Monad.Trans.Class (lift) +import Control.Monad.Trans.Resource (MonadResource) +import Data.Conduit (ConduitT) +import Data.Conduit.Lift (runReaderC) +import Data.Foldable (toList) +import Database.Persist.Sql.Types.Internal (SqlBackend(..)) +import Database.Persist.Sql.Util + ( commaSeparated + , keyAndEntityColumnNames + , parseEntityValues + ) +import Database.Persist.Postgresql +import Database.Persist.Postgresql.Streaming.Internal + +-- | Run a query against a PostgreSQL backend, streaming back the results +-- in constant memory using the PostgreSQL cursor feature. +-- +-- NB: this function is likely to perform worse than 'selectSource' on small +-- result sets. +-- +-- NB: for large result sets, this function is likely to perform very poorly +-- unless you manually configure the @cursor_tuple_fraction@ setting to be +-- close to @1@. This setting decides how much PostgreSQL will prioritise +-- returning the first rows quickly vs. returning all rows efficiently. See +-- +-- https://postgresqlco.nf/doc/en/param/cursor_tuple_fraction/ +-- +-- for more. +selectStream + :: + ( MonadResource m + , MonadReader backend m + , PersistRecordBackend record SqlBackend + , BackendCompatible (RawPostgresql SqlBackend) backend + ) + => [Filter record] + -> [SelectOpt record] + -> ConduitT () (Entity record) m () +selectStream filts opts = do + backend@(RawPostgresql conn _) <- lift $ projectBackend <$> ask + runReaderC backend $ rawSelectStream (parseEntityValues t) (sql conn) (getFiltsValues conn filts) + where + (limit, offset, orders) = limitOffsetOrder opts + t = entityDef $ dummyFromFilts filts + wher conn = filterClause Nothing conn filts + ord conn = orderClause Nothing conn orders + cols = commaSeparated . toList . keyAndEntityColumnNames t + sql conn = connLimitOffset conn (limit,offset) $ mconcat + [ "SELECT " + , cols conn + , " FROM " + , connEscapeTableName conn t + , wher conn + , ord conn + ] + + getFiltsValues conn = snd . filterClauseWithVals Nothing conn + dummyFromFilts :: [Filter record] -> Maybe record + dummyFromFilts _ = Nothing diff --git a/src/Database/Persist/Postgresql/Streaming/Internal.hs b/src/Database/Persist/Postgresql/Streaming/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Database/Persist/Postgresql/Streaming/Internal.hs @@ -0,0 +1,91 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ViewPatterns #-} + +module Database.Persist.Postgresql.Streaming.Internal + ( rawSelectStream + ) where + +import Control.Exception (throwIO) +import Control.Monad.IO.Class (MonadIO(liftIO)) +import Control.Monad.Logger (LoggingT(..), logDebugNS) +import Control.Monad.Reader.Class (MonadReader(ask)) +import Control.Monad.Trans.Class (lift) +import Control.Monad.Trans.Reader (ReaderT(..)) +import Control.Monad.Trans.Resource (MonadResource, release) +import Data.Acquire (Acquire, allocateAcquire, mkAcquire) +import Data.Conduit (ConduitT, (.|)) +import qualified Data.Conduit.Combinators as CC (mapM, yieldMany) +import qualified Data.Text as T (Text, append, pack) +import qualified Data.Text.Encoding as T (encodeUtf8) +import Database.Persist.Sql.Types.Internal (SqlBackend(..)) +import Database.Persist.Postgresql +import Database.Persist.Postgresql.Internal +import qualified Database.PostgreSQL.Simple as PG +import qualified Database.PostgreSQL.Simple.Cursor as PGC +import qualified Database.PostgreSQL.Simple.Types as PG + +-- | Run a @Text@ query, with interpolated @PersistValue@s, against a PostgreSQL +-- backend using cursors, parsing the results with a custom parser function and +-- streaming them back. +-- +-- If the parser function returns @Left@ for any row, a 'PersistException' will +-- be thrown. +rawSelectStream + :: MonadResource m + => ([PersistValue] -> Either T.Text result) + -> T.Text + -> [PersistValue] + -> ConduitT () result (ReaderT (RawPostgresql SqlBackend) m) () +rawSelectStream parseRes query vals = do + srcRes <- lift $ liftPersist $ do + srcRes <- rawQueryResFromCursor query vals + return $ fmap (.| CC.mapM parse) srcRes + (releaseKey, src) <- allocateAcquire srcRes + src + release releaseKey + where + parse resVals = + case parseRes resVals of + Left s -> + liftIO $ throwIO $ + PersistMarshalError ("rawSelectStream: " <> s <> ", vals: " <> T.pack (show vals )) + Right row -> + return row + +rawQueryResFromCursor + :: (MonadIO m1, MonadIO m2, BackendCompatible SqlBackend backend) + => T.Text + -> [PersistValue] + -> ReaderT (RawPostgresql backend) m1 (Acquire (ConduitT () [PersistValue] m2 ())) +rawQueryResFromCursor sql vals = do + RawPostgresql conn' pgConn <- ask + let conn = projectBackend conn' + runLoggingT + (logDebugNS "SQL" $ T.append sql $ T.pack $ "; " ++ show vals) + (connLogFunc conn) + return $ withCursorStmt pgConn (PG.Query $ T.encodeUtf8 sql) vals + +withCursorStmt + :: MonadIO m + => PG.Connection + -> PG.Query + -> [PersistValue] + -> Acquire (ConduitT () [PersistValue] m ()) +withCursorStmt conn query vals = + foldWithCursor `fmap` mkAcquire openC closeC + where + openC = do + rawquery <- liftIO $ PG.formatQuery conn query (map P vals) + PGC.declareCursor conn (PG.Query rawquery) + closeC = PGC.closeCursor + foldWithCursor cursor = go + where + go = do + -- 256 is the default chunk size used for fetching + rows <- liftIO $ PGC.foldForward cursor 256 processRow [] + case rows of + Left final -> CC.yieldMany final + Right nonfinal -> CC.yieldMany nonfinal >> go + processRow s row = pure $ s <> [map unP row]