module Postgres.Connection (connection, connectionIO, SingleOrPool (..), Connection (..)) where

import qualified Control.Exception.Safe as Exception
import qualified Data.Acquire
import qualified Data.Pool
import qualified Data.Text.Encoding
import Database.PostgreSQL.Typed
  ( PGConnection,
    PGDatabase (PGDatabase),
    pgConnect,
    pgDBAddr,
    pgDBName,
    pgDBUser,
    pgDisconnect,
  )
import qualified Log.SqlQuery as SqlQuery
import qualified Network.Socket as Socket
import qualified Postgres.Settings as Settings
import qualified Postgres.Time as Time
import qualified System.Exit
import qualified Prelude

-- | A connection to Postgres. You need this for making Postgres queries.
data Connection = Connection
  { Connection -> DoAnythingHandler
doAnything :: Platform.DoAnythingHandler,
    Connection -> SingleOrPool PGConnection
singleOrPool :: SingleOrPool PGConnection,
    Connection -> Details
connDetails :: SqlQuery.Details,
    Connection -> Interval
timeout :: Time.Interval
  }

-- | A database connection type.
--   Defining our own type makes it easier to change it in the future, without
--   having to fix compilation errors all over the codebase.
data SingleOrPool c
  = -- | By default a connection pool is passed around. It will:
    --   - Create new connections in the pool up to a certain limit.
    --   - Remove connections from the pool after a query in a connection errored.
    Pool (Data.Pool.Pool c)
  | -- | A single connection is only used in the context of a transaction, where
    --   we need to insure several SQL statements happen on the same connection.
    Single c

connectionIO :: Settings.Settings -> Prelude.IO Connection
connectionIO :: Settings -> IO Connection
connectionIO Settings
settings = do
  let database :: PGDatabase
database = Settings -> PGDatabase
Settings.toPGDatabase Settings
settings
  let stripes :: Int
stripes =
        PgPoolStripes -> Int
Settings.unPgPoolStripes (PoolSettings -> PgPoolStripes
Settings.pgPoolStripes (Settings -> PoolSettings
Settings.pgPool Settings
settings))
          Int -> (Int -> Int) -> Int
forall a b. a -> (a -> b) -> b
|> Int -> Int
forall a b. (Integral a, Num b) => a -> b
Prelude.fromIntegral
  let maxIdleTime :: NominalDiffTime
maxIdleTime = PgPoolMaxIdleTime -> NominalDiffTime
Settings.unPgPoolMaxIdleTime (PoolSettings -> PgPoolMaxIdleTime
Settings.pgPoolMaxIdleTime (Settings -> PoolSettings
Settings.pgPool Settings
settings))
  let size :: Int
size =
        PgPoolSize -> Int
Settings.unPgPoolSize (PoolSettings -> PgPoolSize
Settings.pgPoolSize (Settings -> PoolSettings
Settings.pgPool Settings
settings))
          Int -> (Int -> Int) -> Int
forall a b. a -> (a -> b) -> b
|> Int -> Int
forall a b. (Integral a, Num b) => a -> b
Prelude.fromIntegral
  DoAnythingHandler
doAnything <- IO DoAnythingHandler
Platform.doAnythingHandler
  SingleOrPool PGConnection
pool <-
    (Pool PGConnection -> SingleOrPool PGConnection)
-> IO (Pool PGConnection) -> IO (SingleOrPool PGConnection)
forall (m :: * -> *) a value.
Functor m =>
(a -> value) -> m a -> m value
map Pool PGConnection -> SingleOrPool PGConnection
forall c. Pool c -> SingleOrPool c
Pool
      (IO (Pool PGConnection) -> IO (SingleOrPool PGConnection))
