{-# LANGUAGE ApplicativeDo #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE StrictData #-}

-- | Database access for your @App@
module Freckle.App.Database
  ( HasSqlPool(..)
  , SqlPool
  , makePostgresPool
  , makePostgresPoolWith
  , runDB
  , PostgresConnectionConf(..)
  , PostgresPasswordSource(..)
  , PostgresPassword(..)
  , PostgresStatementTimeout(..)
  , postgresStatementTimeoutMilliseconds
  , envParseDatabaseConf
  , envPostgresPasswordSource
  ) where

import Freckle.App.Prelude

import Blammo.Logging
import qualified Control.Immortal as Immortal
import Control.Monad.Reader
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as BS8
import qualified Data.ByteString.Lazy as BSL
import Data.Char (isDigit)
import Data.Pool
import qualified Data.Text as T
import Database.Persist.Postgresql
  ( SqlBackend
  , SqlPersistT
  , createPostgresqlPoolModified
  , createSqlPool
  , openSimpleConn
  , runSqlPool
  )
import Database.PostgreSQL.Simple
  (Connection, Only(..), connectPostgreSQL, execute)
import Database.PostgreSQL.Simple.SqlQQ (sql)
import qualified Freckle.App.Env as Env
import qualified Prelude as Unsafe (read)
import System.Process.Typed (proc, readProcessStdout_)
import UnliftIO.Concurrent (threadDelay)
import UnliftIO.Exception (displayException)
import UnliftIO.IORef

type SqlPool = Pool SqlBackend

class HasSqlPool app where
  getSqlPool :: app -> SqlPool

instance HasSqlPool SqlPool where
  getSqlPool :: SqlPool -> SqlPool
getSqlPool = SqlPool -> SqlPool
forall a. a -> a
id

makePostgresPool :: (MonadUnliftIO m, MonadLoggerIO m) => m SqlPool
makePostgresPool :: m SqlPool
makePostgresPool = do
  PostgresConnectionConf
conf <- IO PostgresConnectionConf -> m PostgresConnectionConf
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO PostgresConnectionConf -> m PostgresConnectionConf)
-> IO PostgresConnectionConf -> m PostgresConnectionConf
forall a b. (a -> b) -> a -> b
$ do
    PostgresPasswordSource
postgresPasswordSource <- (Info Error -> Info Error)
-> Parser Error PostgresPasswordSource -> IO PostgresPasswordSource
forall e a.
AsUnset e =>
(Info Error -> Info e) -> Parser e a -> IO a
Env.parse Info Error -> Info Error
forall a. a -> a
id (Parser Error PostgresPasswordSource -> IO PostgresPasswordSource)
-> Parser Error PostgresPasswordSource -> IO PostgresPasswordSource
forall a b. (a -> b) -> a -> b
$ Parser Error PostgresPasswordSource
-> Parser Error PostgresPasswordSource
forall e a. Parser e a -> Parser e a
Env.kept Parser Error PostgresPasswordSource
envPostgresPasswordSource
    (Info Error -> Info Error)
-> Parser Error PostgresConnectionConf -> IO PostgresConnectionConf
forall e a.
AsUnset e =>
(Info Error -> Info e) -> Parser e a -> IO a
Env.parse Info Error -> Info Error
forall a. a -> a
id (Parser Error PostgresConnectionConf -> IO PostgresConnectionConf)
-> Parser Error PostgresConnectionConf -> IO PostgresConnectionConf
forall a b. (a -> b) -> a -> b
$ Parser Error PostgresConnectionConf
-> Parser Error PostgresConnectionConf
forall e a. Parser e a -> Parser e a
Env.kept (Parser Error PostgresConnectionConf
 -> Parser Error PostgresConnectionConf)
-> Parser Error PostgresConnectionConf
-> Parser Error PostgresConnectionConf
forall a b. (a -> b) -> a -> b
$ PostgresPasswordSource -> Parser Error PostgresConnectionConf
envParseDatabaseConf PostgresPasswordSource
postgresPasswordSource
  PostgresConnectionConf -> m SqlPool
forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
PostgresConnectionConf -> m SqlPool
makePostgresPoolWith PostgresConnectionConf
conf

runDB
  :: (HasSqlPool app, MonadUnliftIO m, MonadReader app m)
  => SqlPersistT m a
  -> m a
runDB :: SqlPersistT m a -> m a
runDB SqlPersistT m a
action = do
  SqlPool
pool <- (app -> SqlPool) -> m SqlPool
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks app -> SqlPool
forall app. HasSqlPool app => app -> SqlPool
getSqlPool
  SqlPersistT m a -> SqlPool -> m a
