Connection.hs

{-# OPTIONS_GHC -fno-warn-orphans -fno-warn-partial-type-signatures #-}

{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE BangPatterns #-}

module Database.Beam.Postgres.Connection
  ( Pg(..), PgF(..)

  , liftIOWithHandle

  , runBeamPostgres, runBeamPostgresDebug

  , pgRenderSyntax, runPgRowReader, getFields

  , withPgDebug

  , postgresUriSyntax ) where

import           Control.Exception (SomeException(..), throwIO, onException, catch)
import           Control.Monad (void)
import           Control.Monad.Base (MonadBase(..))
import           Control.Monad.Free.Church
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Control (MonadBaseControl(..))
import           Data.IORef (newIORef, readIORef, writeIORef)
import           Data.Vector (Vector)
import qualified Data.Vector as V

import           Database.Beam hiding (runDelete, runUpdate, runInsert, insert)
import           Database.Beam.Backend.SQL.BeamExtensions
import           Database.Beam.Backend.SQL.Row ( FromBackendRowF(..), FromBackendRowM(..)
                                               , BeamRowReadError(..), ColumnParseError(..) )
import           Database.Beam.Backend.URI
import           Database.Beam.Schema.Tables

import           Database.Beam.Postgres.Extensions.Copy.File
                   ( PgCopyFromSyntax(..), PgCopyToSyntax(..) )
import           Database.Beam.Postgres.Extensions.Copy.Stream
                   ( PgCopyFromStreamSyntax(..), PgCopyToStreamSyntax(..) )
import           Database.Beam.Postgres.Syntax
import           Database.Beam.Postgres.Full
import           Database.Beam.Postgres.Types

import qualified Database.PostgreSQL.LibPQ as Pg hiding
  (Connection, escapeStringConn, escapeIdentifier, escapeByteaConn, exec)
import qualified Database.PostgreSQL.Simple as Pg
import qualified Database.PostgreSQL.Simple.Copy as PgCopy
import qualified Database.PostgreSQL.Simple.FromField as Pg
import qualified Database.PostgreSQL.Simple.Internal as Pg
  ( Field(..), RowParser(..)
  , escapeStringConn, escapeIdentifier, escapeByteaConn
  , exec, throwResultError )
import qualified Database.PostgreSQL.Simple.Internal as PgI
import qualified Database.PostgreSQL.Simple.Ok as Pg
import qualified Database.PostgreSQL.Simple.Types as Pg (Query(..))

import           Control.Monad.Reader
import           Control.Monad.State
import qualified Control.Monad.Fail as Fail

import           Data.ByteString (ByteString)
import           Data.ByteString.Builder (toLazyByteString, byteString)
import qualified Data.ByteString.Lazy as BL
import           Data.Maybe (listToMaybe, fromMaybe)
import           Data.Proxy
import           Data.String
import qualified Data.Text as T
import           Data.Text.Encoding (decodeUtf8)
import           Data.Typeable (cast)

import           Foreign.C.Types

import           Network.URI (uriToString)

data PgStream a = PgStreamDone     (Either BeamRowReadError a)
                | PgStreamContinue (Maybe PgI.Row -> IO (PgStream a))

-- | 'BeamURIOpeners' for the standard @postgresql:@ URI scheme. See the
-- postgres documentation for more details on the formatting. See documentation
-- for 'BeamURIOpeners' for more information on how to use this with beam
postgresUriSyntax :: c Postgres Pg.Connection Pg
                  -> BeamURIOpeners c
postgresUriSyntax =
    mkUriOpener runBeamPostgres "postgresql:"
        (\uri -> do
            let pgConnStr = fromString (uriToString id uri "")
            hdl <- Pg.connectPostgreSQL pgConnStr
            pure (hdl, Pg.close hdl))

-- * Syntax rendering

pgRenderSyntax ::
  Pg.Connection -> PgSyntax -> IO ByteString
pgRenderSyntax conn (PgSyntax mkQuery) =
  renderBuilder <$> runF mkQuery finish step mempty
  where
    renderBuilder = BL.toStrict . toLazyByteString

    step (EmitBuilder b next) a = next (a <> b)
    step (EmitByteString b next) a = next (a <> byteString b)
    step (EscapeString b next) a = do
      res <- wrapError "EscapeString" (Pg.escapeStringConn conn b)
      next (a <> byteString res)
    step (EscapeBytea b next) a = do
      res <- wrapError "EscapeBytea" (Pg.escapeByteaConn conn b)
      next (a <> byteString res)
    step (EscapeIdentifier b next) a = do
      res <- wrapError "EscapeIdentifier" (Pg.escapeIdentifier conn b)
      next (a <> byteString res)

    finish _ = pure

    wrapError step' go = do
      res <- go
      case res of
        Right res' -> pure res'
        Left res' -> fail (step' <> ": " <> show res')

-- * Run row readers

getFields :: Pg.Result -> IO (Vector Pg.Field)
getFields res = do
  Pg.Col colCount <- Pg.nfields res
  V.generateM (fromIntegral colCount) $ \i ->
    let col = Pg.Col (fromIntegral i)
    in Pg.Field res col <$> Pg.ftype res col

runPgRowReader ::
  Pg.Connection -> Pg.Row -> Pg.Result -> Vector Pg.Field -> FromBackendRowM Postgres a -> IO (Either BeamRowReadError a)
runPgRowReader conn rowIdx res fields (FromBackendRowM readRow) =
  -- 'colCount' and 'fields' are both invariant for the duration of one
  -- 'runPgRowReader' call, so we capture them in the closure of 'step'
  -- rather than threading them through the free-monad result type.
  -- That makes the per-step closure smaller and avoids an O(n) 'length'
  -- per row (Vector stores its length, so 'V.length' is O(1)).
  runF readRow finish step 0
  where
    !colCount = fromIntegral (V.length fields) :: CInt

    step :: forall x. FromBackendRowF Postgres (CInt -> IO (Either BeamRowReadError x))
         -> CInt -> IO (Either BeamRowReadError x)
    step (ParseOneField _) curCol
      | curCol >= colCount =
          pure (Left (BeamRowReadError (Just (fromIntegral curCol))
                                       (ColumnNotEnoughColumns (fromIntegral colCount))))
    step (ParseOneField (next' :: next -> _)) curCol =
      do let field = V.unsafeIndex fields (fromIntegral curCol)
         fieldValue <- Pg.getvalue' res rowIdx (Pg.Col curCol)
         res' <- Pg.runConversion (Pg.fromField field fieldValue) conn
         case res' of
           Pg.Errors errs ->
             let err = fromMaybe (ColumnErrorInternal "Column parse failed with unknown exception") $
                       listToMaybe $
                       do SomeException e <- errs
                          Just pgErr <- pure (cast e)
                          case pgErr of
                            Pg.ConversionFailed { Pg.errSQLType = sql
                                                , Pg.errHaskellType = hs
                                                , Pg.errMessage = msg
                                                , Pg.errSQLField = errField } ->
                              pure (ColumnTypeMismatch hs sql ("Conversion failed for field'" <> errField <> "': " <> msg))
                            Pg.Incompatible { Pg.errSQLType = sql
                                            , Pg.errHaskellType = hs
                                            , Pg.errMessage = msg
                                            , Pg.errSQLField = errField } ->
                              pure (ColumnTypeMismatch hs sql ("Incompatible field: '" <> errField <> "': " <> msg))
                            Pg.UnexpectedNull {} ->
                              pure ColumnUnexpectedNull
             in pure (Left (BeamRowReadError (Just (fromIntegral curCol)) err))
           Pg.Ok x -> next' x (curCol + 1)

    step (Alt (FromBackendRowM a) (FromBackendRowM b) next) curCol =
      do aRes <- runF a (\x curCol' -> pure (Right (next x curCol'))) step curCol
         case aRes of
           Right next' -> next'
           Left aErr -> do
             bRes <- runF b (\x curCol' -> pure (Right (next x curCol'))) step curCol
             case bRes of
               Right next' -> next'
               Left {} -> pure (Left aErr)

    step (FailParseWith err) _ =
      pure (Left err)

    finish x _ = pure (Right x)

