diff --git a/README.md b/README.md index f04438e..ec0c7a2 100644 --- a/README.md +++ b/README.md @@ -12,4 +12,4 @@ - [ ] Implement JSON decoding - [ ] Implement `ByteString` decoding (`bytea`) - Can we make the fromField instance choose whether it wants binary or text? -- [ ] Clean up and document column table stuff +- [x] Clean up and document column table stuff diff --git a/flake.nix b/flake.nix index 805e2ec..aa14d9b 100644 --- a/flake.nix +++ b/flake.nix @@ -21,6 +21,7 @@ postgresql-libpq text transformers + vector ])) pkgs.haskell-language-server diff --git a/lib/Database/PostgreSQL/Opium.hs b/lib/Database/PostgreSQL/Opium.hs index 8739585..0d734e0 100644 --- a/lib/Database/PostgreSQL/Opium.hs +++ b/lib/Database/PostgreSQL/Opium.hs @@ -10,11 +10,13 @@ {-# LANGUAGE TypeOperators #-} module Database.PostgreSQL.Opium - ( Error (..) + ( ColumnTable + , Error (..) , ErrorPosition (..) , FromField (..) , FromRow (..) , fetch_ + , toListColumnTable ) where @@ -24,6 +26,7 @@ import Data.ByteString (ByteString) import Data.IORef (IORef, modifyIORef', newIORef, readIORef) import Data.Proxy (Proxy (..)) import Data.Text (Text) +import Data.Vector (Vector) import Database.PostgreSQL.LibPQ ( Column , Connection @@ -36,6 +39,7 @@ import GHC.TypeLits (KnownSymbol, symbolVal) import qualified Data.Text as Text import qualified Data.Text.Encoding as Encoding +import qualified Data.Vector as Vector import qualified Database.PostgreSQL.LibPQ as LibPQ import Database.PostgreSQL.Opium.Error (Error (..), ErrorPosition (..)) @@ -57,20 +61,26 @@ execParams conn query = do fetch_ :: forall a. FromRow a => Connection -> ByteString -> IO (Either Error [a]) fetch_ conn query = runExceptT $ do result <- execParams conn query - -- TODO: Use unboxed array for columnTable columnTable <- ExceptT $ getColumnTable @a Proxy result nRows <- liftIO $ LibPQ.ntuples result mapM (ExceptT . fromRow result columnTable) [0..nRows - 1] -type ColumnTable = [(Column, Oid)] +newtype ColumnTable = ColumnTable (Vector (Column, Oid)) + deriving (Eq, Show) + +newColumnTable :: [(Column, Oid)] -> ColumnTable +newColumnTable = ColumnTable . Vector.fromList indexColumnTable :: ColumnTable -> Int -> (Column, Oid) -indexColumnTable = (!!) +indexColumnTable (ColumnTable v) i = v `Vector.unsafeIndex` i + +toListColumnTable :: ColumnTable -> [(Column, Oid)] +toListColumnTable (ColumnTable v) = Vector.toList v class FromRow a where getColumnTable :: Proxy a -> Result -> IO (Either Error ColumnTable) default getColumnTable :: (Generic a, GetColumnTable' (Rep a)) => Proxy a -> Result -> IO (Either Error ColumnTable) - getColumnTable Proxy = runExceptT . getColumnTable' @(Rep a) Proxy + getColumnTable Proxy = runExceptT . fmap newColumnTable . getColumnTable' @(Rep a) Proxy fromRow :: Result -> ColumnTable -> Row -> IO (Either Error a) default fromRow :: (Generic a, FromRow' (Rep a)) => Result -> ColumnTable -> Row -> IO (Either Error a) @@ -79,7 +89,7 @@ class FromRow a where runExceptT $ to <$> fromRow' (FromRowCtx result columnTable iRef) row class GetColumnTable' f where - getColumnTable' :: Proxy (f p) -> Result -> ExceptT Error IO ColumnTable + getColumnTable' :: Proxy (f p) -> Result -> ExceptT Error IO [(Column, Oid)] instance GetColumnTable' f => GetColumnTable' (M1 D c f) where getColumnTable' Proxy = getColumnTable' @f Proxy @@ -91,7 +101,7 @@ instance (GetColumnTable' f, GetColumnTable' g) => GetColumnTable' (f :*: g) whe getColumnTable' Proxy result = (++) <$> getColumnTable' @f Proxy result <*> getColumnTable' @g Proxy result -checkColumn :: forall a. FromField a => Proxy a -> String -> Result -> ExceptT Error IO ColumnTable +checkColumn :: forall a. FromField a => Proxy a -> String -> Result -> ExceptT Error IO [(Column, Oid)] checkColumn Proxy nameStr result = do column <- ExceptT $ maybe (Left $ ErrorMissingColumn nameText) Right <$> LibPQ.fnumber result name oid <- liftIO $ LibPQ.ftype result column diff --git a/opium.cabal b/opium.cabal index 4b97eb6..44a17bd 100644 --- a/opium.cabal +++ b/opium.cabal @@ -79,7 +79,8 @@ library containers, postgresql-libpq, text, - transformers + transformers, + vector -- Directories containing source files. hs-source-dirs: lib diff --git a/test/Database/PostgreSQL/OpiumSpec.hs b/test/Database/PostgreSQL/OpiumSpec.hs index 15124e2..3fb3d6b 100644 --- a/test/Database/PostgreSQL/OpiumSpec.hs +++ b/test/Database/PostgreSQL/OpiumSpec.hs @@ -5,6 +5,7 @@ module Database.PostgreSQL.OpiumSpec (spec) where +import Data.ByteString (ByteString) import Data.Proxy (Proxy (..)) import Data.Text (Text) import Database.PostgreSQL.LibPQ (Connection) @@ -41,22 +42,35 @@ isLeft :: Either a b -> Bool isLeft (Left _) = True isLeft _ = False +shouldHaveColumns + :: Opium.FromRow a + => Proxy a + -> Connection + -> ByteString + -> [LibPQ.Column] + -> IO () +shouldHaveColumns proxy conn query expectedColumns = do + Just result <- LibPQ.execParams conn query [] LibPQ.Text + columnTable <- Opium.getColumnTable proxy result + let actualColumns = fmap (map fst . Opium.toListColumnTable) columnTable + actualColumns `shouldBe` Right expectedColumns + spec :: SpecWith Connection spec = do describe "getColumnTable" $ do it "Gets the column table for a result" $ \conn -> do - Just result <- LibPQ.execParams conn "SELECT name, age FROM person" [] LibPQ.Text - columnTable <- Opium.getColumnTable @Person Proxy result - fmap (map fst) columnTable `shouldBe` Right [0, 1] + shouldHaveColumns @Person Proxy conn + "SELECT name, age FROM person" + [0, 1] it "Gets the numbers right for funky configurations" $ \conn -> do - Just result0 <- LibPQ.execParams conn "SELECT age, name FROM person" [] LibPQ.Text - columnTable0 <- Opium.getColumnTable @Person Proxy result0 - fmap (map fst) columnTable0 `shouldBe` Right [1, 0] + shouldHaveColumns @Person Proxy conn + "SELECT age, name FROM person" + [1, 0] - Just result1 <- LibPQ.execParams conn "SELECT 0 AS a, 1 AS b, 2 AS c, age, 4 AS d, name FROM person" [] LibPQ.Text - columnTable1 <- Opium.getColumnTable @Person Proxy result1 - fmap (map fst) columnTable1 `shouldBe` Right [5, 3] + shouldHaveColumns @Person Proxy conn + "SELECT 0 AS a, 1 AS b, 2 AS c, age, 4 AS d, name FROM person" + [5, 3] it "Fails for missing columns" $ \conn -> do Just result <- LibPQ.execParams conn "SELECT 0 AS a FROM person" [] LibPQ.Text