{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DerivingStrategies    #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NumericUnderscores    #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE StrictData            #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}

module Control.Monad.Freer.Extras.Beam.Postgres where

import Cardano.BM.Trace (Trace, logDebug)
import Control.Concurrent (threadDelay)
import Control.Exception (throw, try)
import Control.Monad.Freer (Eff, LastMember, Member, type (~>))
import Control.Monad.Freer.Extras.Beam.Common (BeamError (SqlError), BeamLog (..))
import Control.Monad.Freer.Reader (Reader, ask)
import Data.Aeson (FromJSON, ToJSON)
import Data.Default (Default, def)
import Data.Pool (Pool)
import Data.Pool qualified as Pool
import Data.Text qualified as Text
import Data.Word (Word16)
import Database.Beam (MonadIO (liftIO))
import Database.Beam.Postgres (Connection, Pg, runBeamPostgresDebug)
import Database.PostgreSQL.Simple qualified as Postgres
import GHC.Generics (Generic)

data DbConfig =
    DbConfig
    { DbConfig -> Text
dbConfigUser        :: Text.Text
    , DbConfig -> Text
dbConfigPass        :: Text.Text
    , DbConfig -> Text
dbConfigHost        :: Text.Text
    , DbConfig -> Word16
dbConfigPort        :: Word16
    , DbConfig -> Text
dbConfigDatabase    :: Text.Text
    , DbConfig -> Int
dbConfigPoolSize    :: Int
    , DbConfig -> Text
dbConfigMarconiFile :: Text.Text
    }
    deriving (Int -> DbConfig -> ShowS
[DbConfig] -> ShowS
DbConfig -> String
(Int -> DbConfig -> ShowS)
-> (DbConfig -> String) -> ([DbConfig] -> ShowS) -> Show DbConfig
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DbConfig] -> ShowS
$cshowList :: [DbConfig] -> ShowS
show :: DbConfig -> String
$cshow :: DbConfig -> String
showsPrec :: Int -> DbConfig -> ShowS
$cshowsPrec :: Int -> DbConfig -> ShowS
Show, DbConfig -> DbConfig -> Bool
(DbConfig -> DbConfig -> Bool)
-> (DbConfig -> DbConfig -> Bool) -> Eq DbConfig
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: DbConfig -> DbConfig -> Bool
$c/= :: DbConfig -> DbConfig -> Bool
== :: DbConfig -> DbConfig -> Bool
$c== :: DbConfig -> DbConfig -> Bool
Eq, (forall x. DbConfig -> Rep DbConfig x)
-> (forall x. Rep DbConfig x -> DbConfig) -> Generic DbConfig
forall x. Rep DbConfig x -> DbConfig
forall x. DbConfig -> Rep DbConfig x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep DbConfig x -> DbConfig
$cfrom :: forall x. DbConfig -> Rep DbConfig x
Generic)
    deriving anyclass ([DbConfig] -> Value
[DbConfig] -> Encoding
DbConfig -> Value
DbConfig -> Encoding
(DbConfig -> Value)
-> (DbConfig -> Encoding)
-> ([DbConfig] -> Value)
-> ([DbConfig] -> Encoding)
-> ToJSON DbConfig
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> ToJSON a
toEncodingList :: [DbConfig] -> Encoding
$ctoEncodingList :: [DbConfig] -> Encoding
toJSONList :: [DbConfig] -> Value
$ctoJSONList :: [DbConfig] -> Value
toEncoding :: DbConfig -> Encoding
$ctoEncoding :: DbConfig -> Encoding
toJSON :: DbConfig -> Value
$ctoJSON :: DbConfig -> Value
ToJSON, Value -> Parser [DbConfig]
Value -> Parser DbConfig
(Value -> Parser DbConfig)
-> (Value -> Parser [DbConfig]) -> FromJSON DbConfig
forall a.
(Value -> Parser a) -> (Value -> Parser [a]) -> FromJSON a
parseJSONList :: Value -> Parser [DbConfig]
$cparseJSONList :: Value -> Parser [DbConfig]
parseJSON :: Value -> Parser DbConfig
$cparseJSON :: Value -> Parser DbConfig
FromJSON)