forall backend (m :: * -> *) a.
(MonadUnliftIO m, BackendCompatible SqlBackend backend) =>
ReaderT backend m a -> Pool backend -> m a
runSqlPool SqlPersistT m a
action SqlPool
pool

data PostgresConnectionConf = PostgresConnectionConf
  { PostgresConnectionConf -> String
pccHost :: String
  , PostgresConnectionConf -> Int
pccPort :: Int
  , PostgresConnectionConf -> String
pccUser :: String
  , PostgresConnectionConf -> PostgresPassword
pccPassword :: PostgresPassword
  , PostgresConnectionConf -> String
pccDatabase :: String
  , PostgresConnectionConf -> Int
pccPoolSize :: Int
  , PostgresConnectionConf -> PostgresStatementTimeout
pccStatementTimeout :: PostgresStatementTimeout
  }
  deriving stock (Int -> PostgresConnectionConf -> ShowS
[PostgresConnectionConf] -> ShowS
PostgresConnectionConf -> String
(Int -> PostgresConnectionConf -> ShowS)
-> (PostgresConnectionConf -> String)
-> ([PostgresConnectionConf] -> ShowS)
-> Show PostgresConnectionConf
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PostgresConnectionConf] -> ShowS
$cshowList :: [PostgresConnectionConf] -> ShowS
show :: PostgresConnectionConf -> String
$cshow :: PostgresConnectionConf -> String
showsPrec :: Int -> PostgresConnectionConf -> ShowS
$cshowsPrec :: Int -> PostgresConnectionConf -> ShowS
Show, PostgresConnectionConf -> PostgresConnectionConf -> Bool
(PostgresConnectionConf -> PostgresConnectionConf -> Bool)
-> (PostgresConnectionConf -> PostgresConnectionConf -> Bool)
-> Eq PostgresConnectionConf
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PostgresConnectionConf -> PostgresConnectionConf -> Bool
$c/= :: PostgresConnectionConf -> PostgresConnectionConf -> Bool
== :: PostgresConnectionConf -> PostgresConnectionConf -> Bool
$c== :: PostgresConnectionConf -> PostgresConnectionConf -> Bool
Eq)

data PostgresPasswordSource
  = PostgresPasswordSourceIamAuth
  | PostgresPasswordSourceEnv
  deriving stock (Int -> PostgresPasswordSource -> ShowS
[PostgresPasswordSource] -> ShowS
PostgresPasswordSource -> String
(Int -> PostgresPasswordSource -> ShowS)
-> (PostgresPasswordSource -> String)
-> ([PostgresPasswordSource] -> ShowS)
-> Show PostgresPasswordSource
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PostgresPasswordSource] -> ShowS
$cshowList :: [PostgresPasswordSource] -> ShowS
show :: PostgresPasswordSource -> String
$cshow :: PostgresPasswordSource -> String
showsPrec :: Int -> PostgresPasswordSource -> ShowS
$cshowsPrec :: Int -> PostgresPasswordSource -> ShowS
Show, PostgresPasswordSource -> PostgresPasswordSource -> Bool
(PostgresPasswordSource -> PostgresPasswordSource -> Bool)
-> (PostgresPasswordSource -> PostgresPasswordSource -> Bool)
-> Eq PostgresPasswordSource
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PostgresPasswordSource -> PostgresPasswordSource -> Bool
$c/= :: PostgresPasswordSource -> PostgresPasswordSource -> Bool
== :: PostgresPasswordSource -> PostgresPasswordSource -> Bool
$c== :: PostgresPasswordSource -> PostgresPasswordSource -> Bool
Eq)

data PostgresPassword
  = PostgresPasswordIamAuth
  | PostgresPasswordStatic String
  deriving stock (Int -> PostgresPassword -> ShowS
[PostgresPassword] -> ShowS
PostgresPassword -> String
(Int -> PostgresPassword -> ShowS)
-> (PostgresPassword -> String)
-> ([PostgresPassword] -> ShowS)
-> Show PostgresPassword
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PostgresPassword] -> ShowS
$cshowList :: [PostgresPassword] -> ShowS
show :: PostgresPassword -> String
$cshow :: PostgresPassword -> String
showsPrec :: Int -> PostgresPassword -> ShowS
$cshowsPrec :: Int -> PostgresPassword -> ShowS
Show, PostgresPassword -> PostgresPassword -> Bool
(PostgresPassword -> PostgresPassword -> Bool)
-> (PostgresPassword -> PostgresPassword -> Bool)
-> Eq PostgresPassword
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PostgresPassword -> PostgresPassword -> Bool
$c/= :: PostgresPassword -> PostgresPassword -> Bool
== :: PostgresPassword -> PostgresPassword -> Bool
$c== :: PostgresPassword -> PostgresPassword -> Bool
Eq)

