{-|
Module: Squeal.PostgreSQL.Session.Result
Description: results
Copyright: (c) Eitan Chatav, 2019
Maintainer: eitan@morphism.tech
Stability: experimental

Get values from a `Result`.
-}

{-# LANGUAGE
    FlexibleContexts
  , FlexibleInstances
  , GADTs
  , LambdaCase
  , OverloadedStrings
  , ScopedTypeVariables
  , TypeApplications
  , UndecidableInstances
#-}

module Squeal.PostgreSQL.Session.Result
  ( Result (..)
  , MonadResult (..)
  , liftResult
  , nextRow
  ) where

import Control.Exception (throw)
import Control.Monad (when, (<=<))
import Control.Monad.Catch
import Control.Monad.IO.Class
import Data.ByteString (ByteString)
import Data.Text (Text)
import Data.Traversable (for)
import Text.Read (readMaybe)

import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Char8 as Char8
import qualified Data.Text.Encoding as Text
import qualified Database.PostgreSQL.LibPQ as LibPQ
import qualified Generics.SOP as SOP

import Squeal.PostgreSQL.Session.Decode
import Squeal.PostgreSQL.Session.Exception

{- | `Result`s are generated by executing
`Squeal.PostgreSQL.Session.Statement`s
in a `Squeal.PostgreSQL.Session.Monad.MonadPQ`.

They contain an underlying `LibPQ.Result`
and a `DecodeRow`.
-}
data Result y where
  Result
    :: SOP.SListI row
    => DecodeRow row y
    -> LibPQ.Result
    -> Result y
instance Functor Result where
  fmap :: forall a b. (a -> b) -> Result a -> Result b
fmap a -> b
f (Result DecodeRow row a
decode Result
result) = forall (row :: RowType) y.
SListI row =>
DecodeRow row y -> Result -> Result y
Result (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f DecodeRow row a
decode) Result
result

{- | A `MonadResult` operation extracts values
from the `Result` of a `Squeal.PostgreSQL.Session.Monad.MonadPQ` operation.
There is no need to define instances of `MonadResult`.
An instance of `MonadIO` implies an instance of `MonadResult`.
However, the constraint `MonadResult`
does not imply the constraint `MonadIO`.
-}
class Monad m => MonadResult m where
  -- | Get a row corresponding to a given row number from a `LibPQ.Result`,
  -- throwing an exception if the row number is out of bounds.
  getRow :: LibPQ.Row -> Result y -> m y
  -- | Get all rows from a `LibPQ.Result`.
  getRows :: Result y -> m [y]
  -- | Get the first row if possible from a `LibPQ.Result`.
  firstRow :: Result y -> m (Maybe y)
  -- | Returns the number of rows (tuples) in the query result.
  ntuples :: Result y -> m LibPQ.Row
  -- | Returns the number of columns (fields) in the query result.
  nfields :: Result y -> m LibPQ.Column
  {- |
  Returns the command status tag from the SQL command
  that generated the `Result`.
  Commonly this is just the name of the command,
  but it might include additional data such as the number of rows processed.
  -}
  cmdStatus :: Result y -> m Text
  {- |
  Returns the number of rows affected by the SQL command.
  This function returns `Just` the number of
  rows affected by the SQL statement that generated the `Result`.
  This function can only be used following the execution of a
  SELECT, CREATE TABLE AS, INSERT, UPDATE, DELETE, MOVE, FETCH,
  or COPY statement,or an EXECUTE of a prepared query that
  contains an INSERT, UPDATE, or DELETE statement.
  If the command that generated the PGresult was anything else,
  `cmdTuples` returns `Nothing`.
  -}
  cmdTuples :: Result y -> m (Maybe LibPQ.Row)
  -- | Returns the result status of the command.
  resultStatus :: Result y -> m LibPQ.ExecStatus
  -- | Check if a `Result`'s status is either `LibPQ.CommandOk`
  -- or `LibPQ.TuplesOk` otherwise `throw` a `SQLException`.
  okResult :: Result y -> m ()
  -- | Returns the error message most recently generated by an operation
  -- on the connection.
  resultErrorMessage :: Result y -> m (Maybe ByteString)
  -- | Returns the error code most recently generated by an operation
  -- on the connection.
  --
  -- https://www.postgresql.org/docs/current/static/errcodes-appendix.html
  resultErrorCode :: Result y -> m (Maybe ByteString)

instance (Monad io, MonadIO io) => MonadResult io where
  getRow :: forall y. Row -> Result y -> io y
getRow Row
r (Result DecodeRow row y
decode Result
result) = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    Row
numRows <- Result -> IO Row
LibPQ.ntuples Result
result
    Column
numCols <- Result -> IO Column
LibPQ.nfields Result
result
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Row
numRows forall a. Ord a => a -> a -> Bool
< Row
r) forall a b. (a -> b) -> a -> b
$ forall a e. Exception e => e -> a
throw forall a b. (a -> b) -> a -> b
$ Text -> Row -> Row -> SquealException
RowsException Text
"getRow" Row
r Row
numRows
    [Maybe ByteString]