instance Default DbConfig where
  def :: DbConfig
def = DbConfig :: Text -> Text -> Text -> Word16 -> Text -> Int -> Text -> DbConfig
DbConfig
        { dbConfigUser :: Text
dbConfigUser = Text
"postgres"
        , dbConfigPass :: Text
dbConfigPass = Text
""
        , dbConfigHost :: Text
dbConfigHost = Text
"localhost"
        , dbConfigPort :: Word16
dbConfigPort = Word16
5432
        , dbConfigDatabase :: Text
dbConfigDatabase = Text
"pab"
        , dbConfigPoolSize :: Int
dbConfigPoolSize = Int
20
        , dbConfigMarconiFile :: Text
dbConfigMarconiFile = Text
"marconi.sqlite"
        }

runBeam ::
  forall effs.
  ( LastMember IO effs
  , Member (Reader (Pool Connection)) effs
  )
  => Trace IO BeamLog
  -> Pg
  ~> Eff effs
runBeam :: Trace IO BeamLog -> Pg ~> Eff effs
runBeam Trace IO BeamLog
trace Pg x
action = do
  Pool Connection
pool <- forall (effs :: [* -> *]).
Member (Reader (Pool Connection)) effs =>
Eff effs (Pool Connection)
forall r (effs :: [* -> *]). Member (Reader r) effs => Eff effs r
ask @(Pool Postgres.Connection)
  IO x -> Eff effs x
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO x -> Eff effs x) -> IO x -> Eff effs x
forall a b. (a -> b) -> a -> b
$ Pool Connection -> (Connection -> IO x) -> IO x
forall a r. Pool a -> (a -> IO r) -> IO r
Pool.withResource Pool Connection
pool ((Connection -> IO x) -> IO x) -> (Connection -> IO x) -> IO x
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> Connection -> Int -> IO x
loop Connection
conn ( Int
5 :: Int )
  where
    loop :: Connection -> Int -> IO x
loop Connection
conn Int
retries = do
      let traceSql :: String -> IO ()
traceSql = Trace IO BeamLog -> BeamLog -> IO ()
forall (m :: * -> *) a. MonadIO m => Trace m a -> a -> m ()
logDebug Trace IO BeamLog
trace (BeamLog -> IO ()) -> (String -> BeamLog) -> String -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> BeamLog
SqlLog
      Either SqlError x
resultEither <- IO x -> IO (Either SqlError x)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO x -> IO (Either SqlError x)) -> IO x -> IO (Either SqlError x)
forall a b. (a -> b) -> a -> b
$ Connection -> IO x -> IO x
forall a. Connection -> IO a -> IO a
Postgres.withTransaction Connection
conn (IO x -> IO x) -> IO x -> IO x
forall a b. (a -> b) -> a -> b
$ (String -> IO ()) -> Connection -> Pg x -> IO x
forall a. (String -> IO ()) -> Connection -> Pg a -> IO a
runBeamPostgresDebug String -> IO ()
traceSql Connection
conn Pg x
action
      case Either SqlError x
resultEither of
          -- 'Database.SQLite.Simple.ErrorError' corresponds to an SQL error or
          -- missing database. When this exception is raised, we suppose it's
          -- because the another transaction was already running.
          Left Postgres.SqlError {} | Int
retries Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 -> do
              Int -> IO ()
threadDelay Int
100_000
              Connection -> Int -> IO x
loop Connection
conn (Int
retries Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
          -- We handle and rethrow errors other than
          -- 'Database.SQLite.Simple.ErrorError'.
          Left SqlError
e -> BeamError -> IO x
forall a e. Exception e => e -> a
throw (BeamError -> IO x) -> BeamError -> IO x
forall a b. (a -> b) -> a -> b
$ Text -> BeamError
SqlError (Text -> BeamError) -> Text -> BeamError
forall a b. (a -> b) -> a -> b
$ String -> Text
Text.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ SqlError -> String
forall a. Show a => a -> String
show SqlError
e
          Right x
v -> x -> IO x
forall (m :: * -> *) a. Monad m => a -> m a
return x
v