data PostgresStatementTimeout
  = PostgresStatementTimeoutSeconds Int
  | PostgresStatementTimeoutMilliseconds Int
  deriving stock (Int -> PostgresStatementTimeout -> ShowS
[PostgresStatementTimeout] -> ShowS
PostgresStatementTimeout -> String
(Int -> PostgresStatementTimeout -> ShowS)
-> (PostgresStatementTimeout -> String)
-> ([PostgresStatementTimeout] -> ShowS)
-> Show PostgresStatementTimeout
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PostgresStatementTimeout] -> ShowS
$cshowList :: [PostgresStatementTimeout] -> ShowS
show :: PostgresStatementTimeout -> String
$cshow :: PostgresStatementTimeout -> String
showsPrec :: Int -> PostgresStatementTimeout -> ShowS
$cshowsPrec :: Int -> PostgresStatementTimeout -> ShowS
Show, PostgresStatementTimeout -> PostgresStatementTimeout -> Bool
(PostgresStatementTimeout -> PostgresStatementTimeout -> Bool)
-> (PostgresStatementTimeout -> PostgresStatementTimeout -> Bool)
-> Eq PostgresStatementTimeout
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PostgresStatementTimeout -> PostgresStatementTimeout -> Bool
$c/= :: PostgresStatementTimeout -> PostgresStatementTimeout -> Bool
== :: PostgresStatementTimeout -> PostgresStatementTimeout -> Bool
$c== :: PostgresStatementTimeout -> PostgresStatementTimeout -> Bool
Eq)

postgresStatementTimeoutMilliseconds :: PostgresStatementTimeout -> Int
postgresStatementTimeoutMilliseconds :: PostgresStatementTimeout -> Int
postgresStatementTimeoutMilliseconds = \case
  PostgresStatementTimeoutSeconds Int
s -> Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000
  PostgresStatementTimeoutMilliseconds Int
ms -> Int
ms

-- | Read @PGSTATEMENTTIMEOUT@ as seconds or milliseconds
--
-- >>> readPostgresStatementTimeout "10"
-- Right (PostgresStatementTimeoutSeconds 10)
--
-- >>> readPostgresStatementTimeout "10s"
-- Right (PostgresStatementTimeoutSeconds 10)
--
-- >>> readPostgresStatementTimeout "10ms"
-- Right (PostgresStatementTimeoutMilliseconds 10)
--
-- >>> readPostgresStatementTimeout "20m"
-- Left "..."
--
-- >>> readPostgresStatementTimeout "2m0"
-- Left "..."
--
readPostgresStatementTimeout
  :: String -> Either String PostgresStatementTimeout