withPgDebug :: (String -> IO ()) -> Pg.Connection -> Pg a -> IO (Either BeamRowReadError a)
withPgDebug dbg conn (Pg action) = do
  -- One-entry cache for the cursor-batch path: 'Pg.Result' is constant
  -- within a batch but changes between batches. Caching by 'Pg.Result'
  -- equality (a 'ForeignPtr' comparison, ~free) avoids the redundant
  -- 'getFields' / 'nfields' / 'ftype' calls within each batch.
  --
  -- Default batch size is 256 rows, set by 'postgresql-simple'
  -- 'defaultFoldOptions' (FetchQuantity = Automatic, which resolves to
  -- 256 in 'Database.PostgreSQL.Simple').
  fieldsCache <- newIORef (Nothing :: Maybe (Pg.Result, Vector Pg.Field))

  let cachedGetFields :: Pg.Result -> IO (Vector Pg.Field)
      cachedGetFields res = do
        cached <- readIORef fieldsCache
        case cached of
          Just (cachedRes, fs) | cachedRes == res -> pure fs
          _ -> do
            fs <- getFields res
            writeIORef fieldsCache (Just (res, fs))
            pure fs

      finish x = pure (Right x)
      step (PgLiftIO io next) = io >>= next
      step (PgLiftWithHandle withConn next) = withConn dbg conn >>= next
      step (PgFetchNext next) = next Nothing
      step (PgRunReturning CursorBatching
                           (PgCommandSyntax PgCommandTypeQuery syntax)
                           (mkProcess :: Pg (Maybe x) -> Pg a')
                           next) =
        do query <- pgRenderSyntax conn syntax
           let Pg process = mkProcess (Pg (liftF (PgFetchNext id)))
           dbg (T.unpack (decodeUtf8 query))
           action' <- runF process finishProcess stepProcess Nothing
           case action' of
             PgStreamDone (Right x) -> Pg.execute_ conn (Pg.Query query) >> next x
             PgStreamDone (Left err) -> pure (Left err)
             PgStreamContinue nextStream ->
               let finishUp (PgStreamDone (Right x)) = next x
                   finishUp (PgStreamDone (Left err)) = pure (Left err)
                   finishUp (PgStreamContinue next') = next' Nothing >>= finishUp

                   columnCount = fromIntegral $ valuesNeeded (Proxy @Postgres) (Proxy @x)
               in Pg.foldWith_ (Pg.RP (put columnCount >> ask)) conn (Pg.Query query) (PgStreamContinue nextStream) runConsumer >>= finishUp
      step (PgRunReturning AtOnce
                           (PgCommandSyntax PgCommandTypeQuery syntax)
                           (mkProcess :: Pg (Maybe x) -> Pg a')
                           next) =
        renderExecReturningList "No tuples returned to Postgres query" syntax mkProcess next
      step (PgRunReturning _ (PgCommandSyntax PgCommandTypeDataUpdateReturning syntax) mkProcess next) =
        renderExecReturningList "No tuples returned to Postgres update/insert returning" syntax mkProcess next
      step (PgRunReturning _ (PgCommandSyntax _ syntax) mkProcess next) =
        do query <- pgRenderSyntax conn syntax
           dbg (T.unpack (decodeUtf8 query))
           _ <- Pg.execute_ conn (Pg.Query query)

           let Pg process = mkProcess (Pg (liftF (PgFetchNext id)))
           runF process next stepReturningNone

      renderExecReturningList :: (FromBackendRow Postgres x) => _ -> PgSyntax -> (Pg (Maybe x) -> Pg a') -> _ -> _
      renderExecReturningList errMsg syntax mkProcess next =
        do query <- pgRenderSyntax conn syntax
           dbg (T.unpack (decodeUtf8 query))

           res <- Pg.exec conn query
           sts <- Pg.resultStatus res
           case sts of
             Pg.TuplesOk -> do
               -- Hoist per-query metadata out of the per-row loop: the
               -- same 'Pg.Result' is used for every row, so 'fields'
               -- and 'rowCount' are loop-invariant.
               -- Use getFields directly: unsafeFreeResult below lets libpq
               -- reuse the pointer address, which would cause a false hit in
               -- cachedGetFields on the next query.
               fields <- getFields res
               Pg.Row rowCount <- Pg.ntuples res
               let Pg process = mkProcess (Pg (liftF (PgFetchNext id)))
               runF process (\x _ -> Pg.unsafeFreeResult res >> next x)
                            (stepReturningList fields rowCount res) 0
             _ -> Pg.throwResultError errMsg res sts

      stepReturningNone :: forall a. PgF (IO (Either BeamRowReadError a)) -> IO (Either BeamRowReadError a)
      stepReturningNone (PgLiftIO action' next) = action' >>= next
      stepReturningNone (PgLiftWithHandle withConn next) = withConn dbg conn >>= next
      stepReturningNone (PgFetchNext next) = next Nothing
      stepReturningNone (PgRunReturning {}) = pure (Left (BeamRowReadError Nothing (ColumnErrorInternal  "Nested queries not allowed")))

      stepReturningList :: forall a. Vector Pg.Field -> CInt -> Pg.Result
                        -> PgF (CInt -> IO (Either BeamRowReadError a))
                        -> CInt -> IO (Either BeamRowReadError a)
      stepReturningList _      _        _   (PgLiftIO action' next) rowIdx = action' >>= \x -> next x rowIdx
      stepReturningList fields rowCount res (PgFetchNext next)      rowIdx =
        if rowIdx >= rowCount
          then next Nothing rowIdx
          else runPgRowReader conn (Pg.Row rowIdx) res fields fromBackendRow >>= \case
                 Left err -> pure (Left err)
                 Right r  -> next (Just r) (rowIdx + 1)
      stepReturningList _      _        _   (PgRunReturning {})     _      = pure (Left (BeamRowReadError Nothing (ColumnErrorInternal "Nested queries not allowed")))
      stepReturningList _      _        _   (PgLiftWithHandle {})   _      = pure (Left (BeamRowReadError Nothing (ColumnErrorInternal "Nested queries not allowed")))

      finishProcess :: forall a. a -> Maybe PgI.Row -> IO (PgStream a)
      finishProcess x _ = pure (PgStreamDone (Right x))

      stepProcess :: forall a. PgF (Maybe PgI.Row -> IO (PgStream a)) -> Maybe PgI.Row -> IO (PgStream a)
      stepProcess (PgLiftIO action' next) row = action' >>= flip next row
      stepProcess (PgFetchNext next) Nothing =
        pure . PgStreamContinue $ \res ->
        case res of
          Nothing -> next Nothing Nothing
          Just (PgI.Row rowIdx res') ->
            cachedGetFields res' >>= \fields ->
            runPgRowReader conn rowIdx res' fields fromBackendRow >>= \case
              Left err -> pure (PgStreamDone (Left err))
              Right r -> next (Just r) Nothing
      stepProcess (PgFetchNext next) (Just (PgI.Row rowIdx res)) =
        cachedGetFields res >>= \fields ->
        runPgRowReader conn rowIdx res fields fromBackendRow >>= \case
          Left err -> pure (PgStreamDone (Left err))
          Right r -> pure (PgStreamContinue (next (Just r)))
      stepProcess (PgRunReturning {}) _ = pure (PgStreamDone (Left (BeamRowReadError Nothing (ColumnErrorInternal "Nested queries not allowed"))))
      stepProcess (PgLiftWithHandle _ _) _ = pure (PgStreamDone (Left (BeamRowReadError Nothing (ColumnErrorInternal "Nested queries not allowed"))))

      runConsumer :: forall a. PgStream a -> PgI.Row -> IO (PgStream a)
      runConsumer s@(PgStreamDone {}) _ = pure s
      runConsumer (PgStreamContinue next) row = next (Just row)

  runF action finish step

-- * Beam Monad class

data PgF next where
    PgLiftIO :: IO a -> (a -> next) -> PgF next
    PgRunReturning ::
        FromBackendRow Postgres x =>
        FetchMode -> PgCommandSyntax -> (Pg (Maybe x) -> Pg a) -> (a -> next) -> PgF next
    PgFetchNext ::
        FromBackendRow Postgres x =>
        (Maybe x -> next) -> PgF next
    PgLiftWithHandle :: ((String -> IO ()) -> Pg.Connection -> IO a) -> (a -> next) -> PgF next
instance Functor PgF where
  fmap f = \case
    PgLiftIO io n -> PgLiftIO io $ f . n
    PgRunReturning mode cmd consume n -> PgRunReturning mode cmd consume $ f . n
    PgFetchNext n -> PgFetchNext $ f . n
    PgLiftWithHandle withConn n -> PgLiftWithHandle withConn $ f . n

-- | How to fetch results.
data FetchMode
    = CursorBatching -- ^ Fetch in batches of ~256 rows via cursor for SELECT.
    | AtOnce         -- ^ Fetch all rows at once.

-- | 'MonadBeam' in which we can run Postgres commands. See the documentation
-- for 'MonadBeam' on examples of how to use.
--
-- @beam-postgres@ also provides functions that let you run queries without
-- 'MonadBeam'. These functions may be more efficient and offer a conduit
-- API. See "Database.Beam.Postgres.Conduit" for more information.
--
-- You can execute 'Pg' actions using 'runBeamPostgres' or 'runBeamPostgresDebug'.
newtype Pg a = Pg { runPg :: F PgF a }
    deriving (Monad, Applicative, Functor, MonadFree PgF)

instance Fail.MonadFail Pg where
    fail e =  liftIO (Fail.fail $ "Internal Error with: " <> show e)

instance MonadIO Pg where
    liftIO x = liftF (PgLiftIO x id)

instance MonadBase IO Pg where
    liftBase = liftIO

instance MonadBaseControl IO Pg where
    type StM Pg a = a

    liftBaseWith action =
      liftF (PgLiftWithHandle (\dbg conn -> action (runBeamPostgresDebug dbg conn)) id)

    restoreM = pure

liftIOWithHandle :: (Pg.Connection -> IO a) -> Pg a
liftIOWithHandle f = liftF (PgLiftWithHandle (\_ -> f) id)

runBeamPostgresDebug :: (String -> IO ()) -> Pg.Connection -> Pg a -> IO a
runBeamPostgresDebug dbg conn action =
    withPgDebug dbg conn action >>= either throwIO pure

runBeamPostgres :: Pg.Connection -> Pg a -> IO a
runBeamPostgres = runBeamPostgresDebug (\_ -> pure ())

instance MonadBeam Postgres Pg where
    runReturningMany cmd consume =
        liftF (PgRunReturning CursorBatching cmd consume id)

    runReturningOne cmd =
        liftF (PgRunReturning AtOnce cmd consume id)
      where
        consume next = do
          a <- next
          case a of
            Nothing -> pure Nothing
            Just x -> do
              a' <- next
              case a' of
                Nothing -> pure (Just x)
                Just _ -> pure Nothing

    runReturningFirst cmd =
        liftF (PgRunReturning AtOnce cmd id id)

    runReturningList cmd =
        liftF (PgRunReturning AtOnce cmd consume id)
      where
        consume next =
          let collectM acc = do
                a <- next
                case a of
                  Nothing -> pure (acc [])
                  Just x -> collectM (acc . (x:))
          in collectM id

instance MonadBeamCopyTo Postgres Pg where
    runCopyTo SqlCopyToNoColumns = pure ()
    runCopyTo (SqlCopyTo (PgCopyToSyntax syntax)) =
        runNoReturn (PgCommandSyntax PgCommandTypeDataUpdate syntax)

instance MonadBeamCopyFrom Postgres Pg where
    runCopyFrom SqlCopyFromNoColumns = pure ()
    runCopyFrom (SqlCopyFrom (PgCopyFromSyntax syntax)) =
        runNoReturn (PgCommandSyntax PgCommandTypeDataUpdate syntax)

instance MonadBeamCopyToStream Postgres Pg where
    -- | `runCopyToStream` is exception-safe; if the output stream
    -- produces an exception, the stream will be drained before the exception
    -- is re-thrown
    runCopyToStream SqlCopyToStreamNoColumns _ = pure ()
    runCopyToStream (SqlCopyToStream (PgCopyToStreamSyntax syntax)) sink =
        liftIOWithHandle $ \conn -> do
          query <- pgRenderSyntax conn syntax
          PgCopy.copy_ conn (Pg.Query query)
          let loop onRow = do
                PgCopy.getCopyData conn >>= \case
                  PgCopy.CopyOutRow chunk -> onRow chunk >> loop onRow
                  PgCopy.CopyOutDone _    -> pure ()
              -- Like 'loop', but doesn't use the sink at all. This is used
              -- to drain elements in the COPY stream before re-throwing an exception
              drain = loop (const (pure ()))
          loop sink `catch` (\(e::SomeException) -> drain >> throwIO e)


instance MonadBeamCopyFromStream Postgres Pg where
    -- | `runCopyFromStream` is exception-safe; if the input stream
    -- produces an exception, the connection will be reset to a safe state,
    -- aborting the COPY operation.
    runCopyFromStream SqlCopyFromStreamNoColumns _ = pure ()
    runCopyFromStream (SqlCopyFromStream (PgCopyFromStreamSyntax syntax)) producer =
        liftIOWithHandle $ \conn -> do
          query <- pgRenderSyntax conn syntax
          let loop = 
                producer >>= \case
                  Just chunk -> PgCopy.putCopyData conn chunk >> loop
                  Nothing    -> void $ PgCopy.putCopyEnd conn
          PgCopy.copy_ conn (Pg.Query query)
          loop `onException` PgCopy.putCopyError conn mempty

instance MonadBeamInsertReturning Postgres Pg where
    runInsertReturningListWith i mkProjection = do
        let pgProj tbl = mkProjection (changeBeamRep
              (\(Columnar' (QExpr s) :: Columnar' (QExpr Postgres PostgresInaccessible) ty) ->
                Columnar' (QExpr s) :: Columnar' (QExpr Postgres ()) ty) tbl)
            insertReturningCmd' = i `returning` pgProj
        case insertReturningCmd' of
          PgInsertReturningEmpty ->
            pure []
          PgInsertReturning insertReturningCmd ->
            runReturningList (PgCommandSyntax PgCommandTypeDataUpdateReturning insertReturningCmd)

instance MonadBeamUpdateReturning Postgres Pg where
    runUpdateReturningListWith u mkProjection = do
        let pgProj tbl = mkProjection (changeBeamRep
              (\(Columnar' (QExpr s) :: Columnar' (QExpr Postgres PostgresInaccessible) ty) ->
                Columnar' (QExpr s) :: Columnar' (QExpr Postgres ()) ty) tbl)
            updateReturningCmd' = u `returning` pgProj
        case updateReturningCmd' of
          PgUpdateReturningEmpty ->
            pure []
          PgUpdateReturning updateReturningCmd ->
            runReturningList (PgCommandSyntax PgCommandTypeDataUpdateReturning updateReturningCmd)

instance MonadBeamDeleteReturning Postgres Pg where
    runDeleteReturningListWith d mkProjection = do
        let pgProj tbl = mkProjection (changeBeamRep
              (\(Columnar' (QExpr s) :: Columnar' (QExpr Postgres PostgresInaccessible) ty) ->
                Columnar' (QExpr s) :: Columnar' (QExpr Postgres ()) ty) tbl)
            PgDeleteReturning deleteReturningCmd = d `returning` pgProj
        runReturningList (PgCommandSyntax PgCommandTypeDataUpdateReturning deleteReturningCmd)