{-# 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 :: [Filter record]
-> [SelectOpt record] -> ConduitT () (Entity record) m ()
selectStream [Filter record]
filts [SelectOpt record]
opts = do
  backend :: RawPostgresql SqlBackend
backend@(RawPostgresql SqlBackend
conn Connection
_) <- m (RawPostgresql SqlBackend)
-> ConduitT () (Entity record) m (RawPostgresql SqlBackend)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (RawPostgresql SqlBackend)
 -> ConduitT () (Entity record) m (RawPostgresql SqlBackend))
-> m (RawPostgresql SqlBackend)
-> ConduitT () (Entity record) 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)
-> m backend -> m (RawPostgresql SqlBackend)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m backend
forall r (m :: * -> *). MonadReader r m => m r
ask
  RawPostgresql SqlBackend
-> ConduitT
     () (Entity record) (ReaderT (RawPostgresql SqlBackend) m) ()
-> ConduitT () (Entity record) 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 (ConduitT
   () (Entity record) (ReaderT (RawPostgresql SqlBackend) m) ()
 -> ConduitT () (Entity record) m ())
-> ConduitT
     () (Entity record) (ReaderT (RawPostgresql SqlBackend) m) ()
-> ConduitT () (Entity record) m ()
forall a b. (a -> b) -> a -> b
$ ([PersistValue] -> Either Text (Entity record))
-> Text
-> [PersistValue]
-> ConduitT
     () (Entity record) (ReaderT (RawPostgresql SqlBackend) m) ()
forall (m :: * -> *) result.
MonadResource m =>
([PersistValue] -> Either Text result)
-> Text
-> [PersistValue]
-> ConduitT () result (ReaderT (RawPostgresql SqlBackend) m) ()
rawSelectStream (EntityDef -> [PersistValue] -> Either Text (Entity record)
forall record.
PersistEntity record =>
EntityDef -> [PersistValue] -> Either Text (Entity record)
parseEntityValues EntityDef
t) (SqlBackend -> Text
sql SqlBackend
conn) (SqlBackend -> [Filter record] -> [PersistValue]
forall val.
PersistEntity val =>
SqlBackend -> [Filter val] -> [PersistValue]
getFiltsValues SqlBackend
conn [Filter record]
filts)
  where
    (Int
limit, Int
offset, [SelectOpt record]
orders) = [SelectOpt record] -> (Int, Int, [SelectOpt record])
forall val.
PersistEntity val =>
[SelectOpt val] -> (Int, Int, [SelectOpt val])
limitOffsetOrder [SelectOpt record]
opts
    t :: EntityDef
t = Maybe record -> EntityDef
forall record (proxy :: * -> *).
PersistEntity record =>
proxy record -> EntityDef
entityDef (Maybe record -> EntityDef) -> Maybe record -> EntityDef
forall a b. (a -> b) -> a -> b
$ [Filter record] -> Maybe record
forall record. [Filter record] -> Maybe record
dummyFromFilts [Filter record]
filts
    wher :: SqlBackend -> Text
wher SqlBackend
conn = Maybe FilterTablePrefix -> SqlBackend -> [Filter record] -> Text
forall val.
PersistEntity val =>
Maybe FilterTablePrefix -> SqlBackend -> [Filter val] -> Text
filterClause Maybe FilterTablePrefix
forall a. Maybe a
Nothing SqlBackend
conn [Filter record]
filts
    ord :: SqlBackend -> Text
ord SqlBackend
conn = Maybe FilterTablePrefix -> SqlBackend -> [SelectOpt record] -> Text
forall val.
PersistEntity val =>
Maybe FilterTablePrefix -> SqlBackend -> [SelectOpt val] -> Text
orderClause Maybe FilterTablePrefix
forall a. Maybe a
Nothing SqlBackend
conn [SelectOpt record]
orders
    cols :: SqlBackend -> Text
cols = [Text] -> Text
commaSeparated ([Text] -> Text) -> (SqlBackend -> [Text]) -> SqlBackend -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty Text -> [Text]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (NonEmpty Text -> [Text])
-> (SqlBackend -> NonEmpty Text) -> SqlBackend -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EntityDef -> SqlBackend -> NonEmpty Text
keyAndEntityColumnNames EntityDef
t
    sql :: SqlBackend -> Text
sql SqlBackend
conn = SqlBackend -> (Int, Int) -> Text -> Text
connLimitOffset SqlBackend
conn (Int
limit,Int
offset) (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ [Text] -> Text
forall a. Monoid a => [a] -> a
mconcat
      [ Text
"SELECT "
      , SqlBackend -> Text
cols SqlBackend
conn
      , Text
" FROM "
      , SqlBackend -> EntityDef -> Text
connEscapeTableName SqlBackend
conn EntityDef
t
      , SqlBackend -> Text
wher SqlBackend
conn
      , SqlBackend -> Text
ord SqlBackend
conn
      ]

    getFiltsValues :: SqlBackend -> [Filter val] -> [PersistValue]
getFiltsValues SqlBackend
conn = (Text, [PersistValue]) -> [PersistValue]
forall a b. (a, b) -> b
snd ((Text, [PersistValue]) -> [PersistValue])
-> ([Filter val] -> (Text, [PersistValue]))
-> [Filter val]
-> [PersistValue]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe FilterTablePrefix
-> SqlBackend -> [Filter val] -> (Text, [PersistValue])
forall val.
PersistEntity val =>
Maybe FilterTablePrefix
-> SqlBackend -> [Filter val] -> (Text, [PersistValue])
filterClauseWithVals Maybe FilterTablePrefix
forall a. Maybe a
Nothing SqlBackend
conn
    dummyFromFilts :: [Filter record] -> Maybe record
    dummyFromFilts :: [Filter record] -> Maybe record
dummyFromFilts [Filter record]
_ = Maybe record
forall a. Maybe a
Nothing