readPostgresStatementTimeout :: String -> Either String PostgresStatementTimeout
readPostgresStatementTimeout String
x = case (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
span Char -> Bool
isDigit String
x of
  (String
"", String
_) -> String -> Either String PostgresStatementTimeout
forall a b. a -> Either a b
Left String
"must be {digits}(s|ms)"
  (String
digits, String
"") -> PostgresStatementTimeout -> Either String PostgresStatementTimeout
forall a b. b -> Either a b
Right (PostgresStatementTimeout
 -> Either String PostgresStatementTimeout)
-> PostgresStatementTimeout
-> Either String PostgresStatementTimeout
forall a b. (a -> b) -> a -> b
$ Int -> PostgresStatementTimeout
PostgresStatementTimeoutSeconds (Int -> PostgresStatementTimeout)
-> Int -> PostgresStatementTimeout
forall a b. (a -> b) -> a -> b
$ String -> Int
forall a. Read a => String -> a
Unsafe.read String
digits
  (String
digits, String
"s") -> PostgresStatementTimeout -> Either String PostgresStatementTimeout
forall a b. b -> Either a b
Right (PostgresStatementTimeout
 -> Either String PostgresStatementTimeout)
-> PostgresStatementTimeout
-> Either String PostgresStatementTimeout
forall a b. (a -> b) -> a -> b
$ Int -> PostgresStatementTimeout
PostgresStatementTimeoutSeconds (Int -> PostgresStatementTimeout)
-> Int -> PostgresStatementTimeout
forall a b. (a -> b) -> a -> b
$ String -> Int
forall a. Read a => String -> a
Unsafe.read String
digits
  (String
digits, String
"ms") ->
    PostgresStatementTimeout -> Either String PostgresStatementTimeout
forall a b. b -> Either a b
Right (PostgresStatementTimeout
 -> Either String PostgresStatementTimeout)
-> PostgresStatementTimeout
-> Either String PostgresStatementTimeout
forall a b. (a -> b) -> a -> b
$ Int -> PostgresStatementTimeout
PostgresStatementTimeoutMilliseconds (Int -> PostgresStatementTimeout)
-> Int -> PostgresStatementTimeout
forall a b. (a -> b) -> a -> b
$ String -> Int
forall a. Read a => String -> a
Unsafe.read String
digits
  (String, String)
_ -> String -> Either String PostgresStatementTimeout
forall a b. a -> Either a b
Left String
"must be {digits}(s|ms)"

envPostgresPasswordSource :: Env.Parser Env.Error PostgresPasswordSource
envPostgresPasswordSource :: Parser Error PostgresPasswordSource
envPostgresPasswordSource = Off PostgresPasswordSource
-> On PostgresPasswordSource
-> String
-> Mod Flag PostgresPasswordSource
-> Parser Error PostgresPasswordSource
forall a. Off a -> On a -> String -> Mod Flag a -> Parser Error a
Env.flag
  (PostgresPasswordSource -> Off PostgresPasswordSource
forall a. a -> Off a
Env.Off PostgresPasswordSource
PostgresPasswordSourceEnv)
  (PostgresPasswordSource -> On PostgresPasswordSource
forall a. a -> On a
Env.On PostgresPasswordSource
PostgresPasswordSourceIamAuth)
  String
"USE_RDS_IAM_AUTH"
  Mod Flag PostgresPasswordSource
forall a. Monoid a => a
mempty

envParseDatabaseConf
  :: PostgresPasswordSource -> Env.Parser Env.Error PostgresConnectionConf
envParseDatabaseConf :: PostgresPasswordSource -> Parser Error PostgresConnectionConf
envParseDatabaseConf PostgresPasswordSource
source = do
  String
user <- Reader Error String
-> String -> Mod Var String -> Parser Error String
forall e a.
AsUnset e =>
Reader e a -> String -> Mod Var a -> Parser e a
Env.var Reader Error String
forall e s. (AsEmpty e, IsString s) => Reader e s
Env.nonempty String
"PGUSER" Mod Var String
forall a. Monoid a => a
mempty
  PostgresPassword
password <- case PostgresPasswordSource
source of
    PostgresPasswordSource
PostgresPasswordSourceIamAuth -> PostgresPassword -> Parser Error PostgresPassword
forall (f :: * -> *) a. Applicative f => a -> f a
pure PostgresPassword
PostgresPasswordIamAuth
    PostgresPasswordSource
PostgresPasswordSourceEnv ->
      String -> PostgresPassword
PostgresPasswordStatic (String -> PostgresPassword)
-> Parser Error String -> Parser Error PostgresPassword
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Reader Error String
-> String -> Mod Var String -> Parser Error String
forall e a.
AsUnset e =>
Reader e a -> String -> Mod Var a -> Parser e a
Env.var Reader Error String
forall e s. (AsEmpty e, IsString s) => Reader e s
Env.nonempty String
"PGPASSWORD" Mod Var String
forall a. Monoid a => a
mempty
  String
host <- Reader Error String
-> String -> Mod Var String -> Parser Error String
forall e a.
AsUnset e =>
Reader e a -> String -> Mod Var a -> Parser e a
Env.var Reader Error String
forall e s. (AsEmpty e, IsString s) => Reader e s
Env.nonempty String
"PGHOST" Mod Var String
forall a. Monoid a => a
mempty
  String
database <- Reader Error String
-> String -> Mod Var String -> Parser Error String
forall e a.
AsUnset e =>
Reader e a -> String -> Mod Var a -> Parser e a
Env.var Reader Error String
forall e s. (AsEmpty e, IsString s) => Reader e s
Env.nonempty String
"PGDATABASE" Mod Var String
forall a. Monoid a => a
mempty
  Int
port <- Reader Error Int -> String -> Mod Var Int -> Parser Error Int
forall e a.
AsUnset e =>
Reader e a -> String -> Mod Var a -> Parser e a
Env.var Reader Error Int
forall e a. (AsUnread e, Read a) => Reader e a
Env.auto String
"PGPORT" Mod Var Int
forall a. Monoid a => a
mempty
  Int
poolSize <- Reader Error Int -> String -> Mod Var Int -> Parser Error Int
forall e a.
AsUnset e =>
Reader e a -> String -> Mod Var a -> Parser e a
Env.var Reader Error Int
forall e a. (AsUnread e, Read a) => Reader e a
Env.auto String
"PGPOOLSIZE" (Mod Var Int -> Parser Error Int)
-> Mod Var Int -> Parser Error Int
forall a b. (a -> b) -> a -> b
$ Int -> Mod Var Int
forall a. a -> Mod Var a
Env.def Int
10
  PostgresStatementTimeout
statementTimeout <-
    Reader Error PostgresStatementTimeout
-> String
-> Mod Var PostgresStatementTimeout
-> Parser Error PostgresStatementTimeout
forall e a.
AsUnset e =>
Reader e a -> String -> Mod Var a -> Parser e a
Env.var ((String -> Either String PostgresStatementTimeout)
-> Reader Error PostgresStatementTimeout
forall a. (String -> Either String a) -> Reader Error a
Env.eitherReader String -> Either String PostgresStatementTimeout
readPostgresStatementTimeout) String
"PGSTATEMENTTIMEOUT"
      (Mod Var PostgresStatementTimeout
 -> Parser Error PostgresStatementTimeout)
-> Mod Var PostgresStatementTimeout
-> Parser Error PostgresStatementTimeout
forall a b. (a -> b) -> a -> b
$ PostgresStatementTimeout -> Mod Var PostgresStatementTimeout
forall a. a -> Mod Var a
Env.def (Int -> PostgresStatementTimeout
PostgresStatementTimeoutSeconds Int
120)
  pure PostgresConnectionConf :: String
-> Int
-> String
-> PostgresPassword
-> String
-> Int
-> PostgresStatementTimeout
-> PostgresConnectionConf
PostgresConnectionConf
    { pccHost :: String
pccHost = String
host
    , pccPort :: Int
pccPort = Int
port
    , pccUser :: String
pccUser = String
user
    , pccPassword :: PostgresPassword
pccPassword = PostgresPassword
password
    , pccDatabase :: String
pccDatabase = String
database
    , pccPoolSize :: Int
pccPoolSize = Int
poolSize
    , pccStatementTimeout :: PostgresStatementTimeout
pccStatementTimeout = PostgresStatementTimeout
statementTimeout
    }

data AuroraIamToken = AuroraIamToken
  { AuroraIamToken -> Text
aitToken :: Text
  , AuroraIamToken -> UTCTime
aitCreatedAt :: UTCTime
  , AuroraIamToken -> PostgresConnectionConf
aitPostgresConnectionConf :: PostgresConnectionConf
  }
  deriving stock (Int -> AuroraIamToken -> ShowS
[AuroraIamToken] -> ShowS
AuroraIamToken -> String
(Int -> AuroraIamToken -> ShowS)
-> (AuroraIamToken -> String)
-> ([AuroraIamToken] -> ShowS)
-> Show AuroraIamToken
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AuroraIamToken] -> ShowS
$cshowList :: [AuroraIamToken] -> ShowS
show :: AuroraIamToken -> String
$cshow :: AuroraIamToken -> String
showsPrec :: Int -> AuroraIamToken -> ShowS
$cshowsPrec :: Int -> AuroraIamToken -> ShowS
Show, AuroraIamToken -> AuroraIamToken -> Bool
(AuroraIamToken -> AuroraIamToken -> Bool)
-> (AuroraIamToken -> AuroraIamToken -> Bool) -> Eq AuroraIamToken
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: AuroraIamToken -> AuroraIamToken -> Bool
$c/= :: AuroraIamToken -> AuroraIamToken -> Bool
== :: AuroraIamToken -> AuroraIamToken -> Bool
$c== :: AuroraIamToken -> AuroraIamToken -> Bool
Eq)

