{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Database.Esqueleto.PostgreSQL.Streaming
  ( selectCursor
  ) where

import Control.Monad.IO.Class (MonadIO)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Reader (ReaderT(..), ask)
import Control.Monad.Trans.Resource (MonadResource)
import Data.Conduit (ConduitT)
import Data.Conduit.Lift (runReaderC)
import Database.Esqueleto.Internal.Internal
import Database.Persist.Postgresql
import Database.Persist.Postgresql.Streaming.Internal (rawSelectStream)

-- | Execute an @esqueleto@ @SELECT@ query against a PostgreSQL backend and
-- return a stream of the results in constant memory, using the PostgreSQL
-- cursor feature.
--
-- NB: this function is likely to perform worse than 'select' 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.
selectCursor
  ::
    ( SqlSelect a r
    , MonadIO m
    , MonadResource m
    , BackendCompatible (RawPostgresql SqlBackend) backend
    )
  => SqlQuery a
  -> ConduitT () r (ReaderT backend m) ()
selectCursor :: SqlQuery a -> ConduitT () r (ReaderT backend m) ()
selectCursor SqlQuery a
query = do
  backend :: RawPostgresql SqlBackend
backend@(RawPostgresql SqlBackend
conn Connection
_) <- ReaderT backend m (RawPostgresql SqlBackend)
-> ConduitT () r (ReaderT backend m) (RawPostgresql SqlBackend)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ReaderT backend m (RawPostgresql SqlBackend)
 -> ConduitT () r (ReaderT backend m) (RawPostgresql SqlBackend))
-> ReaderT backend m (RawPostgresql SqlBackend)
-> ConduitT () r (ReaderT backend m) (RawPostgresql SqlBackend)
forall a b. (a -> b) -> a -> b
$ backend -> RawPostgresql SqlBackend
forall sup sub. BackendCompatible sup sub => sub -> sup
projectBackend (backend -> RawPostgresql SqlBackend)
-> ReaderT backend m backend
-> ReaderT backend m (RawPostgresql SqlBackend)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ReaderT backend m backend
forall (m :: * -> *) r. Monad m => ReaderT r m r
ask
  let (Builder
queryTextBuilder, [PersistValue]
vals) = Mode
-> (SqlBackend, IdentState)
-> SqlQuery a
-> (Builder, [PersistValue])
forall a r backend.
(SqlSelect a r, BackendCompatible SqlBackend backend) =>
Mode
-> (backend, IdentState) -> SqlQuery a -> (Builder, [PersistValue])
toRawSql Mode
SELECT (SqlBackend
conn, IdentState
initialIdentState) SqlQuery a
query
      queryText :: Text
queryText = Builder -> Text
builderToText Builder
queryTextBuilder
  RawPostgresql SqlBackend
-> ConduitT
     () r (ReaderT (RawPostgresql SqlBackend) (ReaderT backend m)) ()
-> ConduitT () r (ReaderT backend m) ()
forall (m :: * -> *) r i o res.
Monad m =>
r -> ConduitT i o (ReaderT r m) res -> ConduitT i o m res
runReaderC RawPostgresql SqlBackend
backend
    (([PersistValue] -> Either Text r)
-> Text
-> [PersistValue]
-> ConduitT
     () r (ReaderT (RawPostgresql SqlBackend) (ReaderT backend m)) ()
forall (m :: * -> *) result.
MonadResource m =>
([PersistValue] -> Either Text result)
-> Text
-> [PersistValue]
-> ConduitT () result (ReaderT (RawPostgresql SqlBackend) m) ()
rawSelectStream [PersistValue] -> Either Text r
forall a r. SqlSelect a r => [PersistValue] -> Either Text r
sqlSelectProcessRow Text
queryText [PersistValue]
vals)