{-# LANGUAGE CPP #-}
module Database.PostgreSQL.PQTypes.Internal.Connection (
    Connection(..)
  , ConnectionData(..)
  , withConnectionData
  , ConnectionStats(..)
  , ConnectionSettings(..)
  , defaultConnectionSettings
  , ConnectionSourceM(..)
  , ConnectionSource(..)
  , simpleSource
  , poolSource
  , connect
  , disconnect
  ) where

import Control.Arrow (first)
import Control.Concurrent
import Control.Monad
import Control.Monad.Base
import Control.Monad.Catch
import Data.Function
import Data.IORef
import Data.Kind
import Data.Pool
import Data.Time.Clock
import Foreign.C.String
import Foreign.Ptr
import GHC.Conc (closeFdWith)
import qualified Control.Exception as E
import qualified Data.ByteString as BS
import qualified Data.Foldable as F
import qualified Data.Set as S
import qualified Data.Text as T
import qualified Data.Text.Encoding as T

import Database.PostgreSQL.PQTypes.Internal.C.Interface
import Database.PostgreSQL.PQTypes.Internal.C.Types
import Database.PostgreSQL.PQTypes.Internal.Composite
import Database.PostgreSQL.PQTypes.Internal.Error
import Database.PostgreSQL.PQTypes.Internal.Utils

data ConnectionSettings = ConnectionSettings
  { -- | Connection info string.
    ConnectionSettings -> Text
csConnInfo       :: !T.Text
    -- | Client-side encoding. If set to 'Nothing', database encoding is used.
  , ConnectionSettings -> Maybe Text
csClientEncoding :: !(Maybe T.Text)
    -- | A list of composite types to register. In order to be able to
    -- (de)serialize specific composite types, you need to register them.
  , ConnectionSettings -> [Text]
csComposites     :: ![T.Text]
  } deriving (ConnectionSettings -> ConnectionSettings -> Bool
(ConnectionSettings -> ConnectionSettings -> Bool)
-> (ConnectionSettings -> ConnectionSettings -> Bool)
-> Eq ConnectionSettings
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ConnectionSettings -> ConnectionSettings -> Bool
$c/= :: ConnectionSettings -> ConnectionSettings -> Bool
== :: ConnectionSettings -> ConnectionSettings -> Bool
$c== :: ConnectionSettings -> ConnectionSettings -> Bool
Eq, Eq ConnectionSettings
Eq ConnectionSettings
-> (ConnectionSettings -> ConnectionSettings -> Ordering)
-> (ConnectionSettings -> ConnectionSettings -> Bool)
-> (ConnectionSettings -> ConnectionSettings -> Bool)
-> (ConnectionSettings -> ConnectionSettings -> Bool)
-> (ConnectionSettings -> ConnectionSettings -> Bool)
-> (ConnectionSettings -> ConnectionSettings -> ConnectionSettings)
-> (ConnectionSettings -> ConnectionSettings -> ConnectionSettings)
-> Ord ConnectionSettings
ConnectionSettings -> ConnectionSettings -> Bool
ConnectionSettings -> ConnectionSettings -> Ordering
ConnectionSettings -> ConnectionSettings -> ConnectionSettings
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ConnectionSettings -> ConnectionSettings -> ConnectionSettings
$cmin :: ConnectionSettings -> ConnectionSettings -> ConnectionSettings
max :: ConnectionSettings -> ConnectionSettings -> ConnectionSettings
$cmax :: ConnectionSettings -> ConnectionSettings -> ConnectionSettings
>= :: ConnectionSettings -> ConnectionSettings -> Bool
$c>= :: ConnectionSettings -> ConnectionSettings -> Bool
> :: ConnectionSettings -> ConnectionSettings -> Bool
$c> :: ConnectionSettings -> ConnectionSettings -> Bool
<= :: ConnectionSettings -> ConnectionSettings -> Bool
$c<= :: ConnectionSettings -> ConnectionSettings -> Bool
< :: ConnectionSettings -> ConnectionSettings -> Bool
$c< :: ConnectionSettings -> ConnectionSettings -> Bool
compare :: ConnectionSettings -> ConnectionSettings -> Ordering
$ccompare :: ConnectionSettings -> ConnectionSettings -> Ordering
$cp1Ord :: Eq ConnectionSettings
Ord, Int -> ConnectionSettings -> ShowS
[ConnectionSettings] -> ShowS
ConnectionSettings -> String
(Int -> ConnectionSettings -> ShowS)
-> (ConnectionSettings -> String)
-> ([ConnectionSettings] -> ShowS)
-> Show ConnectionSettings
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ConnectionSettings] -> ShowS
$cshowList :: [ConnectionSettings] -> ShowS
show :: ConnectionSettings -> String
$cshow :: ConnectionSettings -> String
showsPrec :: Int -> ConnectionSettings -> ShowS
$cshowsPrec :: Int -> ConnectionSettings -> ShowS
Show)

