{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralisedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UndecidableInstances #-}
module Database.PostgreSQL.Tx.Squeal.Internal
  ( -- * Disclaimer
    -- $disclaimer

    -- ** Internals
    module Database.PostgreSQL.Tx.Squeal.Internal
  ) where

import Control.Exception (Exception(fromException))
import Control.Monad.IO.Class (MonadIO(liftIO))
import Data.Kind (Constraint)
import Database.PostgreSQL.Tx (TxEnv, TxException, TxM, askTxEnv, mapExceptionTx)
import Database.PostgreSQL.Tx.Squeal.Internal.Reexport
import Database.PostgreSQL.Tx.Unsafe (unsafeLookupTxEnvIO, unsafeMkTxException, unsafeRunIOInTxM, unsafeRunTxM)
import GHC.TypeLits (ErrorMessage(Text), TypeError)
import UnliftIO (MonadUnliftIO)
import qualified Data.ByteString.Char8 as Char8
import qualified Database.PostgreSQL.LibPQ as LibPQ
import qualified Squeal.PostgreSQL as Squeal
import qualified UnliftIO

-- | Runtime environment needed to run @squeal-postgresql@ via @postgresql-tx@.
--
-- @since 0.2.0.0
type SquealEnv r =
  (TxEnv SquealConnection r) :: Constraint

-- | Monad type alias for running @squeal-postgresql@ via @postgresql-tx@.
--
-- @since 0.2.0.0
type SquealM a = forall r. (SquealEnv r) => TxM r a

-- | Alias for 'SquealTxM'' but has the 'SquealEnv' constraint applied to @r@.
--
-- @since 0.2.0.0
type SquealTxM (db :: SchemasType) a =
  forall r. (SquealEnv r) => SquealTxM' db r a

-- | A newtype wrapper around 'TxM' which includes the @squeal@ 'SchemasType'
-- parameter @db@. This is used only as type information.
-- You can easily convert 'TxM' to and from 'SquealTxM'' by using the
-- 'SquealTxM'' constructor and 'fromSquealTxM' function, respectively.
--
-- In practice, you will likely prefer to use the 'SquealTxM' type alias
-- as it includes the 'SquealEnv' constraint on @r@.
--
-- @since 0.2.0.0
newtype SquealTxM' (db :: SchemasType) r a =
  SquealTxM
    { -- | Convert a 'SquealTxM'' to a 'TxM'.
      --
      -- @since 0.2.0.0
      SquealTxM' db r a -> TxM r a
fromSquealTxM :: TxM r a
    }
  deriving newtype (a -> SquealTxM' db r b -> SquealTxM' db r a
(a -> b) -> SquealTxM' db r a -> SquealTxM' db r b
(forall a b. (a -> b) -> SquealTxM' db r a -> SquealTxM' db r b)
-> (forall a b. a -> SquealTxM' db r b -> SquealTxM' db r a)
-> Functor (SquealTxM' db r)
forall (db :: SchemasType) r a b.
a -> SquealTxM' db r b -> SquealTxM' db r a
forall (db :: SchemasType) r a b.
(a -> b) -> SquealTxM' db r a -> SquealTxM' db r b
forall a b. a -> SquealTxM' db r b -> SquealTxM' db r a
forall a b. (a -> b) -> SquealTxM' db r a -> SquealTxM' db r b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> SquealTxM' db r b -> SquealTxM' db r a
$c<$ :: forall (db :: SchemasType) r a b.
a -> SquealTxM' db r b -> SquealTxM' db r a
fmap :: (a -> b) -> SquealTxM' db r a -> SquealTxM' db r b
$cfmap :: forall (db :: SchemasType) r a b.
(a -> b) -> SquealTxM' db r a -> SquealTxM' db r b
Functor, Functor (SquealTxM' db r)
a -> SquealTxM' db r a
Functor (SquealTxM' db r)
-> (forall a. a -> SquealTxM' db r a)
-> (forall a b.
    SquealTxM' db r (a -> b) -> SquealTxM' db r a -> SquealTxM' db r b)
-> (forall a b c.
    (a -> b -> c)
    -> SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r c)
-> (forall a b.
    SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r b)
-> (forall a b.
    SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r a)
-> Applicative (SquealTxM' db r)
SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r b
SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r a
SquealTxM' db r (a -> b) -> SquealTxM' db r a -> SquealTxM' db r b
(a -> b -> c)
-> SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r c
forall (db :: SchemasType) r. Functor (SquealTxM' db r)
forall (db :: SchemasType) r a. a -> SquealTxM' db r a
forall (db :: SchemasType) r a b.
SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r a
forall (db :: SchemasType) r a b.
SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r b
forall (db :: SchemasType) r a b.
SquealTxM' db r (a -> b) -> SquealTxM' db r a -> SquealTxM' db r b
forall (db :: SchemasType) r a b c.
(a -> b -> c)
-> SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r c
forall a. a -> SquealTxM' db r a
forall a b.
SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r a
forall a b.
SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r b
forall a b.
SquealTxM' db r (a -> b) -> SquealTxM' db r a -> SquealTxM' db r b
forall a b c.
(a -> b -> c)
-> SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r a
$c<* :: forall (db :: SchemasType) r a b.
SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r a
*> :: SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r b
$c*> :: forall (db :: SchemasType) r a b.
SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r b
liftA2 :: (a -> b -> c)
-> SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r c
$cliftA2 :: forall (db :: SchemasType) r a b c.
(a -> b -> c)
-> SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r c
<*> :: SquealTxM' db r (a -> b) -> SquealTxM' db r a -> SquealTxM' db r b
$c<*> :: forall (db :: SchemasType) r a b.
SquealTxM' db r (a -> b) -> SquealTxM' db r a -> SquealTxM' db r b
pure :: a -> SquealTxM' db r a
$cpure :: forall (db :: SchemasType) r a. a -> SquealTxM' db r a
$cp1Applicative :: forall (db :: SchemasType) r. Functor (SquealTxM' db r)
Applicative, Applicative (SquealTxM' db r)
a -> SquealTxM' db r a
Applicative (SquealTxM' db r)
-> (forall a b.
    SquealTxM' db r a -> (a -> SquealTxM' db r b) -> SquealTxM' db r b)
-> (forall a b.
    SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r b)
-> (forall a. a -> SquealTxM' db r a)
-> Monad (SquealTxM' db r)
SquealTxM' db r a -> (a -> SquealTxM' db r b) -> SquealTxM' db r b
SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r b
forall (db :: SchemasType) r. Applicative (SquealTxM' db r)
forall (db :: SchemasType) r a. a -> SquealTxM' db r a
forall (db :: SchemasType) r a b.
SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r b
forall (db :: SchemasType) r a b.
SquealTxM' db r a -> (a -> SquealTxM' db r b) -> SquealTxM' db r b
forall a. a -> SquealTxM' db r a
forall a b.
SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r b
forall a b.
SquealTxM' db r a -> (a -> SquealTxM' db r b) -> SquealTxM' db r b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: a -> SquealTxM' db r a
$creturn :: forall (db :: SchemasType) r a. a -> SquealTxM' db r a
>> :: SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r b
$c>> :: forall (db :: SchemasType) r a b.
SquealTxM' db r a -> SquealTxM' db r b -> SquealTxM' db r b
>>= :: SquealTxM' db r a -> (a -> SquealTxM' db r b) -> SquealTxM' db r b
$c>>= :: forall (db :: SchemasType) r a b.
SquealTxM' db r a -> (a -> SquealTxM' db r b) -> SquealTxM' db r b
$cp1Monad :: forall (db :: SchemasType) r. Applicative (SquealTxM' db r)
Monad)

-- | The 'SquealTxM'' monad discourages performing arbitrary 'IO' within a
-- transaction, so this instance generates a type error when client code tries
-- to call 'liftIO'.
--
-- Note that we specialize this instance for 'SquealTxM'' rather than derive it
-- via newtype so we can provide a better error message.
--
-- @since 0.2.0.0
instance
  ( TypeError
      ('Text "MonadIO is banned in SquealTxM'; use 'SquealTxM . unsafeRunIOInTxM' if you are sure this is safe IO")
  ) => MonadIO (SquealTxM' db r)
  where
  liftIO :: IO a -> SquealTxM' db r a
liftIO = IO a -> SquealTxM' db r a
forall a. HasCallStack => a
undefined

-- | Used in the 'SquealEnv' to specify the 'LibPQ.Connection' to use.
-- Should produce the same 'LibPQ.Connection' if called multiple times
-- in the same transaction. Usually you will want to use 'mkSquealConnection'
-- to get one.
--
-- @since 0.2.0.0
newtype SquealConnection =
  UnsafeSquealConnection
    { SquealConnection -> IO Connection
unsafeGetLibPQConnection :: IO LibPQ.Connection
    }

-- | Construct a 'SquealConnection' from a 'LibPQ.Connection'.
--
-- @since 0.2.0.0
mkSquealConnection :: LibPQ.Connection -> SquealConnection
mkSquealConnection :: Connection -> SquealConnection
mkSquealConnection Connection
conn = IO Connection -> SquealConnection
UnsafeSquealConnection (Connection -> IO Connection
forall (f :: * -> *) a. Applicative f => a -> f a
pure Connection
conn)

fromSquealException :: SquealException -> TxException
fromSquealException :: SquealException -> TxException
fromSquealException =
  (SquealException -> Maybe String) -> SquealException -> TxException
forall e. Exception e => (e -> Maybe String) -> e -> TxException
unsafeMkTxException \case
    SQLException SQLState { ByteString
sqlStateCode :: SQLState -> ByteString
sqlStateCode :: ByteString
sqlStateCode } -> String -> Maybe String
forall a. a -> Maybe a
Just (String -> Maybe String) -> String -> Maybe String
forall a b. (a -> b) -> a -> b
$ ByteString -> String
Char8.unpack ByteString
sqlStateCode
    SquealException
_ -> Maybe String
forall a. Maybe a
Nothing

unsafeSquealIOTxM
  :: PQ db db IO a
  -> SquealTxM db a
unsafeSquealIOTxM :: PQ db db IO a -> SquealTxM db a
unsafeSquealIOTxM (PQ K Connection db -> IO (K a db)
f) = TxM r a -> SquealTxM' db r a
forall (db :: SchemasType) r a. TxM r a -> SquealTxM' db r a
SquealTxM (TxM r a -> SquealTxM' db r a) -> TxM r a -> SquealTxM' db r a
forall a b. (a -> b) -> a -> b
$ (SquealException -> Maybe TxException) -> TxM r a -> TxM r a
forall e e' r a.
(Exception e, Exception e') =>
(e -> Maybe e') -> TxM r a -> TxM r a
mapExceptionTx (TxException -> Maybe TxException
forall a. a -> Maybe a
Just (TxException -> Maybe TxException)
-> (SquealException -> TxException)
-> SquealException
-> Maybe TxException
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SquealException -> TxException
fromSquealException) do
  UnsafeSquealConnection { IO Connection
unsafeGetLibPQConnection :: IO Connection
unsafeGetLibPQConnection :: SquealConnection -> IO Connection
unsafeGetLibPQConnection } <- TxM r SquealConnection
forall a r. TxEnv a r => TxM r a
askTxEnv
  IO a -> TxM r a
forall a r. IO a -> TxM r a
unsafeRunIOInTxM do
    Connection
conn <- IO Connection
unsafeGetLibPQConnection
    K a
a <- K Connection db -> IO (K a db)
f (Connection -> K Connection db
forall k a (b :: k). a -> K a b
K Connection
conn)
    a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a

unsafeSquealIOTxM1
  :: (x1 -> PQ db db IO a)
  -> x1 -> SquealTxM db a
unsafeSquealIOTxM1 :: (x1 -> PQ db db IO a) -> x1 -> SquealTxM db a
unsafeSquealIOTxM1 x1 -> PQ db db IO a
f x1
x1 = PQ db db IO a -> SquealTxM db a
forall (db :: SchemasType) a. PQ db db IO a -> SquealTxM db a
unsafeSquealIOTxM (PQ db db IO a -> SquealTxM db a)
-> PQ db db IO a -> SquealTxM db a
forall a b. (a -> b) -> a -> b
$ x1 -> PQ db db IO a
f x1
x1

unsafeSquealIOTxM2
  :: (x1 -> x2 -> PQ db db IO a)
  -> x1 -> x2 -> SquealTxM db a
unsafeSquealIOTxM2 :: (x1 -> x2 -> PQ db db IO a) -> x1 -> x2 -> SquealTxM db a
unsafeSquealIOTxM2 x1 -> x2 -> PQ db db IO a
f x1
x1 x2
x2 = PQ db db IO a -> SquealTxM db a
forall (db :: SchemasType) a. PQ db db IO a -> SquealTxM db a
unsafeSquealIOTxM (PQ db db IO a -> SquealTxM db a)
-> PQ db db IO a -> SquealTxM db a
forall a b. (a -> b) -> a -> b
$ x1 -> x2 -> PQ db db IO a
f x1
x1 x2
x2

unsafeSquealIOTxM3
  :: (x1 -> x2 -> x3 -> PQ db db IO a)
  -> x1 -> x2 -> x3 -> SquealTxM db a
unsafeSquealIOTxM3 :: (x1 -> x2 -> x3 -> PQ db db IO a)
-> x1 -> x2 -> x3 -> SquealTxM db a
unsafeSquealIOTxM3 x1 -> x2 -> x3 -> PQ db db IO a
f x1
x1 x2
x2 x3
x3 = PQ db db IO a -> SquealTxM db a
forall (db :: SchemasType) a. PQ db db IO a -> SquealTxM db a
unsafeSquealIOTxM (PQ db db IO a -> SquealTxM db a)
-> PQ db db IO a -> SquealTxM db a
forall a b. (a -> b) -> a -> b
$ x1 -> x2 -> x3 -> PQ db db IO a
f x1
x1 x2
x2 x3
x3

unsafeRunSquealTransaction
  :: forall r a. (SquealEnv r)
  => (forall db. PQ db db IO a -> PQ db db IO a)
  -> r
  -> TxM r a
  -> IO a
unsafeRunSquealTransaction :: (forall (db :: SchemasType). PQ db db IO a -> PQ db db IO a)
-> r -> TxM r a -> IO a
unsafeRunSquealTransaction forall (db :: SchemasType). PQ db db IO a -> PQ db db IO a
f r
r TxM r a
x = do
  UnsafeSquealConnection { IO Connection
unsafeGetLibPQConnection :: IO Connection
unsafeGetLibPQConnection :: SquealConnection -> IO Connection
unsafeGetLibPQConnection } <- r -> IO SquealConnection
forall a r. TxEnv a r => r -> IO a
unsafeLookupTxEnvIO r
r
  Connection
conn <- IO Connection
unsafeGetLibPQConnection
  (PQ Any Any IO a -> K Connection Any -> IO a)
-> K Connection Any -> PQ Any Any IO a -> IO a
forall a b c. (a -> b -> c) -> b -> a -> c
flip PQ Any Any IO a -> K Connection Any -> IO a
forall (m :: * -> *) (db0 :: SchemasType) (db1 :: SchemasType) x.
Functor m =>
PQ db0 db1 m x -> K Connection db0 -> m x
evalPQ (Connection -> K Connection Any
forall k a (b :: k). a -> K a b
K Connection
conn)
    (PQ Any Any IO a -> IO a) -> PQ Any Any IO a -> IO a
forall a b. (a -> b) -> a -> b
$ PQ Any Any IO a -> PQ Any Any IO a
forall (db :: SchemasType). PQ db db IO a -> PQ db db IO a
f
    (PQ Any Any IO a -> PQ Any Any IO a)
-> PQ Any Any IO a -> PQ Any Any IO a
forall a b. (a -> b) -> a -> b
$ (K Connection Any -> IO (K a Any)) -> PQ Any Any IO a
forall (db0 :: SchemasType) (db1 :: SchemasType) (m :: * -> *) x.
(K Connection db0 -> m (K x db1)) -> PQ db0 db1 m x
PQ \K Connection Any
_ -> a -> K a Any
forall k a (b :: k). a -> K a b
K (a -> K a Any) -> IO a -> IO (K a Any)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> r -> TxM r a -> IO a
forall r a. r -> TxM r a -> IO a
unsafeRunTxM r
r TxM r a
x

-- | A variant of 'transactionallyRetry' which takes a predicate
-- for determining when to retry instead of only doing so on
-- @serialization_failure@.
transactionallyRetry'
  :: (MonadUnliftIO m, MonadPQ db m, Exception e)
  => TransactionMode
  -> (e -> Bool)
  -> m a
  -> m a
transactionallyRetry' :: TransactionMode -> (e -> Bool) -> m a -> m a
transactionallyRetry' TransactionMode
mode e -> Bool
shouldRetry m a
action = ((forall a. m a -> m a) -> m a) -> m a
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> m a) -> m b) -> m b
UnliftIO.mask (((forall a. m a -> m a) -> m a) -> m a)
-> ((forall a. m a -> m a) -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> m a
restore ->
  m (Either SomeException a) -> m a
forall (db :: SchemasType) (m :: * -> *) b.
(MonadPQ db m, MonadIO m) =>
m (Either SomeException b) -> m b
loop (m (Either SomeException a) -> m a)
-> (m a -> m (Either SomeException a)) -> m a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> m (Either SomeException a)
forall (m :: * -> *) e a.
(MonadUnliftIO m, Exception e) =>
m a -> m (Either e a)
UnliftIO.try (m a -> m a) -> m a -> m a
forall a b. (a -> b) -> a -> b
$ do
    a
x <- m a -> m a
forall a. m a -> m a
restore m a
action
    Manipulation '[] db '[] '[] -> m ()
forall (db :: SchemasType) (pq :: * -> *).
MonadPQ db pq =>
Manipulation '[] db '[] '[] -> pq ()
Squeal.manipulate_ Manipulation '[] db '[] '[]
forall (db :: SchemasType). Manipulation_ db () ()
commit
    a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
  where
  loop :: m (Either SomeException b) -> m b
loop m (Either SomeException b)
attempt = do
    Manipulation '[] db '[] '[] -> m ()
forall (db :: SchemasType) (pq :: * -> *).
MonadPQ db pq =>
Manipulation '[] db '[] '[] -> pq ()
Squeal.manipulate_ (Manipulation '[] db '[] '[] -> m ())
-> Manipulation '[] db '[] '[] -> m ()
forall a b. (a -> b) -> a -> b
$ TransactionMode -> Manipulation_ db () ()
forall (db :: SchemasType).
TransactionMode -> Manipulation_ db () ()
begin TransactionMode
mode
    m (Either SomeException b)
attempt m (Either SomeException b)
-> (Either SomeException b -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Right b
a -> b -> m b
forall (m :: * -> *) a. Monad m => a -> m a
return b
a
      Left SomeException
e -> do
        Manipulation '[] db '[] '[] -> m ()
forall (db :: SchemasType) (pq :: * -> *).
MonadPQ db pq =>
Manipulation '[] db '[] '[] -> pq ()
Squeal.manipulate_ Manipulation '[] db '[] '[]
forall (db :: SchemasType). Manipulation_ db () ()
rollback
        if (e -> Bool) -> Maybe e -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any e -> Bool
shouldRetry (SomeException -> Maybe e
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e) then
          m (Either SomeException b) -> m b
loop m (Either SomeException b)
attempt
        else
          SomeException -> m b
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
UnliftIO.throwIO SomeException
e

-- $disclaimer
--
-- Changes to this module will not be reflected in the library's version
-- updates.