{-# LANGUAGE RecordWildCards #-}

{-|
Module      : Database.PostgreSQL.Replicant
Description : A PostgreSQL streaming replication library
Copyright   : (c) James King, 2020, 2021
License     : BSD3
Maintainer  : [email protected]
Stability   : experimental
Portability : POSIX

Connect to a PostgreSQL server as a logical replication client and
receive changes.

The basic API is this:

@
  withLogicalStream defaultSettings $ \change -> do
    print change
    `catch` \err -> do
      show err
@

This is a low-level library meant to give the primitives necessary to
library authors to add streaming replication support.  The API here to
rather simplistic but should be hooked up to something like conduit to
provide better ergonomics.
-}

module Database.PostgreSQL.Replicant
    ( withLogicalStream
    , PgSettings (..)
    ) where

import Control.Concurrent
import Control.Exception
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as B
import qualified Data.Text as T
import qualified Data.Text.Read as T
import Database.PostgreSQL.LibPQ
import Network.Socket.KeepAlive
import System.Posix.Types

import Database.PostgreSQL.Replicant.Exception
import Database.PostgreSQL.Replicant.Protocol
import Database.PostgreSQL.Replicant.Message
import Database.PostgreSQL.Replicant.ReplicationSlot
import Database.PostgreSQL.Replicant.Util

data PgSettings
  = PgSettings
  { PgSettings -> String
pgUser        :: String
  , PgSettings -> String
pgDbName      :: String
  , PgSettings -> String
pgHost        :: String
  , PgSettings -> String
pgPort        :: String
  , PgSettings -> String
pgSlotName    :: String
  , PgSettings -> String
pgUpdateDelay :: String -- ^ Controls how frequently the
                            -- primaryKeepAlive thread updates
                            -- PostgresSQL in @ms@
  }
  deriving (PgSettings -> PgSettings -> Bool
(PgSettings -> PgSettings -> Bool)
-> (PgSettings -> PgSettings -> Bool) -> Eq PgSettings
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PgSettings -> PgSettings -> Bool
$c/= :: PgSettings -> PgSettings -> Bool
== :: PgSettings -> PgSettings -> Bool
$c== :: PgSettings -> PgSettings -> Bool
Eq, Int -> PgSettings -> ShowS
[PgSettings] -> ShowS
PgSettings -> String
(Int -> PgSettings -> ShowS)
-> (PgSettings -> String)
-> ([PgSettings] -> ShowS)
-> Show PgSettings
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PgSettings] -> ShowS
$cshowList :: [PgSettings] -> ShowS
show :: PgSettings -> String
$cshow :: PgSettings -> String
showsPrec :: Int -> PgSettings -> ShowS
$cshowsPrec :: Int -> PgSettings -> ShowS
Show)

pgConnectionString :: PgSettings -> ByteString
pgConnectionString :: PgSettings -> ByteString
pgConnectionString PgSettings {String
pgUpdateDelay :: String
pgSlotName :: String
pgPort :: String
pgHost :: String
pgDbName :: String
pgUser :: String
pgUpdateDelay :: PgSettings -> String
pgSlotName :: PgSettings -> String
pgPort :: PgSettings -> String
pgHost :: PgSettings -> String
pgDbName :: PgSettings -> String
pgUser :: PgSettings -> String
..} = ByteString -> [ByteString] -> ByteString
B.intercalate ByteString
" "
  [ ByteString
"user=" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> (String -> ByteString
B.pack String
pgUser)
  , ByteString
"dbname=" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> (String -> ByteString
B.pack String
pgDbName)
  , ByteString
"host=" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> (String -> ByteString
B.pack String
pgHost)
  , ByteString
"port=" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> (String -> ByteString
B.pack String
pgPort)
  , ByteString
"replication=database"
  ]

-- | Connect to a PostgreSQL database as a user with the replication
-- attribute and start receiving changes using the logical replication
-- protocol.  Logical replication happens at the query level so the
-- changes you get represent the set of queries in a transaction:
-- /insert/, /update/, and /delete/.
--
-- This function will create the replication slot, if it doesn't
-- exist, or reconnect to it otherwise and restart the stream from
-- where the replication slot left off.
--
-- This function can throw exceptions in @IO@ and shut-down the
-- socket in case of any error.
withLogicalStream :: PgSettings -> (Change -> IO a) -> IO ()
withLogicalStream :: PgSettings -> (Change -> IO a) -> IO ()
withLogicalStream PgSettings
settings Change -> IO a
cb = do
  Connection
conn <- ByteString -> IO Connection
connectStart (ByteString -> IO Connection) -> ByteString -> IO Connection
forall a b. (a -> b) -> a -> b
$ PgSettings -> ByteString
pgConnectionString PgSettings
settings
  Maybe Fd
mFd <- Connection -> IO (Maybe Fd)
socket Connection
conn
  Fd