-- | Default connection settings. Note that all strings sent to PostgreSQL by
-- the library are encoded as UTF-8, so don't alter client encoding unless you
-- know what you're doing.
defaultConnectionSettings :: ConnectionSettings
defaultConnectionSettings :: ConnectionSettings
defaultConnectionSettings =
  ConnectionSettings :: Text -> Maybe Text -> [Text] -> ConnectionSettings
ConnectionSettings
  { csConnInfo :: Text
csConnInfo       = Text
T.empty
  , csClientEncoding :: Maybe Text
csClientEncoding = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"UTF-8"
  , csComposites :: [Text]
csComposites     = []
  }

----------------------------------------

-- | Simple connection statistics.
data ConnectionStats = ConnectionStats
  { -- | Number of queries executed so far.
    ConnectionStats -> Int
statsQueries :: !Int
    -- | Number of rows fetched from the database.
  , ConnectionStats -> Int
statsRows    :: !Int
    -- | Number of values fetched from the database.
  , ConnectionStats -> Int
statsValues  :: !Int
    -- | Number of parameters sent to the database.
  , ConnectionStats -> Int
statsParams  :: !Int
  } deriving (ConnectionStats -> ConnectionStats -> Bool
(ConnectionStats -> ConnectionStats -> Bool)
-> (ConnectionStats -> ConnectionStats -> Bool)
-> Eq ConnectionStats
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ConnectionStats -> ConnectionStats -> Bool
$c/= :: ConnectionStats -> ConnectionStats -> Bool
== :: ConnectionStats -> ConnectionStats -> Bool
$c== :: ConnectionStats -> ConnectionStats -> Bool
Eq, Eq ConnectionStats
Eq ConnectionStats
-> (ConnectionStats -> ConnectionStats -> Ordering)
-> (ConnectionStats -> ConnectionStats -> Bool)
-> (ConnectionStats -> ConnectionStats -> Bool)
-> (ConnectionStats -> ConnectionStats -> Bool)
-> (ConnectionStats -> ConnectionStats -> Bool)
-> (ConnectionStats -> ConnectionStats -> ConnectionStats)
-> (ConnectionStats -> ConnectionStats -> ConnectionStats)
-> Ord ConnectionStats
ConnectionStats -> ConnectionStats -> Bool
ConnectionStats -> ConnectionStats -> Ordering
ConnectionStats -> ConnectionStats -> ConnectionStats
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ConnectionStats -> ConnectionStats -> ConnectionStats
$cmin :: ConnectionStats -> ConnectionStats -> ConnectionStats
max :: ConnectionStats -> ConnectionStats -> ConnectionStats
$cmax :: ConnectionStats -> ConnectionStats -> ConnectionStats
>= :: ConnectionStats -> ConnectionStats -> Bool
$c>= :: ConnectionStats -> ConnectionStats -> Bool
> :: ConnectionStats -> ConnectionStats -> Bool
$c> :: ConnectionStats -> ConnectionStats -> Bool
<= :: ConnectionStats -> ConnectionStats -> Bool
$c<= :: ConnectionStats -> ConnectionStats -> Bool
< :: ConnectionStats -> ConnectionStats -> Bool
$c< :: ConnectionStats -> ConnectionStats -> Bool
compare :: ConnectionStats -> ConnectionStats -> Ordering
$ccompare :: ConnectionStats -> ConnectionStats -> Ordering
$cp1Ord :: Eq ConnectionStats
Ord, Int -> ConnectionStats -> ShowS
[ConnectionStats] -> ShowS
ConnectionStats -> String
(Int -> ConnectionStats -> ShowS)
-> (ConnectionStats -> String)
-> ([ConnectionStats] -> ShowS)
-> Show ConnectionStats
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ConnectionStats] -> ShowS
$cshowList :: [ConnectionStats] -> ShowS
show :: ConnectionStats -> String
$cshow :: ConnectionStats -> String
showsPrec :: Int -> ConnectionStats -> ShowS
$cshowsPrec :: Int -> ConnectionStats -> ShowS
Show)