row' <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Result -> Row -> Column -> IO (Maybe ByteString)
LibPQ.getvalue Result
result Row
r) [Column
0 .. Column
numCols forall a. Num a => a -> a -> a
- Column
1]
    case forall {k} (xs :: [k]) a. SListI xs => [a] -> Maybe (NP (K a) xs)
SOP.fromList [Maybe ByteString]
row' of
      Maybe (NP (K (Maybe ByteString)) row)
Nothing -> forall a e. Exception e => e -> a
throw forall a b. (a -> b) -> a -> b
$ Text -> Column -> SquealException
ColumnsException Text
"getRow" Column
numCols
      Just NP (K (Maybe ByteString)) row
row -> case forall (row :: RowType) y.
DecodeRow row y -> NP (K (Maybe ByteString)) row -> Either Text y
execDecodeRow DecodeRow row y
decode NP (K (Maybe ByteString)) row
row of
        Left Text
parseError -> forall a e. Exception e => e -> a
throw forall a b. (a -> b) -> a -> b
$ Text -> Text -> SquealException
DecodingException Text
"getRow" Text
parseError
        Right y
y -> forall (m :: * -> *) a. Monad m => a -> m a
return y
y

  getRows :: forall y. Result y -> io [y]
getRows (Result DecodeRow row y
decode Result
result) = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    Column
numCols <- Result -> IO Column
LibPQ.nfields Result
result
    Row
numRows <- Result -> IO Row
LibPQ.ntuples Result
result
    forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Row
0 .. Row
numRows forall a. Num a => a -> a -> a
- Row
1] forall a b. (a -> b) -> a -> b
$ \ Row
r -> do
      [Maybe ByteString]
row' <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Result -> Row -> Column -> IO (Maybe ByteString)
LibPQ.getvalue Result
result Row
r) [Column
0 .. Column
numCols forall a. Num a => a -> a -> a
- Column
1]
      case forall {k} (xs :: [k]) a. SListI xs => [a] -> Maybe (NP (K a) xs)
SOP.fromList [Maybe ByteString]
row' of
        Maybe (NP (K (Maybe ByteString)) row)
Nothing -> forall a e. Exception e => e -> a
throw forall a b. (a -> b) -> a -> b
$ Text -> Column -> SquealException
ColumnsException Text
"getRows" Column
numCols
        Just NP (K (Maybe ByteString)) row
row -> case forall (row :: RowType) y.
DecodeRow row y -> NP (K (Maybe ByteString)) row -> Either Text y
execDecodeRow DecodeRow row y
decode NP (K (Maybe ByteString)) row
row of
          Left Text
parseError -> forall a e. Exception e => e -> a
throw forall a b. (a -> b) -> a -> b
$ Text -> Text -> SquealException
DecodingException Text
"getRows" Text
parseError
          Right y
y -> forall (m :: * -> *) a. Monad m => a -> m a
return y
y

  firstRow :: forall y. Result y -> io (Maybe y)
firstRow (Result DecodeRow row y
decode Result
result) = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    Row
numRows <- Result -> IO Row
LibPQ.ntuples Result
result
    Column
numCols <- Result -> IO Column
LibPQ.nfields Result
result
    if Row
numRows forall a. Ord a => a -> a -> Bool
<= Row
0 then forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing else do
      [Maybe ByteString]
row' <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Result -> Row -> Column -> IO (Maybe ByteString)
LibPQ.getvalue Result
result Row
0) [Column
0 .. Column
numCols forall a. Num a => a -> a -> a
- Column
1]
      case forall {k} (xs :: [k]) a. SListI xs => [a] -> Maybe (NP (K a) xs)
SOP.fromList [Maybe ByteString]
row' of
        Maybe (NP (K (Maybe ByteString)) row)