createAuroraIamToken :: MonadIO m => PostgresConnectionConf -> m AuroraIamToken
createAuroraIamToken :: PostgresConnectionConf -> m AuroraIamToken
createAuroraIamToken aitPostgresConnectionConf :: PostgresConnectionConf
aitPostgresConnectionConf@PostgresConnectionConf {Int
String
PostgresStatementTimeout
PostgresPassword
pccStatementTimeout :: PostgresStatementTimeout
pccPoolSize :: Int
pccDatabase :: String
pccPassword :: PostgresPassword
pccUser :: String
pccPort :: Int
pccHost :: String
pccStatementTimeout :: PostgresConnectionConf -> PostgresStatementTimeout
pccPoolSize :: PostgresConnectionConf -> Int
pccDatabase :: PostgresConnectionConf -> String
pccPassword :: PostgresConnectionConf -> PostgresPassword
pccUser :: PostgresConnectionConf -> String
pccPort :: PostgresConnectionConf -> Int
pccHost :: PostgresConnectionConf -> String
..} = do
  -- TODO: Consider recording how long creating an auth token takes
  -- somewhere, even if it is just in the logs, so we get an idea of how long
  -- it takes in prod.
  Text
aitToken <- Text -> Text
T.strip (Text -> Text) -> (ByteString -> Text) -> ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Text
decodeUtf8 (ByteString -> Text)
-> (ByteString -> ByteString) -> ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
BSL.toStrict (ByteString -> Text) -> m ByteString -> m Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ProcessConfig () () () -> m ByteString
forall (m :: * -> *) stdin stdoutIgnored stderr.
MonadIO m =>
ProcessConfig stdin stdoutIgnored stderr -> m ByteString
readProcessStdout_
    (String -> [String] -> ProcessConfig () () ()
proc
      String
"aws"
      [ String
"rds"
      , String
"generate-db-auth-token"
      , String
"--hostname"
      , String
pccHost
      , String
"--port"
      , Int -> String
forall a. Show a => a -> String
show Int
pccPort
      , String
"--username"
      , String
pccUser
      ]
    )
  UTCTime