-- | Initial connection statistics.
initialStats :: ConnectionStats
initialStats :: ConnectionStats
initialStats = ConnectionStats :: Int -> Int -> Int -> Int -> ConnectionStats
ConnectionStats {
  statsQueries :: Int
statsQueries = Int
0
, statsRows :: Int
statsRows    = Int
0
, statsValues :: Int
statsValues  = Int
0
, statsParams :: Int
statsParams  = Int
0
}

-- | Representation of a connection object.
--
-- /Note:/ PGconn is not managed with a ForeignPtr because finalizers are broken
-- and at program exit might run even though another thread is inside the
-- relevant withForeignPtr block, executing a safe FFI call (in this case
-- executing an SQL query).
--
-- See https://gitlab.haskell.org/ghc/ghc/-/issues/10975 for more info.
data ConnectionData = ConnectionData
  { ConnectionData -> Ptr PGconn
cdPtr      :: !(Ptr PGconn)
  -- ^ Pointer to connection object.
  , ConnectionData -> ConnectionStats
cdStats    :: !ConnectionStats
  -- ^ Statistics associated with the connection.
  , ConnectionData -> IORef (Set Text)
cdPreparedQueries :: !(IORef (S.Set T.Text))
  -- ^ A set of named prepared statements of the connection.
  }

-- | Wrapper for hiding representation of a connection object.
newtype Connection = Connection {
  Connection -> MVar (Maybe ConnectionData)
unConnection :: MVar (Maybe ConnectionData)
}

withConnectionData
  :: Connection
  -> String
  -> (ConnectionData -> IO (ConnectionData, r))
  -> IO r
withConnectionData :: Connection
-> String -> (ConnectionData -> IO (ConnectionData, r)) -> IO r
withConnectionData (Connection MVar (Maybe ConnectionData)
mvc) String
fname ConnectionData -> IO (ConnectionData, r)
f =
  MVar (Maybe ConnectionData)
-> (Maybe ConnectionData -> IO (Maybe ConnectionData, r)) -> IO r
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar (Maybe ConnectionData)
mvc ((Maybe ConnectionData -> IO (Maybe ConnectionData, r)) -> IO r)
-> (Maybe ConnectionData -> IO (Maybe ConnectionData, r)) -> IO r
forall a b. (a -> b) -> a -> b
$ \Maybe ConnectionData
mc -> case Maybe ConnectionData
mc of
    Maybe ConnectionData
Nothing -> String -> IO (Maybe ConnectionData, r)
forall a. String -> IO a
hpqTypesError (String -> IO (Maybe ConnectionData, r))
-> String -> IO (Maybe ConnectionData, r)
forall a b. (a -> b) -> a -> b
$ String
fname String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
": no connection"
    Just ConnectionData
cd -> (ConnectionData -> Maybe ConnectionData)
-> (ConnectionData, r) -> (Maybe ConnectionData, r)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first ConnectionData -> Maybe ConnectionData
forall a. a -> Maybe a
Just ((ConnectionData, r) -> (Maybe ConnectionData, r))
-> IO (ConnectionData, r) -> IO (Maybe ConnectionData, r)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ConnectionData -> IO (ConnectionData, r)
f ConnectionData
cd

-- | Database connection supplier.
newtype ConnectionSourceM m = ConnectionSourceM {
  ConnectionSourceM m -> forall r. (Connection -> m r) -> m r
withConnection :: forall r. (Connection -> m r) -> m r
}