-> IO (Pool PGConnection) -> IO (SingleOrPool PGConnection)
forall a b. (a -> b) -> a -> b
<| IO PGConnection
-> (PGConnection -> IO ())
-> Int
-> NominalDiffTime
-> Int
-> IO (Pool PGConnection)
forall a.
IO a
-> (a -> IO ()) -> Int -> NominalDiffTime -> Int -> IO (Pool a)
Data.Pool.createPool
        (PGDatabase -> IO PGConnection
pgConnect PGDatabase
database IO PGConnection
-> (IOException -> IO PGConnection) -> IO PGConnection
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> (e -> m a) -> m a
`Exception.catch` Text -> IOException -> IO PGConnection
forall a. Text -> IOException -> IO a
handleError (PGDatabase -> Text
toConnectionString PGDatabase
database))
        PGConnection -> IO ()
pgDisconnect
        Int
stripes
        NominalDiffTime
maxIdleTime
        Int
size
  Connection -> IO Connection
forall (f :: * -> *) a. Applicative f => a -> f a
Prelude.pure
    ( DoAnythingHandler
-> SingleOrPool PGConnection -> Details -> Interval -> Connection
Connection
        DoAnythingHandler
doAnything
        SingleOrPool PGConnection
pool
        (PGDatabase -> Details
connectionDetails PGDatabase
database)
        (Settings -> Interval
Settings.pgQueryTimeout Settings
settings)
    )

-- | Create a 'Connection'.
connection :: Settings.Settings -> Data.Acquire.Acquire Connection
connection :: Settings -> Acquire Connection
connection Settings
settings =
  IO Connection -> (Connection -> IO ()) -> Acquire Connection
forall a. IO a -> (a -> IO ()) -> Acquire a
Data.Acquire.mkAcquire (Settings -> IO Connection
connectionIO Settings
settings) Connection -> IO ()
release
  where
    release :: Connection -> IO ()
release Connection {SingleOrPool PGConnection
singleOrPool :: SingleOrPool PGConnection
singleOrPool :: Connection -> SingleOrPool PGConnection
singleOrPool} =
      case SingleOrPool PGConnection
singleOrPool of
        Pool Pool PGConnection
pool -> Pool PGConnection -> IO ()
forall a. Pool a -> IO ()
Data.Pool.destroyAllResources Pool PGConnection
pool
        Single PGConnection
single -> PGConnection -> IO ()
pgDisconnect PGConnection
single

handleError :: Text -> Exception.IOException -> Prelude.IO a
handleError :: Text -> IOException -> IO a
handleError Text
connectionString IOException
err = do
  String -> IO ()
Prelude.putStrLn String
"I couldn't connect to Postgres"
  String -> IO ()
Prelude.putStrLn String
""
  String -> IO ()
Prelude.putStrLn String
"Is the database running?"
  String -> IO ()
Prelude.putStrLn (String
"I tried to connect to: " String -> String -> String
forall appendable.
Semigroup appendable =>
appendable -> appendable -> appendable
++ Text -> String
Text.toList Text
connectionString)
  String -> IO a
forall a. String -> IO a
System.Exit.die (IOException -> String
forall e. Exception e => e -> String
Exception.displayException IOException
err)

toConnectionString :: PGDatabase -> Text
toConnectionString :: PGDatabase -> Text
toConnectionString PGDatabase {ByteString
pgDBUser :: ByteString
pgDBUser :: PGDatabase -> ByteString
pgDBUser, Either (String, String) SockAddr
pgDBAddr :: Either (String, String) SockAddr
pgDBAddr :: PGDatabase -> Either (String, String) SockAddr
pgDBAddr, ByteString
pgDBName :: ByteString
pgDBName :: PGDatabase -> ByteString
pgDBName} =
  Text -> List Text -> Text
Text.join
    Text