Nothing -> forall a e. Exception e => e -> a
throw forall a b. (a -> b) -> a -> b
$ Text -> Column -> SquealException
ColumnsException Text
"firstRow" Column
numCols
        Just NP (K (Maybe ByteString)) row
row -> case forall (row :: RowType) y.
DecodeRow row y -> NP (K (Maybe ByteString)) row -> Either Text y
execDecodeRow DecodeRow row y
decode NP (K (Maybe ByteString)) row
row of
          Left Text
parseError -> forall a e. Exception e => e -> a
throw forall a b. (a -> b) -> a -> b
$ Text -> Text -> SquealException
DecodingException Text
"firstRow" Text
parseError
          Right y
y -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just y
y

  ntuples :: forall y. Result y -> io Row
ntuples = forall (io :: * -> *) x y.
MonadIO io =>
(Result -> IO x) -> Result y -> io x
liftResult Result -> IO Row
LibPQ.ntuples

  nfields :: forall y. Result y -> io Column
nfields = forall (io :: * -> *) x y.
MonadIO io =>
(Result -> IO x) -> Result y -> io x
liftResult Result -> IO Column
LibPQ.nfields

  resultStatus :: forall y. Result y -> io ExecStatus
resultStatus = forall (io :: * -> *) x y.
MonadIO io =>
(Result -> IO x) -> Result y -> io x
liftResult Result -> IO ExecStatus
LibPQ.resultStatus

  cmdStatus :: forall y. Result y -> io Text
cmdStatus = forall (io :: * -> *) x y.
MonadIO io =>
(Result -> IO x) -> Result y -> io x
liftResult (Maybe ByteString -> IO Text
getCmdStatus forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Result -> IO (Maybe ByteString)
LibPQ.cmdStatus)
    where
      getCmdStatus :: Maybe ByteString -> IO Text
getCmdStatus = \case
        Maybe ByteString
Nothing -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ Text -> SquealException
ConnectionException Text
"LibPQ.cmdStatus"
        Just ByteString
bytes -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ ByteString -> Text
Text.decodeUtf8 ByteString
bytes

  cmdTuples :: forall y. Result y -> io (Maybe Row)
cmdTuples = forall (io :: * -> *) x y.
MonadIO io =>
(Result -> IO x) -> Result y -> io x
liftResult (Maybe ByteString -> IO (Maybe Row)
getCmdTuples forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Result -> IO (Maybe ByteString)
LibPQ.cmdTuples)
    where
      getCmdTuples :: Maybe ByteString -> IO (Maybe Row)
getCmdTuples = \case
        Maybe ByteString
Nothing -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ Text -> SquealException
ConnectionException Text
"LibPQ.cmdTuples"
        Just ByteString
bytes -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
          if ByteString -> Bool
ByteString.null ByteString
bytes
          then forall a. Maybe a
Nothing
          else forall a. Num a => Integer -> a
fromInteger forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Read a => String -> Maybe a
readMaybe (ByteString -> String
Char8.unpack ByteString
bytes)

  okResult :: forall y. Result y -> io ()
okResult = forall (io :: * -> *) x y.
MonadIO io =>
(Result -> IO x) -> Result y -> io x
liftResult forall (io :: * -> *). MonadIO io => Result -> io ()
okResult_ 

  resultErrorMessage :: forall y. Result y -> io (Maybe ByteString)
resultErrorMessage = forall (io :: * -> *) x y.
MonadIO io =>
(Result -> IO x) -> Result y -> io x
liftResult Result -> IO (Maybe ByteString)
LibPQ.resultErrorMessage

  resultErrorCode :: forall y. Result y -> io (Maybe ByteString)
resultErrorCode = forall (io :: * -> *) x y.
MonadIO io =>
(Result -> IO x) -> Result y -> io x
liftResult (forall a b c. (a -> b -> c) -> b -> a -> c
flip Result -> FieldCode -> IO (Maybe ByteString)
LibPQ.resultErrorField FieldCode
LibPQ.DiagSqlstate)

-- | Intended to be used for unfolding in streaming libraries, `nextRow`
-- takes a total number of rows (which can be found with `ntuples`)
-- and a `LibPQ.Result` and given a row number if it's too large returns `Nothing`,
-- otherwise returning the row along with the next row number.
nextRow
  :: MonadIO io
  => LibPQ.Row -- ^ total number of rows
  -> Result y -- ^ result
  -> LibPQ.Row -- ^ row number
  -> io (Maybe (LibPQ.Row, y))