-- | Wrapper for a polymorphic connection source.
newtype ConnectionSource (cs :: [(Type -> Type) -> Constraint]) = ConnectionSource {
  ConnectionSource cs
-> forall (m :: * -> *). MkConstraint m cs => ConnectionSourceM m
unConnectionSource :: forall m. MkConstraint m cs => ConnectionSourceM m
}

-- | Default connection supplier. It establishes new
-- database connection each time 'withConnection' is called.
simpleSource
  :: ConnectionSettings
  -> ConnectionSource [MonadBase IO, MonadMask]
simpleSource :: ConnectionSettings -> ConnectionSource '[MonadBase IO, MonadMask]
simpleSource ConnectionSettings
cs = (forall (m :: * -> *).
 MkConstraint m '[MonadBase IO, MonadMask] =>
 ConnectionSourceM m)
-> ConnectionSource '[MonadBase IO, MonadMask]
forall (cs :: [(* -> *) -> Constraint]).
(forall (m :: * -> *). MkConstraint m cs => ConnectionSourceM m)
-> ConnectionSource cs
ConnectionSource ((forall (m :: * -> *).
  MkConstraint m '[MonadBase IO, MonadMask] =>
  ConnectionSourceM m)
 -> ConnectionSource '[MonadBase IO, MonadMask])
-> (forall (m :: * -> *).
    MkConstraint m '[MonadBase IO, MonadMask] =>
    ConnectionSourceM m)
-> ConnectionSource '[MonadBase IO, MonadMask]
forall a b. (a -> b) -> a -> b
$ ConnectionSourceM :: forall (m :: * -> *).
(forall r. (Connection -> m r) -> m r) -> ConnectionSourceM m
ConnectionSourceM {
  withConnection :: forall r. (Connection -> m r) -> m r
withConnection = m Connection -> (Connection -> m ()) -> (Connection -> m r) -> m r
forall (m :: * -> *) a c b.
MonadMask m =>
m a -> (a -> m c) -> (a -> m b) -> m b
bracket (IO Connection -> m Connection
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase (IO Connection -> m Connection) -> IO Connection -> m Connection
forall a b. (a -> b) -> a -> b
$ ConnectionSettings -> IO Connection
connect ConnectionSettings
cs) (IO () -> m ()
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase (IO () -> m ()) -> (Connection -> IO ()) -> Connection -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Connection -> IO ()
disconnect)
}

-- | Pooled source. It uses striped pool from resource-pool
-- package to cache established connections and reuse them.
poolSource
  :: ConnectionSettings
  -> Int -- ^ Stripe count. The number of distinct sub-pools
  -- to maintain. The smallest acceptable value is 1.
  -> NominalDiffTime -- ^ Amount of time for which an unused database
  -- connection is kept open. The smallest acceptable value is 0.5
  -- seconds.
  --
  -- The elapsed time before closing database connection may be
  -- a little longer than requested, as the reaper thread wakes
  -- at 1-second intervals.
  -> Int -- ^ Maximum number of database connections to keep open
  -- per stripe. The smallest acceptable value is 1.
  --
  -- Requests for database connections will block if this limit is
  -- reached on a single stripe, even if other stripes have idle
  -- connections available.
  -> IO (ConnectionSource [MonadBase IO, MonadMask])
poolSource :: ConnectionSettings
-> Int
-> NominalDiffTime
-> Int
-> IO (ConnectionSource '[MonadBase IO, MonadMask])
poolSource ConnectionSettings
cs Int
numStripes NominalDiffTime
idleTime Int
maxResources = do
  Pool Connection
pool <- IO Connection
-> (Connection -> IO ())
-> Int
-> NominalDiffTime
-> Int
-> IO (Pool Connection)
forall a.
IO a
-> (a -> IO ()) -> Int -> NominalDiffTime -> Int -> IO (Pool a)
createPool (ConnectionSettings -> IO Connection
connect ConnectionSettings
cs) Connection -> IO ()
disconnect Int
numStripes NominalDiffTime
idleTime Int
maxResources
  ConnectionSource '[MonadBase IO, MonadMask]
-> IO (ConnectionSource '[MonadBase IO, MonadMask])
forall (m :: * -> *) a. Monad m => a -> m a
return (ConnectionSource '[MonadBase IO, MonadMask]
 -> IO (ConnectionSource '[MonadBase IO, MonadMask]))
-> ConnectionSource '[MonadBase IO, MonadMask]
-> IO (ConnectionSource '[MonadBase IO, MonadMask])
forall a b. (a -> b) -> a -> b
$ (forall (m :: * -> *).
 MkConstraint m '[MonadBase IO, MonadMask] =>
 ConnectionSourceM m)
-> ConnectionSource '[MonadBase IO, MonadMask]
forall (cs :: [(* -> *) -> Constraint]).
(forall (m :: * -> *). MkConstraint m cs => ConnectionSourceM m)
-> ConnectionSource cs
ConnectionSource ((forall (m :: * -> *).
  MkConstraint m '[MonadBase IO, MonadMask] =>
  ConnectionSourceM m)
 -> ConnectionSource '[MonadBase IO, MonadMask])
-> (forall (m :: * -> *).
    MkConstraint m '[MonadBase IO, MonadMask] =>
    ConnectionSourceM m)
-> ConnectionSource '[MonadBase IO, MonadMask]
forall a b. (a -> b) -> a -> b
$ ConnectionSourceM :: forall (m :: * -> *).
(forall r. (Connection -> m r) -> m r) -> ConnectionSourceM m
ConnectionSourceM {
    withConnection :: forall r. (Connection -> m r) -> m r
withConnection = Pool Connection -> (Connection -> m r) -> m r
forall (f :: * -> *) t b.
(MonadMask f, MonadBase IO f) =>
Pool t -> (t -> f b) -> f b
doWithConnection Pool Connection
pool ((Connection -> m r) -> m r)
-> ((Connection -> m r) -> Connection -> m r)
-> (Connection -> m r)
-> m r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Connection -> m Connection
forall (m :: * -> *). MonadBase IO m => Connection -> m Connection
clearStats (Connection -> m Connection)
-> (Connection -> m r) -> Connection -> m r
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=>)
  }
  where
    doWithConnection :: Pool t -> (t -> f b) -> f b
doWithConnection Pool t
pool t -> f b
m = (b, ()) -> b
forall a b. (a, b) -> a
fst ((b, ()) -> b) -> f (b, ()) -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (t, LocalPool t)
-> ((t, LocalPool t) -> ExitCase b -> f ())
-> ((t, LocalPool t) -> f b)
-> f (b, ())
forall (m :: * -> *) a b c.
MonadMask m =>
m a -> (a -> ExitCase b -> m c) -> (a -> m b) -> m (b, c)
generalBracket
      (IO (t, LocalPool t) -> f (t, LocalPool t)
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase (IO (t, LocalPool t) -> f (t, LocalPool t))
-> IO (t, LocalPool t) -> f (t, LocalPool t)
forall a b. (a -> b) -> a -> b
$ Pool t -> IO (t, LocalPool t)
forall a. Pool a -> IO (a, LocalPool a)
takeResource Pool t
pool)
      (\(t
resource, LocalPool t
local) -> \case
          ExitCaseSuccess b
_ -> IO () -> f ()
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase (IO () -> f ()) -> IO () -> f ()
forall a b. (a -> b) -> a -> b
$ LocalPool t -> t -> IO ()
forall a. LocalPool a -> a -> IO ()
putResource LocalPool t
local t
resource
          ExitCase b
_                 -> IO () -> f ()
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase (IO () -> f ()) -> IO () -> f ()
forall a b. (a -> b) -> a -> b
$ Pool t -> LocalPool t -> t -> IO ()
forall a. Pool a -> LocalPool a -> a -> IO ()
destroyResource Pool t
pool LocalPool t
local t
resource
      )
      (\(t
resource, LocalPool t
_) -> t -> f b
m t
resource)

    clearStats :: Connection -> m Connection
clearStats conn :: Connection
conn@(Connection MVar (Maybe ConnectionData)
mv) = do
      IO () -> m ()
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase (IO () -> m ())
-> ((Maybe ConnectionData -> IO (Maybe ConnectionData)) -> IO ())
-> (Maybe ConnectionData -> IO (Maybe ConnectionData))
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MVar (Maybe ConnectionData)
-> (Maybe ConnectionData -> IO (Maybe ConnectionData)) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar (Maybe ConnectionData)
mv ((Maybe ConnectionData -> IO (Maybe ConnectionData)) -> m ())
-> (Maybe ConnectionData -> IO (Maybe ConnectionData)) -> m ()
forall a b. (a -> b) -> a -> b
$ \Maybe ConnectionData
mconn ->
        Maybe ConnectionData -> IO (Maybe ConnectionData)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ConnectionData -> IO (Maybe ConnectionData))
-> Maybe ConnectionData -> IO (Maybe ConnectionData)
forall a b. (a -> b) -> a -> b
$ (\ConnectionData
cd -> ConnectionData
cd { cdStats :: ConnectionStats
cdStats = ConnectionStats
initialStats }) (ConnectionData -> ConnectionData)
-> Maybe ConnectionData -> Maybe ConnectionData
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe ConnectionData
mconn
      Connection -> m Connection
forall (m :: * -> *) a. Monad m => a -> m a
return Connection
conn

----------------------------------------

-- | Low-level function for connecting to the database. Useful if one wants to
-- implement custom connection source.
--
-- /Warning:/ the 'Connection' needs to be explicitly destroyed with
-- 'disconnect', otherwise there will be a resource leak.
connect :: ConnectionSettings -> IO Connection
connect :: ConnectionSettings -> IO Connection
connect ConnectionSettings{[Text]
Maybe Text
Text
csComposites :: [Text]
csClientEncoding :: Maybe Text
csConnInfo :: Text
csComposites :: ConnectionSettings -> [Text]
csClientEncoding :: ConnectionSettings -> Maybe Text
csConnInfo :: ConnectionSettings -> Text
..} = ((forall a. IO a -> IO a) -> IO Connection) -> IO Connection
forall (m :: * -> *) b.
MonadMask m =>
((forall a. m a -> m a) -> m b) -> m b
mask (((forall a. IO a -> IO a) -> IO Connection) -> IO Connection)
-> ((forall a. IO a -> IO a) -> IO Connection) -> IO Connection
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> do
  Ptr PGconn
conn <- ByteString -> (CString -> IO (Ptr PGconn)) -> IO (Ptr PGconn)
forall a. ByteString -> (CString -> IO a) -> IO a
BS.useAsCString (Text -> ByteString
T.encodeUtf8 Text
csConnInfo) ((forall a. IO a -> IO a) -> CString -> IO (Ptr PGconn)
openConnection forall a. IO a -> IO a
unmask)
  (IO Connection -> IO () -> IO Connection
forall (m :: * -> *) a b. MonadCatch m => m a -> m b -> m a
`onException` Ptr PGconn -> IO ()
c_PQfinish Ptr PGconn
conn) (IO Connection -> IO Connection)
-> (IO Connection -> IO Connection)
-> IO Connection
-> IO Connection
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO Connection -> IO Connection
forall a. IO a -> IO a
unmask (IO Connection -> IO Connection) -> IO Connection -> IO Connection
forall a b. (a -> b) -> a -> b
$ do
    ConnStatusType
status <- Ptr PGconn -> IO ConnStatusType
c_PQstatus Ptr PGconn
conn
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ConnStatusType
status ConnStatusType -> ConnStatusType -> Bool
forall a. Eq a => a -> a -> Bool
/= ConnStatusType
c_CONNECTION_OK) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
      Ptr PGconn -> String -> IO ()
forall a. Ptr PGconn -> String -> IO a
throwLibPQError Ptr PGconn
conn String
fname
    Maybe Text -> (Text -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
F.forM_ Maybe Text
csClientEncoding ((Text -> IO ()) -> IO ()) -> (Text -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Text
enc -> do
      CInt
res <- ByteString -> (CString -> IO CInt) -> IO CInt
forall a. ByteString -> (CString -> IO a) -> IO a
BS.useAsCString (Text -> ByteString
T.encodeUtf8 Text
enc) (Ptr PGconn -> CString -> IO CInt
c_PQsetClientEncoding Ptr PGconn
conn)
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (CInt
res CInt -> CInt -> Bool
forall a. Eq a => a -> a -> Bool
== -CInt
1) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
        Ptr PGconn -> String -> IO ()
forall a. Ptr PGconn -> String -> IO a
throwLibPQError Ptr PGconn
conn String
fname
    Ptr PGconn -> IO ()
c_PQinitTypes Ptr PGconn
conn
    Ptr PGconn -> [Text] -> IO ()
registerComposites Ptr PGconn
conn [Text]
csComposites
    IORef (Set Text)
preparedQueries <- Set Text -> IO (IORef (Set Text))
forall a. a -> IO (IORef a)
newIORef Set Text
forall a. Set a
S.empty
    (MVar (Maybe ConnectionData) -> Connection)
-> IO (MVar (Maybe ConnectionData)) -> IO Connection
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap MVar (Maybe ConnectionData) -> Connection
Connection (IO (MVar (Maybe ConnectionData)) -> IO Connection)
-> (Maybe ConnectionData -> IO (MVar (Maybe ConnectionData)))
-> Maybe ConnectionData
-> IO Connection
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe ConnectionData -> IO (MVar (Maybe ConnectionData))
forall a. a -> IO (MVar a)
newMVar (Maybe ConnectionData -> IO Connection)
-> Maybe ConnectionData -> IO Connection
forall a b. (a -> b) -> a -> b
$ ConnectionData -> Maybe ConnectionData
forall a. a -> Maybe a
Just ConnectionData :: Ptr PGconn -> ConnectionStats -> IORef (Set Text) -> ConnectionData
ConnectionData
      { cdPtr :: Ptr PGconn
cdPtr = Ptr PGconn
conn
      , cdStats :: ConnectionStats
cdStats = ConnectionStats
initialStats
      , cdPreparedQueries :: IORef (Set Text)
cdPreparedQueries = IORef (Set Text)
preparedQueries
      }
  where
    fname :: String
fname = String
"connect"

    openConnection :: (forall r. IO r -> IO r) -> CString -> IO (Ptr PGconn)
    openConnection :: (forall a. IO a -> IO a) -> CString -> IO (Ptr PGconn)
openConnection forall a. IO a -> IO a
unmask CString
conninfo = do
      -- We want to use non-blocking C functions to be able to observe incoming
      -- asynchronous exceptions, hence we don't use PQconnectdb here.
      Ptr PGconn
conn <- CString -> IO (Ptr PGconn)
c_PQconnectStart CString
conninfo
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ptr PGconn
conn Ptr PGconn -> Ptr PGconn -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr PGconn
forall a. Ptr a
nullPtr) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
        String -> IO ()
forall a. String -> IO a
throwError String
"PQconnectStart returned a null pointer"
      Fd
fd <- Ptr PGconn -> IO Fd
getFd Ptr PGconn
conn
      (IO (Ptr PGconn) -> IO () -> IO (Ptr PGconn)
forall (m :: * -> *) a b. MonadCatch m => m a -> m b -> m a
`onException` Ptr PGconn -> IO ()
c_PQfinish Ptr PGconn
conn) (IO (Ptr PGconn) -> IO (Ptr PGconn))
-> (IO (Ptr PGconn) -> IO (Ptr PGconn))
-> IO (Ptr PGconn)
-> IO (Ptr PGconn)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Ptr PGconn) -> IO (Ptr PGconn)
forall a. IO a -> IO a
unmask (IO (Ptr PGconn) -> IO (Ptr PGconn))
-> IO (Ptr PGconn) -> IO (Ptr PGconn)
forall a b. (a -> b) -> a -> b
$ (IO (Ptr PGconn) -> IO (Ptr PGconn)) -> IO (Ptr PGconn)
forall a. (a -> a) -> a
fix ((IO (Ptr PGconn) -> IO (Ptr PGconn)) -> IO (Ptr PGconn))
-> (IO (Ptr PGconn) -> IO (Ptr PGconn)) -> IO (Ptr PGconn)
forall a b. (a -> b) -> a -> b
$ \IO (Ptr PGconn)
loop -> do
        PostgresPollingStatusType
ps <- Ptr PGconn -> IO PostgresPollingStatusType
c_PQconnectPoll Ptr PGconn
conn
        if | PostgresPollingStatusType
ps PostgresPollingStatusType -> PostgresPollingStatusType -> Bool
forall a. Eq a => a -> a -> Bool
== PostgresPollingStatusType
c_PGRES_POLLING_READING -> Fd -> IO ()
threadWaitRead  Fd
fd IO () -> IO (Ptr PGconn) -> IO (Ptr PGconn)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO (Ptr PGconn)
loop
           | PostgresPollingStatusType
ps PostgresPollingStatusType -> PostgresPollingStatusType -> Bool
forall a. Eq a => a -> a -> Bool
== PostgresPollingStatusType
c_PGRES_POLLING_WRITING -> Fd -> IO ()
threadWaitWrite Fd
fd IO () -> IO (Ptr PGconn) -> IO (Ptr PGconn)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO (Ptr PGconn)
loop
           | PostgresPollingStatusType
ps PostgresPollingStatusType -> PostgresPollingStatusType -> Bool
forall a. Eq a => a -> a -> Bool
== PostgresPollingStatusType
c_PGRES_POLLING_OK      -> Ptr PGconn -> IO (Ptr PGconn)
forall (m :: * -> *) a. Monad m => a -> m a
return Ptr PGconn
conn
           | Bool
otherwise                     -> String -> IO (Ptr PGconn)
forall a. String -> IO a
throwError String
"openConnection failed"
      where
        getFd :: Ptr PGconn -> IO Fd
getFd Ptr PGconn
conn = do
          Fd
fd <- Ptr PGconn -> IO Fd
c_PQsocket Ptr PGconn
conn
          Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Fd
fd Fd -> Fd -> Bool
forall a. Eq a => a -> a -> Bool
== -Fd
1) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
            String -> IO ()
forall a. String -> IO a
throwError String
"invalid file descriptor"
          Fd -> IO Fd
forall (m :: * -> *) a. Monad m => a -> m a
return Fd
fd

        throwError :: String -> IO a
        throwError :: String -> IO a
throwError = String -> IO a
forall a. String -> IO a
hpqTypesError (String -> IO a) -> ShowS -> String -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
fname String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
": " String -> ShowS
forall a. [a] -> [a] -> [a]
++)

-- | Low-level function for disconnecting from the database. Useful if one wants
-- to implement custom connection source.
disconnect :: Connection -> IO ()
disconnect :: Connection -> IO ()
disconnect (Connection MVar (Maybe ConnectionData)
mvconn) = MVar (Maybe ConnectionData)
-> (Maybe ConnectionData -> IO (Maybe ConnectionData)) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar (Maybe ConnectionData)
mvconn ((Maybe ConnectionData -> IO (Maybe ConnectionData)) -> IO ())
-> (Maybe ConnectionData -> IO (Maybe ConnectionData)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Maybe ConnectionData
mconn -> do
  case Maybe ConnectionData
mconn of
    Just ConnectionData
cd -> do
      let conn :: Ptr PGconn
conn = ConnectionData -> Ptr PGconn
cdPtr ConnectionData
cd
      -- This covers the case when a connection is closed while other Haskell
      -- threads are using GHC's IO manager to wait on the descriptor. This is
      -- commonly the case with asynchronous notifications, for example. Since
      -- libpq is responsible for opening and closing the file descriptor, GHC's
      -- IO manager needs to be informed that the file descriptor has been
      -- closed. The IO manager will then raise an exception in those threads.
      Ptr PGconn -> IO Fd
c_PQsocket Ptr PGconn
conn IO Fd -> (Fd -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        -1 -> Ptr PGconn -> IO ()
c_PQfinish Ptr PGconn
conn -- can happen if the connection is bad/lost
        Fd
fd -> (Fd -> IO ()) -> Fd -> IO ()
closeFdWith (\Fd
_ -> Ptr PGconn -> IO ()
c_PQfinish Ptr PGconn
conn) Fd
fd

    Maybe ConnectionData
Nothing -> HPQTypesError -> IO ()
forall e a. Exception e => e -> IO a
E.throwIO (String -> HPQTypesError
HPQTypesError String
"disconnect: no connection (shouldn't happen)")
  Maybe ConnectionData -> IO (Maybe ConnectionData)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ConnectionData
forall a. Maybe a
Nothing