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

module Control.Monad.Freer.Extras.Beam.Sqlite 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 Database.Beam (MonadIO (liftIO))
import Database.Beam.Sqlite (SqliteM, runBeamSqliteDebug)
import Database.SQLite.Simple qualified as Sqlite
import GHC.Generics (Generic)

data DbConfig =
    DbConfig
    { DbConfig -> Text
dbConfigFile     :: Text.Text
    -- ^ The path to the sqlite database file. May be absolute or relative.
    , DbConfig -> Int
dbConfigPoolSize :: Int
    -- ^ Max number of concurrent sqlite database connections.
    }
    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] -> Encoding
[DbConfig] -> Value
DbConfig -> Encoding
DbConfig -> Value
(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)

-- | Default database config uses an in-memory sqlite database that is shared
-- between all threads in the process.
defaultDbConfig :: DbConfig
defaultDbConfig :: DbConfig
defaultDbConfig = DbConfig :: Text -> Int -> DbConfig
DbConfig
                  { dbConfigFile :: Text
dbConfigFile = Text
"file::memory:?cache=shared"
                  , dbConfigPoolSize :: Int
dbConfigPoolSize = Int
20
                  }

instance Default DbConfig where
  def :: DbConfig
def = DbConfig
defaultDbConfig

runBeam ::
  forall effs.
  ( LastMember IO effs
  , Member (Reader (Pool Sqlite.Connection)) effs
  )
  => Trace IO BeamLog
  -> SqliteM
  ~> Eff effs
runBeam :: Trace IO BeamLog -> SqliteM ~> Eff effs
runBeam Trace IO BeamLog
trace SqliteM 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 Sqlite.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
Sqlite.withTransaction Connection
conn (IO x -> IO x) -> IO x -> IO x
forall a b. (a -> b) -> a -> b
$ (String -> IO ()) -> Connection -> SqliteM x -> IO x
forall a. (String -> IO ()) -> Connection -> SqliteM a -> IO a
runBeamSqliteDebug String -> IO ()
traceSql Connection
conn SqliteM 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 (Sqlite.SQLError Error
Sqlite.ErrorError Text
_ Text
_) | 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