{-# LANGUAGE CPP #-}

------------------------------------------------------------------------------
-- |
-- Module:      Database.PostgreSQL.Simple.Internal.PQResultUtils
-- Copyright:   (c) 2011 MailRank, Inc.
--              (c) 2011-2012 Leon P Smith
-- License:     BSD3
-- Maintainer:  Leon P Smith <leon@melding-monads.com>
-- Stability:   experimental
--
------------------------------------------------------------------------------


module Database.PostgreSQL.Simple.Internal.PQResultUtils
    ( finishQueryWith
    , finishQueryWithV
    , finishQueryWithVU
    , getRowWith
    ) where

import           Control.Exception as E
import           Data.ByteString (ByteString)
import           Data.Foldable (for_)
import           Database.PostgreSQL.Simple.FromField (ResultError(..))
import           Database.PostgreSQL.Simple.Ok
import           Database.PostgreSQL.Simple.Types (Query(..))
import           Database.PostgreSQL.Simple.Internal as Base hiding (result, row)
import           Database.PostgreSQL.Simple.TypeInfo
import qualified Database.PostgreSQL.LibPQ as PQ
import qualified Data.ByteString.Char8 as B
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as MVU
import           Control.Monad.Trans.Reader
import           Control.Monad.Trans.State.Strict

finishQueryWith :: RowParser r -> Connection -> Query -> PQ.Result -> IO [r]
finishQueryWith :: RowParser r -> Connection -> Query -> Result -> IO [r]
finishQueryWith RowParser r
parser Connection
conn Query
q Result
result = Query -> Result -> IO [r] -> IO [r]
forall a. Query -> Result -> IO a -> IO a
finishQueryWith' Query
q Result
result (IO [r] -> IO [r]) -> IO [r] -> IO [r]
forall a b. (a -> b) -> a -> b
$ do
    Row
nrows <- Result -> IO Row
PQ.ntuples Result
result
    Column
ncols <- Result -> IO Column
PQ.nfields Result
result
    Row -> Row -> (Row -> IO r) -> IO [r]
forall n a. (Ord n, Num n) => n -> n -> (n -> IO a) -> IO [a]
forM' Row
0 (Row
nrowsRow -> Row -> Row
forall a. Num a => a -> a -> a
-Row
1) ((Row -> IO r) -> IO [r]) -> (Row -> IO r) -> IO [r]
forall a b. (a -> b) -> a -> b
$ \Row
row ->
        RowParser r -> Row -> Column -> Connection -> Result -> IO r
forall r.
RowParser r -> Row -> Column -> Connection -> Result -> IO r
getRowWith RowParser r
parser Row
row Column
ncols Connection
conn Result
result

finishQueryWithV :: RowParser r -> Connection -> Query -> PQ.Result -> IO (V.Vector r)
finishQueryWithV :: RowParser r -> Connection -> Query -> Result -> IO (Vector r)
finishQueryWithV RowParser r
parser Connection
conn Query
q Result
result = Query -> Result -> IO (Vector r) -> IO (Vector r)
forall a. Query -> Result -> IO a -> IO a
finishQueryWith' Query
q Result
result (IO (Vector r) -> IO (Vector r)) -> IO (Vector r) -> IO (Vector r)
forall a b. (a -> b) -> a -> b
$ do
    Row
nrows <- Result -> IO Row
PQ.ntuples Result
result
    let PQ.Row CInt
nrows' = Row
nrows
    Column
ncols <- Result -> IO Column
PQ.nfields Result
result
    MVector RealWorld r
mv <- Int -> IO (MVector (PrimState IO) r)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
MV.unsafeNew (CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CInt
nrows')
    [Row] -> (Row -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [ Row
0 .. Row
nrowsRow -> Row -> Row
forall a. Num a => a -> a -> a
-Row
1 ] ((Row -> IO ()) -> IO ()) -> (Row -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Row
row -> do
        let PQ.Row CInt
row' = Row
row
        r
value <- RowParser r -> Row -> Column -> Connection -> Result -> IO r
forall r.
RowParser r -> Row -> Column -> Connection -> Result -> IO r
getRowWith RowParser r
parser Row
row Column
ncols Connection
conn Result
result
        MVector (PrimState IO) r -> Int -> r -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite MVector RealWorld r
MVector (PrimState IO) r
mv (CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CInt
row') r
value
    MVector (PrimState IO) r -> IO (Vector r)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.unsafeFreeze MVector RealWorld r
MVector (PrimState IO) r
mv

finishQueryWithVU :: VU.Unbox r => RowParser r -> Connection -> Query -> PQ.Result -> IO (VU.Vector r)
finishQueryWithVU :: RowParser r -> Connection -> Query -> Result -> IO (Vector r)
finishQueryWithVU RowParser r
parser Connection
conn Query
q Result
result = Query -> Result -> IO (Vector r) -> IO (Vector r)
forall a. Query -> Result -> IO a -> IO a
finishQueryWith' Query
q Result
result (IO (Vector r) -> IO (Vector r)) -> IO (Vector r) -> IO (Vector r)
forall a b. (a -> b) -> a -> b
$ do
    Row
nrows <- Result -> IO Row
PQ.ntuples Result
result
    let PQ.Row CInt
nrows' = Row
nrows
    Column
ncols <- Result -> IO Column
PQ.nfields Result
result
    MVector RealWorld r
mv <- Int -> IO (MVector (PrimState IO) r)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MVU.unsafeNew (CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CInt
nrows')
    [Row] -> (Row -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [ Row
0 .. Row
nrowsRow -> Row -> Row
forall a. Num a => a -> a -> a
-Row
1 ] ((Row -> IO ()) -> IO ()) -> (Row -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Row
row -> do
        let PQ.Row CInt
row' = Row
row
        r
value <- RowParser r -> Row -> Column -> Connection -> Result -> IO r
forall r.
RowParser r -> Row -> Column -> Connection -> Result -> IO r
getRowWith RowParser r
parser Row
row Column
ncols Connection
conn Result
result
        MVector (PrimState IO) r -> Int -> r -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MVU.unsafeWrite MVector RealWorld r
MVector (PrimState IO) r
mv (CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CInt
row') r
value
    MVector (PrimState IO) r -> IO (Vector r)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld r
MVector (PrimState IO) r
mv

finishQueryWith' :: Query -> PQ.Result -> IO a -> IO a
finishQueryWith' :: Query -> Result -> IO a -> IO a
finishQueryWith' Query
q Result
result IO a
k = do
  ExecStatus
status <- Result -> IO ExecStatus
PQ.resultStatus Result
result
  case ExecStatus
status of
    ExecStatus
PQ.TuplesOk -> IO a
k
    ExecStatus
PQ.EmptyQuery    -> String -> IO a
forall a. String -> IO a
queryErr String
"query: Empty query"
    ExecStatus
PQ.CommandOk     -> String -> IO a
forall a. String -> IO a
queryErr String
"query resulted in a command response (did you mean to use `execute` or forget a RETURNING?)"
    ExecStatus
PQ.CopyOut       -> String -> IO a
forall a. String -> IO a
queryErr String
"query: COPY TO is not supported"
    ExecStatus
PQ.CopyIn        -> String -> IO a
forall a. String -> IO a
queryErr String
"query: COPY FROM is not supported"
#if MIN_VERSION_postgresql_libpq(0,9,3)
    ExecStatus
PQ.CopyBoth      -> String -> IO a
forall a. String -> IO a
queryErr String
"query: COPY BOTH is not supported"
#endif
#if MIN_VERSION_postgresql_libpq(0,9,2)
    ExecStatus
PQ.SingleTuple   -> String -> IO a
forall a. String -> IO a
queryErr String
"query: single-row mode is not supported"
#endif
    ExecStatus
PQ.BadResponse   -> ByteString -> Result -> ExecStatus -> IO a
forall a. ByteString -> Result -> ExecStatus -> IO a
throwResultError ByteString
"query" Result
result ExecStatus
status
    ExecStatus
PQ.NonfatalError -> ByteString -> Result -> ExecStatus -> IO a
forall a. ByteString -> Result -> ExecStatus -> IO a
throwResultError ByteString
"query" Result
result ExecStatus
status
    ExecStatus
PQ.FatalError    -> ByteString -> Result -> ExecStatus -> IO a
forall a. ByteString -> Result -> ExecStatus -> IO a
throwResultError ByteString
"query" Result
result ExecStatus
status
  where
    queryErr :: String -> IO a
queryErr String
msg = QueryError -> IO a
forall e a. Exception e => e -> IO a
throwIO (QueryError -> IO a) -> QueryError -> IO a
forall a b. (a -> b) -> a -> b
$ String -> Query -> QueryError
QueryError String
msg Query
q

getRowWith :: RowParser r -> PQ.Row -> PQ.Column -> Connection -> PQ.Result -> IO r
getRowWith :: RowParser r -> Row -> Column -> Connection -> Result -> IO r
getRowWith RowParser r
parser Row
row Column
ncols Connection
conn Result
result = do
  let rw :: Row
rw = Row -> Result -> Row
Row Row
row Result
result
  let unCol :: Column -> Int
unCol (PQ.Col CInt
x) = CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CInt
x :: Int
  Ok (r, Column)
okvc <- Conversion (r, Column) -> Connection -> IO (Ok (r, Column))
forall a. Conversion a -> Connection -> IO (Ok a)
runConversion (StateT Column Conversion r -> Column -> Conversion (r, Column)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (ReaderT Row (StateT Column Conversion) r
-> Row -> StateT Column Conversion r
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (RowParser r -> ReaderT Row (StateT Column Conversion) r
forall a. RowParser a -> ReaderT Row (StateT Column Conversion) a
unRP RowParser r
parser) Row
rw) Column
0) Connection
conn
  case Ok (r, Column)
okvc of
    Ok (r
val,Column
col) | Column
col Column -> Column -> Bool
forall a. Eq a => a -> a -> Bool
== Column
ncols -> r -> IO r
forall (m :: * -> *) a. Monad m => a -> m a
return r
val
                 | Bool
otherwise -> do
                     [(TypeInfo, Maybe ByteString)]
vals <- Column
-> Column
-> (Column -> IO (TypeInfo, Maybe ByteString))
-> IO [(TypeInfo, Maybe ByteString)]
forall n a. (Ord n, Num n) => n -> n -> (n -> IO a) -> IO [a]
forM' Column
0 (Column
ncolsColumn -> Column -> Column
forall a. Num a => a -> a -> a
-Column
1) ((Column -> IO (TypeInfo, Maybe ByteString))
 -> IO [(TypeInfo, Maybe ByteString)])
-> (Column -> IO (TypeInfo, Maybe ByteString))
-> IO [(TypeInfo, Maybe ByteString)]
forall a b. (a -> b) -> a -> b
$ \Column
c -> do
                         TypeInfo
tinfo <- Connection -> Oid -> IO TypeInfo
getTypeInfo Connection
conn (Oid -> IO TypeInfo) -> IO Oid -> IO TypeInfo
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Result -> Column -> IO Oid
PQ.ftype Result
result Column
c
                         Maybe ByteString
v <- Result -> Row -> Column -> IO (Maybe ByteString)
PQ.getvalue Result
result Row
row Column
c
                         (TypeInfo, Maybe ByteString) -> IO (TypeInfo, Maybe ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return ( TypeInfo
tinfo
                                , (ByteString -> ByteString) -> Maybe ByteString -> Maybe ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> ByteString
ellipsis Maybe ByteString
v       )
                     ResultError -> IO r
forall a e. Exception e => e -> a
throw (String -> Maybe Oid -> String -> String -> String -> ResultError
ConversionFailed
                      (Int -> String
forall a. Show a => a -> String
show (Column -> Int
unCol Column
ncols) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" values: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [(TypeInfo, Maybe ByteString)] -> String
forall a. Show a => a -> String
show [(TypeInfo, Maybe ByteString)]
vals)
                      Maybe Oid
forall a. Maybe a
Nothing
                      String
""
                      (Int -> String
forall a. Show a => a -> String
show (Column -> Int
unCol Column
col) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" slots in target type")
                      String
"mismatch between number of columns to convert and number in target type")
    Errors []  -> ResultError -> IO r
forall e a. Exception e => e -> IO a
throwIO (ResultError -> IO r) -> ResultError -> IO r
forall a b. (a -> b) -> a -> b
$ String -> Maybe Oid -> String -> String -> String -> ResultError
ConversionFailed String
"" Maybe Oid
forall a. Maybe a
Nothing String
"" String
"" String
"unknown error"
    Errors [SomeException
x] -> SomeException -> IO r
forall e a. Exception e => e -> IO a
throwIO SomeException
x
    Errors [SomeException]
xs  -> ManyErrors -> IO r
forall e a. Exception e => e -> IO a
throwIO (ManyErrors -> IO r) -> ManyErrors -> IO r
forall a b. (a -> b) -> a -> b
$ [SomeException] -> ManyErrors
ManyErrors [SomeException]
xs

ellipsis :: ByteString -> ByteString
ellipsis :: ByteString -> ByteString
ellipsis ByteString
bs
    | ByteString -> Int
B.length ByteString
bs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
15 = Int -> ByteString -> ByteString
B.take Int
10 ByteString
bs ByteString -> ByteString -> ByteString
`B.append` ByteString
"[...]"
    | Bool
otherwise        = ByteString
bs

forM' :: (Ord n, Num n) => n -> n -> (n -> IO a) -> IO [a]
forM' :: n -> n -> (n -> IO a) -> IO [a]
forM' n
lo n
hi n -> IO a
m = n -> [a] -> IO [a]
loop n
hi []
  where
    loop :: n -> [a] -> IO [a]
loop !n
n ![a]
as
      | n
n n -> n -> Bool
forall a. Ord a => a -> a -> Bool
< n
lo = [a] -> IO [a]
forall (m :: * -> *) a. Monad m => a -> m a
return [a]
as
      | Bool
otherwise = do
           a
a <- n -> IO a
m n
n
           n -> [a] -> IO [a]
loop (n
nn -> n -> n
forall a. Num a => a -> a -> a
-n
1) (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
as)
{-# INLINE forM' #-}