aitCreatedAt <- IO UTCTime -> m UTCTime
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UTCTime
getCurrentTime
  pure AuroraIamToken :: Text -> UTCTime -> PostgresConnectionConf -> AuroraIamToken
AuroraIamToken { UTCTime
Text
PostgresConnectionConf
aitCreatedAt :: UTCTime
aitToken :: Text
aitPostgresConnectionConf :: PostgresConnectionConf
aitPostgresConnectionConf :: PostgresConnectionConf
aitCreatedAt :: UTCTime
aitToken :: Text
.. }

-- | Spawns a thread that refreshes the IAM auth token every minute
--
-- The IAM auth token lasts 15 minutes, but we refresh it every minute just to
-- be super safe.
--
spawnIamTokenRefreshThread
  :: (MonadUnliftIO m, MonadLogger m)
  => PostgresConnectionConf
  -> m (IORef AuroraIamToken)
spawnIamTokenRefreshThread :: PostgresConnectionConf -> m (IORef AuroraIamToken)
spawnIamTokenRefreshThread PostgresConnectionConf
conf = do
  Message -> m ()
forall (m :: * -> *).
(HasCallStack, MonadLogger m) =>
Message -> m ()
logInfo Message
"Spawning thread to refresh IAM auth token"
  IORef AuroraIamToken
tokenIORef <- AuroraIamToken -> m (IORef AuroraIamToken)
forall (m :: * -> *) a. MonadIO m => a -> m (IORef a)
newIORef (AuroraIamToken -> m (IORef AuroraIamToken))
-> m AuroraIamToken -> m (IORef AuroraIamToken)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PostgresConnectionConf -> m AuroraIamToken
forall (m :: * -> *).
MonadIO m =>
PostgresConnectionConf -> m AuroraIamToken
createAuroraIamToken PostgresConnectionConf
conf
  m Thread -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Thread -> m ()) -> m Thread -> m ()
forall a b. (a -> b) -> a -> b
$ (Thread -> m ()) -> m Thread
forall (m :: * -> *).
MonadUnliftIO m =>
(Thread -> m ()) -> m Thread
Immortal.create ((Thread -> m ()) -> m Thread) -> (Thread -> m ()) -> m Thread
forall a b. (a -> b) -> a -> b
$ \Thread
_ -> (Either SomeException () -> m ()) -> m () -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
(Either SomeException () -> m ()) -> m () -> m ()
Immortal.onFinish Either SomeException () -> m ()
forall (m :: * -> *) e.
(MonadLogger m, Exception e) =>
Either e () -> m ()
onFinishCallback (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    Message -> m ()
forall (m :: * -> *).
(HasCallStack, MonadLogger m) =>
Message -> m ()
logDebug Message
"Refreshing IAM auth token"
    PostgresConnectionConf -> IORef AuroraIamToken -> m ()
forall (m :: * -> *).
MonadIO m =>
PostgresConnectionConf -> IORef AuroraIamToken -> m ()
refreshIamToken PostgresConnectionConf
conf IORef AuroraIamToken
tokenIORef
    Int -> m ()
forall (m :: * -> *). MonadIO m => Int -> m ()
threadDelay Int
forall a. Num a => a
oneMinuteInMicroseconds
  pure IORef AuroraIamToken
tokenIORef
 where
  oneMinuteInMicroseconds :: a
oneMinuteInMicroseconds = a
60 a -> a -> a
forall a. Num a => a -> a -> a
* a
1000000

  onFinishCallback :: Either e () -> m ()
onFinishCallback = \case
    Left e
ex ->
      Message -> m ()
forall (m :: * -> *).
(HasCallStack, MonadLogger m) =>
Message -> m ()
logError
        (Message -> m ()) -> Message -> m ()
forall a b. (a -> b) -> a -> b
$ Text
"Error refreshing IAM auth token"
        Text -> [SeriesElem] -> Message
:# [Key
"exception" Key -> String -> SeriesElem
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= e -> String
forall e. Exception e => e -> String
displayException e
ex]
    Right () -> () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

