RowSpec.hs

{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}

module PostgreSQL.Result.RowSpec (spec) where

import Test.Hspec (Spec, describe, it, shouldBe)

import           Control.Monad (when)
import qualified Control.Monad.Except as Except
import qualified Control.Monad.Reader as Reader
import           Data.Bifunctor (first)
import           Data.ByteString (ByteString)
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Massiv.Array as Massiv
import           Data.Traversable (for)
import qualified Data.Vector as Vector
import           Database.PostgreSQL.LibPQ (Format (Text), Oid, invalidOid)
import qualified PostgreSQL.Result.Cell as Cell
import qualified PostgreSQL.Result.Column as Column
import qualified PostgreSQL.Result.Row as Row
import           PostgreSQL.Types (ColumnNum, ProcessorError (..), ProcessorErrors, Value (..))

type Matrix = Massiv.Array Massiv.B Massiv.Ix2

runMatrixRow
  :: Vector.Vector (Oid, Format)
  -> HashMap.HashMap ByteString ColumnNum
  -> Matrix Value
  -> Row.Row a
  -> Either ProcessorErrors [a]
runMatrixRow columns names datas row = Except.runExcept $ do
  let
    Massiv.Sz2 numRows' numCols' = Massiv.size datas

    numRows = fromIntegral numRows'

    numCols = min (fromIntegral numCols') $ fromIntegral $ Vector.length columns

  runner <- Row.runRow row $ \req -> do
    col <-
      case Row.columnRequest_position req of
        Row.FixedColumn col -> pure col

        Row.NamedColumn name ->
          case HashMap.lookup name names of
            Just col -> pure col
            Nothing -> Except.throwError [MissingNamedColumn name]

    when (col >= numCols) $ Except.throwError [NotEnoughColumns col numCols]

    let (oid, format) = columns Vector.! fromIntegral col

    cell <-
      Except.liftEither $ first (fmap (ColumnParserError col oid format)) $
        Column.parseColumn (Row.columnRequest_parser req) oid format

    pure $ Reader.ReaderT $ \row -> do
      let value = Massiv.index' datas $ Massiv.Ix2 (fromIntegral row) $ fromIntegral col
      Except.liftEither $ first (fmap (CellParserError col oid format row value)) $
        Cell.parseCell cell value

  for [0 .. numRows - 1] (Reader.runReaderT runner)

spec :: Spec
spec =
  describe "Row" $ do
    it "works as an Applicative" $
      runMatrixRow [] [] [[], []] (pure ())
        `shouldBe` Right [(), ()]

    it "works on single columns" $
      runMatrixRow
        [(invalidOid, Text)]
        []
        [["Hello"], ["World"]]
        (Row.columnWith Column.raw)
        `shouldBe` Right ["Hello", "World"]

    it "works on multiple columns" $
      runMatrixRow
        [(invalidOid, Text), (invalidOid, Text)]
        []
        [["1", "Hello"], ["2", "World"]]
        (replicate <$> Row.column <*> Row.columnWith Column.raw)
        `shouldBe` Right [["Hello"], ["World", "World"]]

    it "resolves named columns" $
      runMatrixRow
        [(invalidOid, Text), (invalidOid, Text)]
        [("b", 1)]
        [["1", "Hello"], ["2", "World"]]
        (Row.namedColumnWith "b" Column.raw)
        `shouldBe` Right ["Hello", "World"]

    it "works with fixed columns" $
      runMatrixRow
        [(invalidOid, Text), (invalidOid, Text)]
        []
        [["1", "Hello"], ["2", "World"]]
        (Row.fixedColumnWith 1 Column.raw)
        `shouldBe` Right ["Hello", "World"]

    it "fails on too few columns" $
      runMatrixRow
        [(invalidOid, Text)]
        []
        [["Hello"]]
        (Row.fixedColumnWith 1 Column.raw)
        `shouldBe` Left [NotEnoughColumns 1 1]

    it "fails on unknown named column" $
      runMatrixRow
        [(invalidOid, Text)]
        []
        [["Hello"]]
        (Row.namedColumnWith "unknown" Column.raw)
        `shouldBe` Left [MissingNamedColumn "unknown"]