From cc4844167f83987458b707d1d0df2c6bd535942c Mon Sep 17 00:00:00 2001
From: Alberto Valverde <alberto@toscat.net>
Date: Tue, 23 Dec 2014 14:13:07 +0100
Subject: [PATCH] Implemented support for custom/composite primary keys

---
 src/Database/Esqueleto/Internal/Sql.hs | 20 ++++++++++----------
 test/Test.hs                           | 19 +++++++++++++++++++
 2 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/src/Database/Esqueleto/Internal/Sql.hs b/src/Database/Esqueleto/Internal/Sql.hs
index 342b69e..e62eade 100644
--- a/src/Database/Esqueleto/Internal/Sql.hs
+++ b/src/Database/Esqueleto/Internal/Sql.hs
@@ -7,6 +7,7 @@
            , MultiParamTypeClasses
            , OverloadedStrings
            , UndecidableInstances
+           , ScopedTypeVariables
  #-}
 -- | This is an internal module, anything exported by this module
 -- may change without a major version bump.  Please use only
@@ -58,6 +59,8 @@ import Data.List (intersperse)
 import Data.Monoid (Monoid(..), (<>))
 import Data.Proxy (Proxy(..))
 import Database.Esqueleto.Internal.PersistentImport
+import Database.Persist.Sql.Util (
+    entityColumnNames, entityColumnCount, parseEntityValues)
 import qualified Control.Monad.Trans.State as S
 import qualified Control.Monad.Trans.Writer as W
 import qualified Data.Conduit as C
@@ -993,11 +996,9 @@ instance SqlSelect () () where
 instance PersistEntity a => SqlSelect (SqlExpr (Entity a)) (Entity a) where
   sqlSelectCols info expr@(EEntity ident) = ret
       where
-        process ed = uncommas $
-                     map ((name <>) . fromDBName info) $
-                     map fieldDB $
-                     entityId ed :
-                     entityFields ed
+        conn = fst info
+        process ed = uncommas . map ((name <>) . TLB.fromText) $
+                     entityColumnNames ed conn
         -- 'name' is the biggest difference between 'RawSql' and
         -- 'SqlSelect'.  We automatically create names for tables
         -- (since it's not the user who's writing the FROM
@@ -1007,11 +1008,10 @@ instance PersistEntity a => SqlSelect (SqlExpr (Entity a)) (Entity a) where
         name = useIdent info ident <> "."
         ret = let ed = entityDef $ getEntityVal $ return expr
               in (process ed, mempty)
-  sqlSelectColCount = (+1) . length . entityFields . entityDef . getEntityVal
-  sqlSelectProcessRow (idCol:ent) =
-    Entity <$> fromPersistValue idCol
-           <*> fromPersistValues ent
-  sqlSelectProcessRow _ = Left "SqlSelect (Entity a): wrong number of columns."
+  sqlSelectColCount = entityColumnCount . entityDef . getEntityVal
+  sqlSelectProcessRow = parseEntityValues ed
+    where ed = entityDef $ getEntityVal $ (Proxy :: Proxy (SqlExpr (Entity a)))
+
 
 getEntityVal :: Proxy (SqlExpr (Entity a)) -> Proxy a
 getEntityVal = const Proxy
diff --git a/test/Test.hs b/test/Test.hs
index 41ba06f..b362b7b 100644
--- a/test/Test.hs
+++ b/test/Test.hs
@@ -1,6 +1,8 @@
 {-# LANGUAGE ConstraintKinds
            , EmptyDataDecls
            , FlexibleContexts
+           , FlexibleInstances
+           , DeriveGeneric
            , GADTs
            , GeneralizedNewtypeDeriving
            , MultiParamTypeClasses
@@ -63,6 +65,11 @@ share [mkPersist sqlSettings, mkMigrate "migrateAll"] [persistUpperCase|
     number Int
     Primary number
     deriving Eq Show
+  Point
+    x Int
+    y Int
+    Primary x y
+    deriving Eq Show
 |]
 
 -- | this could be achieved with S.fromList, but not all lists
@@ -924,6 +931,18 @@ main = do
             ret `shouldBe` fc
             fcPk `shouldBe` thePk
 
+      it "works with composite primary key" $
+        run $ do
+          let p = Point x y
+              x = 10
+              y = 15
+              Right thePk = keyFromValues [ PersistInt64 $ fromIntegral x
+                                          , PersistInt64 $ fromIntegral y]
+          pPk <- insert p
+          [Entity _ ret] <- select $ from $ return
+          liftIO $ do
+            ret `shouldBe` p
+            pPk `shouldBe` thePk
 ----------------------------------------------------------------------