refreshIamToken
  :: MonadIO m => PostgresConnectionConf -> IORef AuroraIamToken -> m ()
refreshIamToken :: PostgresConnectionConf -> IORef AuroraIamToken -> m ()
refreshIamToken PostgresConnectionConf
conf IORef AuroraIamToken
tokenIORef = do
  AuroraIamToken
token' <- PostgresConnectionConf -> m AuroraIamToken
forall (m :: * -> *).
MonadIO m =>
PostgresConnectionConf -> m AuroraIamToken
createAuroraIamToken PostgresConnectionConf
conf
  IORef AuroraIamToken -> AuroraIamToken -> m ()
forall (m :: * -> *) a. MonadIO m => IORef a -> a -> m ()
writeIORef IORef AuroraIamToken
tokenIORef AuroraIamToken
token'

setTimeout :: MonadIO m => PostgresConnectionConf -> Connection -> m ()
setTimeout :: PostgresConnectionConf -> Connection -> m ()
setTimeout PostgresConnectionConf {Int
String
PostgresStatementTimeout
PostgresPassword
pccStatementTimeout :: PostgresStatementTimeout
pccPoolSize :: Int
pccDatabase :: String
pccPassword :: PostgresPassword
pccUser :: String
pccPort :: Int
pccHost :: String
pccStatementTimeout :: PostgresConnectionConf -> PostgresStatementTimeout
pccPoolSize :: PostgresConnectionConf -> Int
pccDatabase :: PostgresConnectionConf -> String
pccPassword :: PostgresConnectionConf -> PostgresPassword
pccUser :: PostgresConnectionConf -> String
pccPort :: PostgresConnectionConf -> Int
pccHost :: PostgresConnectionConf -> String
..} Connection
conn = do
  let timeoutMillis :: Int
timeoutMillis = PostgresStatementTimeout -> Int
postgresStatementTimeoutMilliseconds PostgresStatementTimeout
pccStatementTimeout
  m Int64 -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Int64 -> m ()) -> m Int64 -> m ()
forall a b. (a -> b) -> a -> b
$ IO Int64 -> m Int64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int64 -> m Int64) -> IO Int64 -> m Int64
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> Only Int -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
execute
    Connection
conn
    [sql| SET statement_timeout = ? |]
    (Int -> Only Int
forall a. a -> Only a
Only Int
timeoutMillis)

makePostgresPoolWith
  :: (MonadUnliftIO m, MonadLoggerIO m) => PostgresConnectionConf -> m SqlPool
makePostgresPoolWith :: PostgresConnectionConf -> m SqlPool
makePostgresPoolWith conf :: PostgresConnectionConf
conf@PostgresConnectionConf {Int
String
PostgresStatementTimeout
PostgresPassword
pccStatementTimeout :: PostgresStatementTimeout
pccPoolSize :: Int
pccDatabase :: String
pccPassword :: PostgresPassword
pccUser :: String
pccPort :: Int
pccHost :: String
pccStatementTimeout :: PostgresConnectionConf -> PostgresStatementTimeout
pccPoolSize :: PostgresConnectionConf -> Int
pccDatabase :: PostgresConnectionConf -> String
pccPassword :: PostgresConnectionConf -> PostgresPassword
pccUser :: PostgresConnectionConf -> String
pccPort :: PostgresConnectionConf -> Int
pccHost :: PostgresConnectionConf -> String
..} = case PostgresPassword
pccPassword of
  PostgresPassword
PostgresPasswordIamAuth -> PostgresConnectionConf -> m SqlPool
forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
PostgresConnectionConf -> m SqlPool
makePostgresPoolWithIamAuth PostgresConnectionConf
conf
  PostgresPasswordStatic String
password -> (Connection -> IO ()) -> ByteString -> Int -> m SqlPool
forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
(Connection -> IO ()) -> ByteString -> Int -> m SqlPool
createPostgresqlPoolModified
    (PostgresConnectionConf -> Connection -> IO ()
forall (m :: * -> *).
MonadIO m =>
PostgresConnectionConf -> Connection -> m ()
setTimeout PostgresConnectionConf
conf)
    (PostgresConnectionConf -> String -> ByteString
postgresConnectionString PostgresConnectionConf
conf String
password)
    Int