sockFd <- ReplicantException -> Maybe Fd -> IO Fd
forall e a. Exception e => e -> Maybe a -> IO a
maybeThrow (String -> ReplicantException
ReplicantException String
"withLogicalStream: could not get socket fd") Maybe Fd
mFd
  PollingStatus
pollResult <- Connection -> Fd -> IO PollingStatus
pollConnectStart Connection
conn Fd
sockFd
  let updateFreq :: Int
updateFreq = PgSettings -> Int
getUpdateDelay PgSettings
settings
  case PollingStatus
pollResult of
    PollingStatus
PollingFailed -> ReplicantException -> IO ()
forall e a. Exception e => e -> IO a
throwIO (ReplicantException -> IO ()) -> ReplicantException -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> ReplicantException
ReplicantException String
"withLogicalStream: Unable to connect to the database"
    PollingStatus
PollingOk -> do
      Maybe IdentifySystem
maybeInfo <- Connection -> IO (Maybe IdentifySystem)
identifySystemSync Connection
conn
      IdentifySystem
_ <- ReplicantException -> Maybe IdentifySystem -> IO IdentifySystem
forall e a. Exception e => e -> Maybe a -> IO a
maybeThrow (String -> ReplicantException
ReplicantException String
"withLogicalStream: could not get system information") Maybe IdentifySystem
maybeInfo
      ReplicationSlotInfo
repSlot <- Connection -> ByteString -> IO ReplicationSlotInfo
setupReplicationSlot Connection
conn (ByteString -> IO ReplicationSlotInfo)
-> ByteString -> IO ReplicationSlotInfo
forall a b. (a -> b) -> a -> b
$ String -> ByteString
B.pack (String -> ByteString)
-> (PgSettings -> String) -> PgSettings -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PgSettings -> String
pgSlotName (PgSettings -> ByteString) -> PgSettings -> ByteString
forall a b. (a -> b) -> a -> b
$ PgSettings
settings
      Connection -> ByteString -> LSN -> Int -> (Change -> IO a) -> IO ()
forall a.
Connection -> ByteString -> LSN -> Int -> (Change -> IO a) -> IO ()
startReplicationStream Connection
conn (ReplicationSlotInfo -> ByteString
slotName ReplicationSlotInfo
repSlot) (ReplicationSlotInfo -> LSN
slotRestart ReplicationSlotInfo
repSlot) Int
updateFreq Change -> IO a
cb
      () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  where
    pollConnectStart :: Connection -> Fd -> IO PollingStatus
    pollConnectStart :: Connection -> Fd -> IO PollingStatus
pollConnectStart Connection
conn fd :: Fd
fd@(Fd CInt
cint) = do
      PollingStatus
pollStatus <- Connection -> IO PollingStatus
connectPoll Connection
conn
      case PollingStatus
pollStatus of
        PollingStatus
PollingReading -> do
          Fd -> IO ()
threadWaitRead Fd
fd
          Connection -> Fd -> IO PollingStatus
pollConnectStart Connection
conn Fd
fd
        PollingStatus
PollingWriting -> do
          Fd -> IO ()
threadWaitWrite Fd
fd
          Connection -> Fd -> IO PollingStatus
pollConnectStart Connection
conn Fd
fd
        PollingStatus
PollingOk -> do
          Either KeepAliveError ()
_ <- CInt -> KeepAlive -> IO (Either KeepAliveError ())
setKeepAlive CInt
cint (KeepAlive -> IO (Either KeepAliveError ()))
-> KeepAlive -> IO (Either KeepAliveError ())
forall a b. (a -> b) -> a -> b
$ Bool -> Word32 -> Word32 -> KeepAlive
KeepAlive Bool
True Word32
60 Word32
2
          PollingStatus -> IO PollingStatus
forall (f :: * -> *) a. Applicative f => a -> f a
pure PollingStatus
PollingOk
        PollingStatus
PollingFailed -> PollingStatus -> IO PollingStatus
forall (f :: * -> *) a. Applicative f => a -> f a
pure PollingStatus
PollingFailed
    getUpdateDelay :: PgSettings -> Int
    getUpdateDelay :: PgSettings -> Int
getUpdateDelay PgSettings {String
pgUpdateDelay :: String
pgSlotName :: String
pgPort :: String
pgHost :: String
pgDbName :: String
pgUser :: String
pgUpdateDelay :: PgSettings -> String
pgSlotName :: PgSettings -> String
pgPort :: PgSettings -> String
pgHost :: PgSettings -> String
pgDbName :: PgSettings -> String
pgUser :: PgSettings -> String
..} =
      case Reader Int
forall a. Integral a => Reader a
T.decimal Reader Int
-> (String -> Text) -> String -> Either String (Int, Text)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
T.pack (String -> Either String (Int, Text))
-> String -> Either String (Int, Text)
forall a b. (a -> b) -> a -> b
$ String
pgUpdateDelay of
        Left String
_ -> Int
3000
        Right (Int
i, Text
_) -> Int
i