From e7d494c034f2edb78d792f8031e8dfc9ad44fbb8 Mon Sep 17 00:00:00 2001
From: Paul Brinkmeier <hallo@pbrinkmeier.de>
Date: Sat, 16 Sep 2023 04:36:42 +0200
Subject: [PATCH] Add getColumnTable function

---
 lib/Database/PostgreSQL/Opium.hs      | 36 +++++++++++++++++++++++++++
 test/Database/PostgreSQL/OpiumSpec.hs | 16 ++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/lib/Database/PostgreSQL/Opium.hs b/lib/Database/PostgreSQL/Opium.hs
index 49339b5..d15a32c 100644
--- a/lib/Database/PostgreSQL/Opium.hs
+++ b/lib/Database/PostgreSQL/Opium.hs
@@ -6,6 +6,7 @@
 {-# LANGUAGE KindSignatures #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeOperators #-}
 
 module Database.PostgreSQL.Opium
@@ -61,10 +62,45 @@ fetchResult result = do
   runExceptT $ mapM (ExceptT . flip fromRow result) [0..nRows - 1]
 
 class FromRow a where
+  getColumnTable :: Proxy a -> Result -> IO [LibPQ.Column]
+  default getColumnTable :: (Generic a, GetColumnTable' (Rep a)) => Proxy a -> Result -> IO [LibPQ.Column]
+  getColumnTable Proxy = getColumnTable' @(Rep a) Proxy
+
   fromRow :: Row -> Result -> IO (Either Error a)
   default fromRow :: (Generic a, FromRow' (Rep a)) => Row -> Result -> IO (Either Error a)
   fromRow row result = fmap to <$> fromRow' row result
 
+class GetColumnTable' f where
+  getColumnTable' :: Proxy (f p) -> Result -> IO [LibPQ.Column]
+
+instance GetColumnTable' f => GetColumnTable' (M1 D c f) where
+  getColumnTable' Proxy = getColumnTable' @f Proxy
+
+instance GetColumnTable' f => GetColumnTable' (M1 C c f) where
+  getColumnTable' Proxy = getColumnTable' @f Proxy
+
+instance (GetColumnTable' f, GetColumnTable' g) => GetColumnTable' (f :*: g) where
+  getColumnTable' Proxy result =
+    (++) <$> getColumnTable' @f Proxy result <*> getColumnTable' @g Proxy result
+
+checkColumn :: FromField f => Proxy f -> String -> Result -> IO [Column]
+checkColumn Proxy nameStr result = LibPQ.fnumber result name >>= \case
+    Just column -> do
+      -- TODO: Rewrite FromField to check whether oid works for decoding t
+      _oid <- LibPQ.ftype result column
+      pure [column]
+    Nothing -> do
+      -- TODO: Return ErrorMissingColumn
+      undefined
+    where
+      name = Encoding.encodeUtf8 $ Text.pack nameStr
+
+instance {-# OVERLAPPABLE #-} (FromField t, KnownSymbol nameSym) => GetColumnTable' (M1 S ('MetaSel ('Just nameSym) nu ns dl) (Rec0 t)) where
+  getColumnTable' Proxy = checkColumn @t Proxy $ symbolVal @nameSym Proxy
+
+instance {-# OVERLAPPING #-} (KnownSymbol nameSym, FromField t) => GetColumnTable' (M1 S ('MetaSel ('Just nameSym) nu ns dl) (Rec0 (Maybe t))) where
+  getColumnTable' Proxy = checkColumn @t Proxy $ symbolVal @nameSym Proxy
+
 class FromRow' f where
   fromRow' :: Row -> Result -> IO (Either Error (f p))
 
diff --git a/test/Database/PostgreSQL/OpiumSpec.hs b/test/Database/PostgreSQL/OpiumSpec.hs
index f09eeb3..bd61898 100644
--- a/test/Database/PostgreSQL/OpiumSpec.hs
+++ b/test/Database/PostgreSQL/OpiumSpec.hs
@@ -4,6 +4,7 @@
 
 module Database.PostgreSQL.OpiumSpec (spec) where
 
+import Data.Proxy (Proxy (..))
 import Data.Text (Text)
 import Database.PostgreSQL.LibPQ (Connection)
 import GHC.Generics (Generic)
@@ -31,6 +32,21 @@ isLeft _ = False
 
 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 (Proxy :: Proxy Person) result
+      columnTable `shouldBe` [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
+      columnTable0 `shouldBe` [1, 0]
+
+      Just result1 <- LibPQ.execParams conn "SELECT 0 a, 1 b, 2 c, age, 4 d, name FROM person" [] LibPQ.Text
+      columnTable1 <- Opium.getColumnTable @Person Proxy result1
+      columnTable1 `shouldBe` [5, 3]
+
   describe "fromRow" $ do
     it "decodes rows in a Result" $ \conn -> do
       Just result <- LibPQ.execParams conn "SELECT * FROM person" [] LibPQ.Text