pccPoolSize

-- | Creates a PostgreSQL pool using IAM auth for the password
makePostgresPoolWithIamAuth
  :: (MonadUnliftIO m, MonadLoggerIO m) => PostgresConnectionConf -> m SqlPool
makePostgresPoolWithIamAuth :: PostgresConnectionConf -> m SqlPool
makePostgresPoolWithIamAuth conf :: PostgresConnectionConf
conf@PostgresConnectionConf {Int
String
PostgresStatementTimeout
PostgresPassword
pccStatementTimeout :: PostgresStatementTimeout
pccPoolSize :: Int
pccDatabase :: String
pccPassword :: PostgresPassword
pccUser :: String
pccPort :: Int
pccHost :: String
pccStatementTimeout :: PostgresConnectionConf -> PostgresStatementTimeout
pccPoolSize :: PostgresConnectionConf -> Int
pccDatabase :: PostgresConnectionConf -> String
pccPassword :: PostgresConnectionConf -> PostgresPassword
pccUser :: PostgresConnectionConf -> String
pccPort :: PostgresConnectionConf -> Int
pccHost :: PostgresConnectionConf -> String
..} = do
  IORef AuroraIamToken
tokenIORef <- PostgresConnectionConf -> m (IORef AuroraIamToken)
forall (m :: * -> *).
(MonadUnliftIO m, MonadLogger m) =>
PostgresConnectionConf -> m (IORef AuroraIamToken)
spawnIamTokenRefreshThread PostgresConnectionConf
conf
  (LogFunc -> IO SqlBackend) -> Int -> m SqlPool
forall backend (m :: * -> *).
(MonadLoggerIO m, MonadUnliftIO m,
 BackendCompatible SqlBackend backend) =>
(LogFunc -> IO backend) -> Int -> m (Pool backend)
createSqlPool (IORef AuroraIamToken -> LogFunc -> IO SqlBackend
mkConn IORef AuroraIamToken
tokenIORef) Int
pccPoolSize
 where
  mkConn :: IORef AuroraIamToken -> LogFunc -> IO SqlBackend
mkConn IORef AuroraIamToken
tokenIORef LogFunc
logFunc = do
    AuroraIamToken
token <- IORef AuroraIamToken -> IO AuroraIamToken
forall (m :: * -> *) a. MonadIO m => IORef a -> m a
readIORef IORef AuroraIamToken
tokenIORef
    let connStr :: ByteString
connStr = PostgresConnectionConf -> String -> ByteString
postgresConnectionString PostgresConnectionConf
conf (Text -> String
unpack (Text -> String) -> Text -> String
forall a b. (a -> b) -> a -> b
$ AuroraIamToken -> Text
aitToken AuroraIamToken
token)
    Connection
conn <- ByteString -> IO Connection
connectPostgreSQL ByteString
connStr
    PostgresConnectionConf -> Connection -> IO ()
forall (m :: * -> *).
MonadIO m =>
PostgresConnectionConf -> Connection -> m ()
setTimeout PostgresConnectionConf
conf Connection
conn
    LogFunc -> Connection -> IO SqlBackend
openSimpleConn LogFunc
logFunc Connection
conn

postgresConnectionString :: PostgresConnectionConf -> String -> ByteString
postgresConnectionString :: PostgresConnectionConf -> String -> ByteString
postgresConnectionString PostgresConnectionConf {Int
String
PostgresStatementTimeout
PostgresPassword
pccStatementTimeout :: PostgresStatementTimeout
pccPoolSize :: Int
pccDatabase :: String
pccPassword :: PostgresPassword
pccUser :: String
pccPort :: Int
pccHost :: String
pccStatementTimeout :: PostgresConnectionConf -> PostgresStatementTimeout
pccPoolSize :: PostgresConnectionConf -> Int
pccDatabase :: PostgresConnectionConf -> String
pccPassword :: PostgresConnectionConf -> PostgresPassword
pccUser :: PostgresConnectionConf -> String
pccPort :: PostgresConnectionConf -> Int
pccHost :: PostgresConnectionConf -> String
..} String
password =
  String -> ByteString
BS8.pack (String -> ByteString) -> String -> ByteString
forall a b. (a -> b) -> a -> b
$ [String] -> String
unwords
    [ String
"host=" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
pccHost
    , String
"port=" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
pccPort
    , String
"user=" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
pccUser
    , String
"password=" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
password
    , String
"dbname=" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
pccDatabase
    ]