module Database.PostgreSQL.Query.Functions
       ( -- * Raw query execution
         pgQuery
       , pgQueryWithMasker
       , pgExecute
       , pgExecuteWithMasker
         -- * Transactions
       , pgWithTransaction
       , pgWithSavepoint
       , pgWithTransactionMode
       , pgWithTransactionModeRetry
       , pgWithTransactionSerializable
         -- * Auxiliary
       , pgRepsertRow
       ) where

import Data.Int ( Int64 )
import Database.PostgreSQL.Query.Import
import Database.PostgreSQL.Query.Internal
import Database.PostgreSQL.Query.SqlBuilder
import Database.PostgreSQL.Query.TH
import Database.PostgreSQL.Query.Types
import Database.PostgreSQL.Simple
import Database.PostgreSQL.Simple.Transaction

import qualified Data.Text.Encoding as T

{- | Execute query generated by 'SqlBuilder'. Typical use case:

@
let userName = "Vovka Erohin" :: Text
pgQuery [sqlExp| SELECT id, name FROM users WHERE name = #{userName}|]
@

Or

@
let userName = "Vovka Erohin" :: Text
pgQuery $ Qp "SELECT id, name FROM users WHERE name = ?" [userName]
@

Which is almost the same. In both cases proper value escaping is performed
so you stay protected from sql injections.

-}

pgQuery
  :: (MonadPostgres m, ToSqlBuilder q, FromRow r, HasCallStack)
  => q
  -> m [r]
pgQuery :: q -> m [r]
pgQuery = (HasCallStack => q -> m [r]) -> q -> m [r]
forall a. HasCallStack => (HasCallStack => a) -> a
withFrozenCallStack ((HasCallStack => q -> m [r]) -> q -> m [r])
-> (HasCallStack => q -> m [r]) -> q -> m [r]
forall a b. (a -> b) -> a -> b
$ LogMasker -> q -> m [r]
forall (m :: * -> *) q r.
(MonadPostgres m, ToSqlBuilder q, FromRow r, HasCallStack) =>
LogMasker -> q -> m [r]
pgQueryWithMasker LogMasker
defaultLogMasker

-- | Execute arbitrary query and return count of affected rows
pgExecute
  :: (MonadPostgres m, ToSqlBuilder q, HasCallStack)
  => q
  -> m Int64
pgExecute :: q -> m Int64
pgExecute = (HasCallStack => q -> m Int64) -> q -> m Int64
forall a. HasCallStack => (HasCallStack => a) -> a
withFrozenCallStack ((HasCallStack => q -> m Int64) -> q -> m Int64)
-> (HasCallStack => q -> m Int64) -> q -> m Int64
forall a b. (a -> b) -> a -> b
$ LogMasker -> q -> m Int64
forall (m :: * -> *) q.
(MonadPostgres m, ToSqlBuilder q, HasCallStack) =>
LogMasker -> q -> m Int64
pgExecuteWithMasker LogMasker
defaultLogMasker

pgQueryWithMasker
  :: (MonadPostgres m, ToSqlBuilder q, FromRow r, HasCallStack)
  => LogMasker
  -> q
  -> m [r]
pgQueryWithMasker :: LogMasker -> q -> m [r]
pgQueryWithMasker LogMasker
masker q
q = (HasCallStack => m [r]) -> m [r]
forall a. HasCallStack => (HasCallStack => a) -> a
withFrozenCallStack ((HasCallStack => m [r]) -> m [r])
-> (HasCallStack => m [r]) -> m [r]
forall a b. (a -> b) -> a -> b
$ (Connection -> m [r]) -> m [r]
forall (m :: * -> *) a. HasPostgres m => (Connection -> m a) -> m a
withPGConnection ((Connection -> m [r]) -> m [r]) -> (Connection -> m [r]) -> m [r]
forall a b. (a -> b) -> a -> b
$ \Connection
c -> do
    (Query
queryBs, ByteString
logBs) <- IO (Query, ByteString) -> m (Query, ByteString)
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase (IO (Query, ByteString) -> m (Query, ByteString))
-> IO (Query, ByteString) -> m (Query, ByteString)
forall a b. (a -> b) -> a -> b
$ Connection -> LogMasker -> SqlBuilder -> IO (Query, ByteString)
runSqlBuilder Connection
c LogMasker
masker (SqlBuilder -> IO (Query, ByteString))
-> SqlBuilder -> IO (Query, ByteString)
forall a b. (a -> b) -> a -> b
$ q -> SqlBuilder
forall a. ToSqlBuilder a => a -> SqlBuilder
toSqlBuilder q
q
    Text -> m ()
forall (m :: * -> *). (HasCallStack, MonadLogger m) => Text -> m ()
logDebug (Text -> m ()) -> Text -> m ()
forall a b. (a -> b) -> a -> b
$ ByteString -> Text
T.decodeUtf8 ByteString
logBs
    IO [r] -> m [r]
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase (IO [r] -> m [r]) -> IO [r] -> m [r]
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> IO [r]
forall r. FromRow r => Connection -> Query -> IO [r]
query_ Connection
c Query
queryBs

pgExecuteWithMasker
  :: (MonadPostgres m, ToSqlBuilder q, HasCallStack)
  => LogMasker
  -> q
  -> m Int64
pgExecuteWithMasker :: LogMasker -> q -> m Int64
pgExecuteWithMasker LogMasker
masker q
q = (HasCallStack => m Int64) -> m Int64
forall a. HasCallStack => (HasCallStack => a) -> a
withFrozenCallStack ((HasCallStack => m Int64) -> m Int64)
-> (HasCallStack => m Int64) -> m Int64
forall a b. (a -> b) -> a -> b
$ (Connection -> m Int64) -> m Int64
forall (m :: * -> *) a. HasPostgres m => (Connection -> m a) -> m a
withPGConnection ((Connection -> m Int64) -> m Int64)
-> (Connection -> m Int64) -> m Int64
forall a b. (a -> b) -> a -> b
$ \Connection
c -> do
    (Query
queryBs, ByteString
logBs) <- IO (Query, ByteString) -> m (Query, ByteString)
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase (IO (Query, ByteString) -> m (Query, ByteString))
-> IO (Query, ByteString) -> m (Query, ByteString)
forall a b. (a -> b) -> a -> b
$ Connection -> LogMasker -> SqlBuilder -> IO (Query, ByteString)
runSqlBuilder Connection
c LogMasker
masker (SqlBuilder -> IO (Query, ByteString))
-> SqlBuilder -> IO (Query, ByteString)
forall a b. (a -> b) -> a -> b
$ q -> SqlBuilder
forall a. ToSqlBuilder a => a -> SqlBuilder
toSqlBuilder q
q
    Text -> m ()
forall (m :: * -> *). (HasCallStack, MonadLogger m) => Text -> m ()
logDebug (Text -> m ()) -> Text -> m ()
forall a b. (a -> b) -> a -> b
$ ByteString -> Text
T.decodeUtf8 ByteString
logBs
    IO Int64 -> m Int64
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase (IO Int64 -> m Int64) -> IO Int64 -> m Int64
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> IO Int64
execute_ Connection
c Query
queryBs

-- | Execute all queries inside one transaction. Rollback transaction on exceptions
pgWithTransaction
  :: (HasPostgres m, MonadBaseControl IO m, TransactionSafe m, HasCallStack)
  => (HasCallStack => m a)
  -> m a
pgWithTransaction :: (HasCallStack => m a) -> m a
pgWithTransaction HasCallStack => m a
action = (Connection -> m a) -> m a
forall (m :: * -> *) a. HasPostgres m => (Connection -> m a) -> m a
withPGConnection ((Connection -> m a) -> m a) -> (Connection -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \Connection
con -> do
    (RunInBase m IO -> IO (StM m a)) -> m a
forall (b :: * -> *) (m :: * -> *) a.
MonadBaseControl b m =>
(RunInBase m b -> b (StM m a)) -> m a
control ((RunInBase m IO -> IO (StM m a)) -> m a)
-> (RunInBase m IO -> IO (StM m a)) -> m a
forall a b. (a -> b) -> a -> b
$ \RunInBase m IO
runInIO -> do
        Connection -> IO (StM m a) -> IO (StM m a)
forall a. Connection -> IO a -> IO a
withTransaction Connection
con (IO (StM m a) -> IO (StM m a)) -> IO (StM m a) -> IO (StM m a)
forall a b. (a -> b) -> a -> b
$ m a -> IO (StM m a)
RunInBase m IO
runInIO m a
HasCallStack => m a
action

-- | Same as `pgWithTransaction` but executes queries inside savepoint
pgWithSavepoint
  :: (HasPostgres m, MonadBaseControl IO m, TransactionSafe m, HasCallStack)
  => (HasCallStack => m a)
  -> m a
pgWithSavepoint :: (HasCallStack => m a) -> m a
pgWithSavepoint HasCallStack => m a
action = (Connection -> m a) -> m a
forall (m :: * -> *) a. HasPostgres m => (Connection -> m a) -> m a
withPGConnection ((Connection -> m a) -> m a) -> (Connection -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \Connection
con -> do
    (RunInBase m IO -> IO (StM m a)) -> m a
forall (b :: * -> *) (m :: * -> *) a.
MonadBaseControl b m =>
(RunInBase m b -> b (StM m a)) -> m a
control ((RunInBase m IO -> IO (StM m a)) -> m a)
-> (RunInBase m IO -> IO (StM m a)) -> m a
forall a b. (a -> b) -> a -> b
$ \RunInBase m IO
runInIO -> do
        Connection -> IO (StM m a) -> IO (StM m a)
forall a. Connection -> IO a -> IO a
withSavepoint Connection
con (IO (StM m a) -> IO (StM m a)) -> IO (StM m a) -> IO (StM m a)
forall a b. (a -> b) -> a -> b
$ m a -> IO (StM m a)
RunInBase m IO
runInIO m a
HasCallStack => m a
action

-- | Wrapper for 'withTransactionMode': Execute an action inside a SQL
-- transaction with a given transaction mode.
pgWithTransactionMode
  :: (HasPostgres m, MonadBaseControl IO m, TransactionSafe m, HasCallStack)
  => TransactionMode
  -> (HasCallStack => m a)
  -> m a
pgWithTransactionMode :: TransactionMode -> (HasCallStack => m a) -> m a
pgWithTransactionMode TransactionMode
tmode HasCallStack => m a
ma = (Connection -> m a) -> m a
forall (m :: * -> *) a. HasPostgres m => (Connection -> m a) -> m a
withPGConnection ((Connection -> m a) -> m a) -> (Connection -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \Connection
con -> do
    (RunInBase m IO -> IO (StM m a)) -> m a
forall (b :: * -> *) (m :: * -> *) a.
MonadBaseControl b m =>
(RunInBase m b -> b (StM m a)) -> m a
control ((RunInBase m IO -> IO (StM m a)) -> m a)
-> (RunInBase m IO -> IO (StM m a)) -> m a
forall a b. (a -> b) -> a -> b
$ \RunInBase m IO
runInIO -> do
        TransactionMode -> Connection -> IO (StM m a) -> IO (StM m a)
forall a. TransactionMode -> Connection -> IO a -> IO a
withTransactionMode TransactionMode
tmode Connection
con (IO (StM m a) -> IO (StM m a)) -> IO (StM m a) -> IO (StM m a)
forall a b. (a -> b) -> a -> b
$ m a -> IO (StM m a)
RunInBase m IO
runInIO m a
HasCallStack => m a
ma

-- | Wrapper for 'withTransactionModeRetry': Like 'pgWithTransactionMode',
-- but also takes a custom callback to determine if a transaction
-- should be retried if an SqlError occurs. If the callback returns
-- True, then the transaction will be retried. If the callback returns
-- False, or an exception other than an SqlError occurs then the
-- transaction will be rolled back and the exception rethrown.
pgWithTransactionModeRetry
  :: (HasPostgres m, MonadBaseControl IO m, TransactionSafe m, HasCallStack)
  => TransactionMode
  -> (SqlError -> Bool)
  -> (HasCallStack => m a)
  -> m a
pgWithTransactionModeRetry :: TransactionMode
-> (SqlError -> Bool) -> (HasCallStack => m a) -> m a
pgWithTransactionModeRetry TransactionMode
tmode SqlError -> Bool
epred HasCallStack => m a
ma = (Connection -> m a) -> m a
forall (m :: * -> *) a. HasPostgres m => (Connection -> m a) -> m a
withPGConnection ((Connection -> m a) -> m a) -> (Connection -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \Connection
con -> do
    (RunInBase m IO -> IO (StM m a)) -> m a
forall (b :: * -> *) (m :: * -> *) a.
MonadBaseControl b m =>
(RunInBase m b -> b (StM m a)) -> m a
control ((RunInBase m IO -> IO (StM m a)) -> m a)
-> (RunInBase m IO -> IO (StM m a)) -> m a
forall a b. (a -> b) -> a -> b
$ \RunInBase m IO
runInIO -> do
        TransactionMode
-> (SqlError -> Bool) -> Connection -> IO (StM m a) -> IO (StM m a)
forall a.
TransactionMode -> (SqlError -> Bool) -> Connection -> IO a -> IO a
withTransactionModeRetry TransactionMode
tmode SqlError -> Bool
epred Connection
con (IO (StM m a) -> IO (StM m a)) -> IO (StM m a) -> IO (StM m a)
forall a b. (a -> b) -> a -> b
$ m a -> IO (StM m a)
RunInBase m IO
runInIO m a
HasCallStack => m a
ma

-- | Wrapper for 'withTransactionSerializable': Execute an action
-- inside of a 'Serializable' transaction. If a serialization failure
-- occurs, roll back the transaction and try again. Be warned that
-- this may execute the IO action multiple times.
--
-- A Serializable transaction creates the illusion that your program
-- has exclusive access to the database. This means that, even in a
-- concurrent setting, you can perform queries in sequence without
-- having to worry about what might happen between one statement and
-- the next.
pgWithTransactionSerializable
  :: (HasPostgres m, MonadBaseControl IO m, TransactionSafe m)
  => (HasCallStack => m a)
  -> m a
pgWithTransactionSerializable :: (HasCallStack => m a) -> m a
pgWithTransactionSerializable HasCallStack => m a
ma = (Connection -> m a) -> m a
forall (m :: * -> *) a. HasPostgres m => (Connection -> m a) -> m a
withPGConnection ((Connection -> m a) -> m a) -> (Connection -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \Connection
con -> do
    (RunInBase m IO -> IO (StM m a)) -> m a
forall (b :: * -> *) (m :: * -> *) a.
MonadBaseControl b m =>
(RunInBase m b -> b (StM m a)) -> m a
control ((RunInBase m IO -> IO (StM m a)) -> m a)
-> (RunInBase m IO -> IO (StM m a)) -> m a
forall a b. (a -> b) -> a -> b
$ \RunInBase m IO
runInIO -> do
        Connection -> IO (StM m a) -> IO (StM m a)
forall a. Connection -> IO a -> IO a
withTransactionSerializable Connection
con (IO (StM m a) -> IO (StM m a)) -> IO (StM m a) -> IO (StM m a)
forall a b. (a -> b) -> a -> b
$ m a -> IO (StM m a)
RunInBase m IO
runInIO m a
HasCallStack => m a
ma



{- | Perform repsert of the same row, first trying "update where" then
"insert" with concatenated fields. Which means that if you run

@
pgRepsertRow "emails" (MR [("user_id", mkValue uid)]) (MR [("email", mkValue email)])
@

Then firstly will be performed

@
UPDATE "emails" SET email = 'foo@bar.com' WHERE "user_id" = 1234
@

And if no one row is affected (which is returned by 'pgExecute'), then

@
INSERT INTO "emails" ("user_id", "email") VALUES (1234, 'foo@bar.com')
@

will be performed

-}

pgRepsertRow
  :: ( MonadPostgres m, MonadLogger m
     , ToMarkedRow wrow, ToMarkedRow urow, HasCallStack)
  => FN                         -- ^ Table name
  -> wrow                       -- ^ where condition
  -> urow                       -- ^ update row
  -> m ()
pgRepsertRow :: FN -> wrow -> urow -> m ()
pgRepsertRow FN
tname wrow
wrow urow
urow = do
    let wmr :: MarkedRow
wmr = wrow -> MarkedRow
forall a. ToMarkedRow a => a -> MarkedRow
toMarkedRow wrow
wrow
    Int64
aff <- SqlBuilder -> m Int64
forall (m :: * -> *) q.
(MonadPostgres m, ToSqlBuilder q, HasCallStack) =>
q -> m Int64
pgExecute (SqlBuilder -> m Int64) -> SqlBuilder -> m Int64
forall a b. (a -> b) -> a -> b
$ FN -> urow -> SqlBuilder -> SqlBuilder
forall q flds.
(ToSqlBuilder q, ToMarkedRow flds) =>
FN -> flds -> q -> SqlBuilder
updateTable FN
tname urow
urow
           [sqlExp|WHERE ^{mrToBuilder "AND" wmr}|]
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int64
aff Int64 -> Int64 -> Bool
forall a. Eq a => a -> a -> Bool
== Int64
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        let umr :: MarkedRow
umr = urow -> MarkedRow
forall a. ToMarkedRow a => a -> MarkedRow
toMarkedRow urow
urow
            imr :: MarkedRow
imr = MarkedRow
wmr MarkedRow -> MarkedRow -> MarkedRow
forall a. Semigroup a => a -> a -> a
<> MarkedRow
umr
        Int64
_ <- SqlBuilder -> m Int64
forall (m :: * -> *) q.
(MonadPostgres m, ToSqlBuilder q, HasCallStack) =>
q -> m Int64
pgExecute (SqlBuilder -> m Int64) -> SqlBuilder -> m Int64
forall a b. (a -> b) -> a -> b
$ FN -> MarkedRow -> SqlBuilder
forall b. ToMarkedRow b => FN -> b -> SqlBuilder
insertInto FN
tname MarkedRow
imr
        () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()