From df63cd864b477da4d83b91fcc4b680548a9ec869 Mon Sep 17 00:00:00 2001
From: Felipe Lessa <felipe.lessa@gmail.com>
Date: Wed, 5 Sep 2012 23:17:37 -0300
Subject: [PATCH] COUNT(*).

---
 src/Database/Esqueleto.hs                   |  4 ++--
 src/Database/Esqueleto/Internal/Language.hs |  3 +++
 src/Database/Esqueleto/Internal/Sql.hs      |  3 ++-
 test/Test.hs                                | 24 ++++++++++++++++++++-
 4 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/Database/Esqueleto.hs b/src/Database/Esqueleto.hs
index b1f2e44..8e55ca7 100644
--- a/src/Database/Esqueleto.hs
+++ b/src/Database/Esqueleto.hs
@@ -18,8 +18,8 @@ module Database.Esqueleto
   ( -- * @esqueleto@'s Language
     Esqueleto( where_, on, orderBy, asc, desc
              , sub_select, sub_selectDistinct, (^.), (?.)
-             , val, isNothing, just, nothing, not_, (==.), (>=.)
-             , (>.), (<=.), (<.), (!=.), (&&.), (||.)
+             , val, isNothing, just, nothing, countRows, not_
+             , (==.), (>=.), (>.), (<=.), (<.), (!=.), (&&.), (||.)
              , (+.), (-.), (/.), (*.)
              , set, (=.), (+=.), (-=.), (*=.), (/=.) )
   , from
diff --git a/src/Database/Esqueleto/Internal/Language.hs b/src/Database/Esqueleto/Internal/Language.hs
index 7bf2c40..9d1e09a 100644
--- a/src/Database/Esqueleto/Internal/Language.hs
+++ b/src/Database/Esqueleto/Internal/Language.hs
@@ -157,6 +157,9 @@ class (Functor query, Applicative query, Monad query) =>
   -- | @NULL@ value.
   nothing :: expr (Single (Maybe typ))
 
+  -- | @COUNT(*)@ value.
+  countRows :: Num a => expr (Single a)
+
   not_ :: expr (Single Bool) -> expr (Single Bool)
 
   (==.) :: PersistField typ => expr (Single typ) -> expr (Single typ) -> expr (Single Bool)
diff --git a/src/Database/Esqueleto/Internal/Sql.hs b/src/Database/Esqueleto/Internal/Sql.hs
index 458fec4..ebc056a 100644
--- a/src/Database/Esqueleto/Internal/Sql.hs
+++ b/src/Database/Esqueleto/Internal/Sql.hs
@@ -264,7 +264,8 @@ instance Esqueleto SqlQuery SqlExpr SqlPersist where
   isNothing _ = error "Esqueleto/Sql/isNothing: never here (see GHC #6124)"
   just (ERaw p f) = ERaw p f
   just _ = error "Esqueleto/Sql/just: never here (see GHC #6124)"
-  nothing = ERaw Never $ \_ -> ("NULL", mempty)
+  nothing   = ERaw Never $ \_ -> ("NULL",     mempty)
+  countRows = ERaw Never $ \_ -> ("COUNT(*)", mempty)
 
   not_ (ERaw p f) = ERaw Never $ \esc -> let (b, vals) = f esc
                                          in ("NOT " <> parensM p b, vals)
diff --git a/test/Test.hs b/test/Test.hs
index 00b3183..687640e 100644
--- a/test/Test.hs
+++ b/test/Test.hs
@@ -13,6 +13,7 @@
 module Main (main) where
 
 import Control.Applicative (Applicative(..), (<$>))
+import Control.Monad (replicateM_)
 import Control.Monad.Base (MonadBase(..))
 import Control.Monad.IO.Class (MonadIO(liftIO))
 import Control.Monad.Logger (MonadLogger(..), LogLevel(..))
@@ -363,7 +364,7 @@ main = do
                   return p
           liftIO $ ret2 `shouldBe` [ p3e, p2e ]
 
-    describe "update" $
+    describe "update" $ do
       it "works on a simple example" $
         run $ do
           p1k <- insert p1
@@ -382,6 +383,27 @@ main = do
                                   , Entity p1k (Person anon (Just 72))
                                   , Entity p3k p3 ]
 
+      it "works with a subexpression having COUNT(*)" $
+        run $ do
+          p1k <- insert p1
+          p2k <- insert p2
+          p3k <- insert p3
+          replicateM_ 3 (insert $ BlogPost "" p1k)
+          replicateM_ 7 (insert $ BlogPost "" p3k)
+          let blogPostsBy p =
+                from $ \b -> do
+                where_ (b ^. BlogPostAuthorId ==. p ^. PersonId)
+                return countRows
+          ()  <- update $ \p -> do
+                 set p [ PersonAge =. just (sub_select (blogPostsBy p)) ]
+          ret <- select $
+                 from $ \p -> do
+                 orderBy [ asc (p ^. PersonName) ]
+                 return p
+          liftIO $ ret `shouldBe` [ Entity p1k p1 { personAge = Just 3 }
+                                  , Entity p3k p3 { personAge = Just 7 }
+                                  , Entity p2k p2 { personAge = Just 0 } ]
+
 
 ----------------------------------------------------------------------