""
    [ ByteString -> Text
Data.Text.Encoding.decodeUtf8 ByteString
pgDBUser,
      Text
":*****@",
      case Either (String, String) SockAddr
pgDBAddr of
        Prelude.Right SockAddr
sockAddr ->
          String -> Text
Text.fromList (SockAddr -> String
forall a. Show a => a -> String
Prelude.show SockAddr
sockAddr)
        Prelude.Left (String
hostName, String
serviceName) ->
          String -> Text
Text.fromList String
hostName
            Text -> Text -> Text
forall appendable.
Semigroup appendable =>
appendable -> appendable -> appendable
++ Text
":"
            Text -> Text -> Text
forall appendable.
Semigroup appendable =>
appendable -> appendable -> appendable
++ String -> Text
Text.fromList String
serviceName,
      Text
"/",
      ByteString -> Text
Data.Text.Encoding.decodeUtf8 ByteString
pgDBName
    ]

connectionDetails :: PGDatabase -> SqlQuery.Details
connectionDetails :: PGDatabase -> Details
connectionDetails PGDatabase
db =
  case PGDatabase -> Either (String, String) SockAddr
pgDBAddr PGDatabase
db of
    Prelude.Left (String
hostName, String
serviceName) ->
      Details
SqlQuery.emptyDetails
        { databaseType :: Maybe Text
SqlQuery.databaseType = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
SqlQuery.postgresql,
          host :: Maybe Text
SqlQuery.host = Text -> Maybe Text
forall a. a -> Maybe a
Just (String -> Text
Text.fromList String
hostName),
          port :: Maybe Int
SqlQuery.port = Text -> Maybe Int
Text.toInt (String -> Text
Text.fromList String
serviceName),
          database :: Maybe Text
SqlQuery.database = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
databaseName
        }
    Prelude.Right (Socket.SockAddrInet PortNumber
portNum HostAddress
hostAddr) ->
      Details
SqlQuery.emptyDetails
        { databaseType :: Maybe Text
SqlQuery.databaseType = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
SqlQuery.postgresql,
          host :: Maybe Text
SqlQuery.host = Text -> Maybe Text
forall a. a -> Maybe a
Just (String -> Text
Text.fromList (HostAddress -> String
forall a. Show a => a -> String
Prelude.show HostAddress
hostAddr)),
          port :: Maybe Int
SqlQuery.port = Int -> Maybe Int
forall a. a -> Maybe a
Just (PortNumber -> Int
forall a b. (Integral a, Num b) => a -> b
Prelude.fromIntegral PortNumber
portNum),
          database :: Maybe Text
SqlQuery.database = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
databaseName
        }
    Prelude.Right (Socket.SockAddrInet6 PortNumber
portNum HostAddress
_flowInfo HostAddress6
hostAddr HostAddress
_scopeId) ->
      Details
SqlQuery.emptyDetails
        { databaseType :: Maybe Text
SqlQuery.databaseType = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
SqlQuery.postgresql,
          host :: Maybe Text
SqlQuery.host = Text -> Maybe Text
forall a. a -> Maybe a
Just (String -> Text
Text.fromList (HostAddress6 -> String
forall a. Show a => a -> String
Prelude.show HostAddress6
hostAddr)),
          port :: Maybe Int
SqlQuery.port = Int -> Maybe Int
forall a. a -> Maybe a
Just (PortNumber -> Int
forall a b. (Integral a, Num b) => a -> b
Prelude.fromIntegral PortNumber
portNum),
          database :: Maybe Text
SqlQuery.database = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
databaseName
        }
    Prelude.Right (Socket.SockAddrUnix String
sockPath) ->
      Details
SqlQuery.emptyDetails
        { databaseType :: Maybe Text
SqlQuery.databaseType = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
SqlQuery.postgresql,
          host :: Maybe Text
SqlQuery.host = Text -> Maybe Text
forall a. a -> Maybe a
Just (String -> Text
Text.fromList String
sockPath),
          port :: Maybe Int
SqlQuery.port = Maybe Int
forall a. Maybe a
Nothing,
          database :: Maybe Text
SqlQuery.database = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
databaseName
        }
  where
    databaseName :: Text
databaseName = PGDatabase -> ByteString
pgDBName PGDatabase
db ByteString -> (ByteString -> Text) -> Text
forall a b. a -> (a -> b) -> b
|> ByteString -> Text
Data.Text.Encoding.decodeUtf8