nextRow :: forall (io :: * -> *) y.
MonadIO io =>
Row -> Result y -> Row -> io (Maybe (Row, y))
nextRow Row
total (Result DecodeRow row y
decode Result
result) Row
r
  = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ if Row
r forall a. Ord a => a -> a -> Bool
>= Row
total then forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing else do
    Column
numCols <- Result -> IO Column
LibPQ.nfields Result
result
    [Maybe ByteString]
row' <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Result -> Row -> Column -> IO (Maybe ByteString)
LibPQ.getvalue Result
result Row
r) [Column
0 .. Column
numCols forall a. Num a => a -> a -> a
- Column
1]
    case forall {k} (xs :: [k]) a. SListI xs => [a] -> Maybe (NP (K a) xs)
SOP.fromList [Maybe ByteString]
row' of
      Maybe (NP (K (Maybe ByteString)) row)
Nothing -> forall a e. Exception e => e -> a
throw forall a b. (a -> b) -> a -> b
$ Text -> Column -> SquealException
ColumnsException Text
"nextRow" Column
numCols
      Just NP (K (Maybe ByteString)) row
row -> case forall (row :: RowType) y.
DecodeRow row y -> NP (K (Maybe ByteString)) row -> Either Text y
execDecodeRow DecodeRow row y
decode NP (K (Maybe ByteString)) row
row of
        Left Text
parseError -> forall a e. Exception e => e -> a
throw forall a b. (a -> b) -> a -> b
$ Text -> Text -> SquealException
DecodingException Text
"nextRow" Text
parseError
        Right y
y -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Row
rforall a. Num a => a -> a -> a
+Row
1, y
y)

okResult_ :: MonadIO io => LibPQ.Result -> io ()
okResult_ :: forall (io :: * -> *). MonadIO io => Result -> io ()
okResult_ Result
result = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
  ExecStatus
status <- Result -> IO ExecStatus
LibPQ.resultStatus Result
result
  case ExecStatus
status of
    ExecStatus
LibPQ.CommandOk -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
    ExecStatus
LibPQ.TuplesOk -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
    ExecStatus
_ -> do
      Maybe ByteString
stateCodeMaybe <- Result -> FieldCode -> IO (Maybe ByteString)
LibPQ.resultErrorField Result
result FieldCode
LibPQ.DiagSqlstate
      case Maybe ByteString
stateCodeMaybe of
        Maybe ByteString
Nothing -> forall a e. Exception e => e -> a
throw forall a b. (a -> b) -> a -> b
$ Text -> SquealException
ConnectionException Text
"LibPQ.resultErrorField"
        Just ByteString
stateCode -> do
          Maybe ByteString
msgMaybe <- Result -> IO (Maybe ByteString)
LibPQ.resultErrorMessage Result
result
          case Maybe ByteString
msgMaybe of
            Maybe ByteString
Nothing -> forall a e. Exception e => e -> a
throw forall a b. (a -> b) -> a -> b
$ Text -> SquealException
ConnectionException Text
"LibPQ.resultErrorMessage"
            Just ByteString
msg -> forall a e. Exception e => e -> a
throw forall b c a. (b -> c) -> (a -> b) -> a -> c
. SQLState -> SquealException
SQLException forall a b. (a -> b) -> a -> b
$ ExecStatus -> ByteString -> ByteString -> SQLState
SQLState ExecStatus
status ByteString
stateCode ByteString
msg

-- | Lifts actions on results from @LibPQ@.
liftResult
  :: MonadIO io
  => (LibPQ.Result -> IO x)
  -> Result y -> io x
liftResult :: forall (io :: * -> *) x y.
MonadIO io =>
(Result -> IO x) -> Result y -> io x
liftResult Result -> IO x
f (Result DecodeRow row y
_ Result
result) = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Result -> IO x
f Result
result

execDecodeRow
  :: DecodeRow row y
  -> SOP.NP (SOP.K (Maybe ByteString)) row
  -> Either Text y
execDecodeRow :: forall (row :: RowType) y.
DecodeRow row y -> NP (K (Maybe ByteString)) row -> Either Text y
execDecodeRow DecodeRow row y
decode = forall (row :: RowType) y.
DecodeRow row y -> NP (K (Maybe ByteString)) row -> Either Text y
runDecodeRow DecodeRow row y
decode