Generalize postgresql aggregate functions

src/Database/Esqueleto/Internal/Sql.hs

@@ -17,7 +17,7 @@
 module Database.Esqueleto.Internal.Sql
   ( -- * The pretty face
     SqlQuery
-  , SqlExpr
+  , SqlExpr(..)
   , SqlEntity
   , select
   , selectSource
@@ -35,17 +35,25 @@ module Database.Esqueleto.Internal.Sql
   , unsafeSqlFunction
   , unsafeSqlExtractSubField
   , UnsafeSqlFunctionArgument
+  , OrderByClause
   , rawSelectSource
   , runSource
   , rawEsqueleto
   , toRawSql
   , Mode(..)
+  , NeedParens(..)
   , IdentState
   , initialIdentState
   , IdentInfo
   , SqlSelect(..)
   , veryUnsafeCoerceSqlExprValue
   , veryUnsafeCoerceSqlExprValueList
+  -- * Helper functions
+  , makeOrderByNoNewline
+  , uncommas'
+  , parens
+  , toArgList
+  , builderToText
   ) where
 
 import Control.Arrow ((***), first)
@@ -1145,9 +1153,10 @@ makeHaving info (Where (ERaw _ f))         = first ("\nHAVING " <>) (f info)
 makeHaving _    (Where (ECompositeKey _)) = throw (CompositeKeyErr MakeHavingError)
 
 -- makeHaving, makeWhere and makeOrderBy
-makeOrderBy :: IdentInfo -> [OrderByClause] -> (TLB.Builder, [PersistValue])
+makeOrderByNoNewline ::
-makeOrderBy _    [] = mempty
+     IdentInfo -> [OrderByClause] -> (TLB.Builder, [PersistValue])
-makeOrderBy info os = first ("\nORDER BY " <>) . uncommas' $ concatMap mk os
+makeOrderByNoNewline _    [] = mempty
+makeOrderByNoNewline info os = first ("ORDER BY " <>) . uncommas' $ concatMap mk os
   where
     mk :: OrderByClause -> [(TLB.Builder, [PersistValue])]
     mk (EOrderBy t (ERaw p f)) = [first ((<> orderByType t) . parensM p) (f info)]
@@ -1159,6 +1168,12 @@ makeOrderBy info os = first ("\nORDER BY " <>) . uncommas' $ concatMap mk os
     orderByType ASC  = " ASC"
     orderByType DESC = " DESC"
 
+makeOrderBy :: IdentInfo -> [OrderByClause] -> (TLB.Builder, [PersistValue])
+makeOrderBy _ [] = mempty
+makeOrderBy info is =
+  let (tlb, vals) = makeOrderByNoNewline info is
+  in ("\n" <> tlb, vals)
+
 {-# DEPRECATED EOrderRandom "Since 2.6.0: `rand` ordering function is not uniform across all databases! To avoid accidental partiality it will be removed in the next major version." #-}
 
 makeLimit :: IdentInfo -> LimitClause -> [OrderByClause] -> (TLB.Builder, [PersistValue])

src/Database/Esqueleto/PostgreSQL.hs

@@ -1,32 +1,32 @@
+{-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE OverloadedStrings
+           , GADTs
  #-}
 -- | This module contain PostgreSQL-specific functions.
 --
 -- /Since: 2.2.8/
 module Database.Esqueleto.PostgreSQL
-  ( arrayAggDistinct
+  ( AggMode(..)
+  , arrayAggDistinct
   , arrayAgg
+  , arrayAggWith
   , arrayRemove
   , arrayRemoveNull
   , stringAgg
+  , stringAggWith
   , chr
   , now_
   , random_
+  -- * Internal
+  , unsafeSqlAggregateFunction
   ) where
 
-import Database.Esqueleto.Internal.Language hiding (random_)
+import           Data.Monoid
-import Database.Esqueleto.Internal.PersistentImport
+import qualified Data.Text.Internal.Builder                   as TLB
-import Database.Esqueleto.Internal.Sql
+import           Data.Time.Clock                              (UTCTime)
-import Data.Time.Clock (UTCTime)
+import           Database.Esqueleto.Internal.Language         hiding (random_)
-
+import           Database.Esqueleto.Internal.PersistentImport
--- | (@array_agg@) Concatenate distinct input values, including @NULL@s, into
+import           Database.Esqueleto.Internal.Sql
--- an array.
---
--- /Since: 2.5.3/
-arrayAggDistinct :: SqlExpr (Value a) -> SqlExpr (Value [a])
-arrayAggDistinct = arrayAgg . distinct'
-  where
-    distinct' = unsafeSqlBinOp " " (unsafeSqlValue "DISTINCT")
 
 -- | (@random()@) Split out into database specific modules
 -- because MySQL uses `rand()`.
@@ -35,12 +35,57 @@ arrayAggDistinct = arrayAgg . distinct'
 random_ :: (PersistField a, Num a) => SqlExpr (Value a)
 random_ = unsafeSqlValue "RANDOM()"
 
--- | (@array_agg@) Concatenate input values, including @NULL@s,
+-- | Aggregate mode
--- into an array.
+data AggMode = AggModeAll -- ^ ALL
+             | AggModeDistinct -- ^ DISTINCT
+  deriving (Show)
+
+-- | (Internal) Create a custom aggregate functions with aggregate mode
 --
--- /Since: 2.2.8/
+-- /Do/ /not/ use this function directly, instead define a new function and give
+-- it a type (see `unsafeSqlBinOp`)
+unsafeSqlAggregateFunction ::
+     UnsafeSqlFunctionArgument a
+  => TLB.Builder
+  -> AggMode
+  -> a
+  -> [OrderByClause]
+  -> SqlExpr (Value b)
+unsafeSqlAggregateFunction name mode args orderByClauses =
+  ERaw Never $ \info ->
+    let (orderTLB, orderVals) = makeOrderByNoNewline info orderByClauses
+        -- Don't add a space if we don't have order by clauses
+        orderTLBSpace = case orderByClauses of
+                          [] -> ""
+                          (_:_) -> " "
+        (argsTLB, argsVals) =
+          uncommas' $ map (\(ERaw _ f) -> f info) $ toArgList args
+        aggMode = case mode of
+                    AggModeAll -> "" -- ALL is the default, so we don't need to
+                                     -- specify it
+                    AggModeDistinct -> "DISTINCT "
+    in ( name <> parens (aggMode <> argsTLB <> orderTLBSpace <> orderTLB)
+       , argsVals <> orderVals
+       )
+
+--- | (@array_agg@) Concatenate input values, including @NULL@s,
+--- into an array.
+arrayAggWith ::
+     AggMode -> SqlExpr (Value a) -> [OrderByClause] -> SqlExpr (Value [a])
+arrayAggWith = unsafeSqlAggregateFunction "array_agg"
+
+--- | (@array_agg@) Concatenate input values, including @NULL@s,
+--- into an array.
 arrayAgg :: SqlExpr (Value a) -> SqlExpr (Value [a])
-arrayAgg = unsafeSqlFunction "array_agg"
+arrayAgg x = arrayAggWith AggModeAll x []
+
+-- | (@array_agg@) Concatenate distinct input values, including @NULL@s, into
+-- an array.
+--
+-- /Since: 2.5.3/
+arrayAggDistinct :: SqlExpr (Value a) -> SqlExpr (Value [a])
+arrayAggDistinct x = arrayAggWith AggModeDistinct x []
+
 
 -- | (@array_remove@) Remove all elements equal to the given value from the
 -- array.
@@ -51,20 +96,32 @@ arrayRemove arr elem' = unsafeSqlFunction "array_remove" (arr, elem')
 
 -- | Remove @NULL@ values from an array
 arrayRemoveNull :: SqlExpr (Value [Maybe a]) -> SqlExpr (Value [a])
+-- This can't be a call to arrayRemove because it changes the value type
 arrayRemoveNull x = unsafeSqlFunction "array_remove" (x, unsafeSqlValue "NULL")
 
 
+-- | (@string_agg@) Concatenate input values separated by a
+-- delimiter.
+stringAggWith ::
+     SqlString s
+  => AggMode -- ^ Aggregate mode (ALL or DISTINCT)
+  -> SqlExpr (Value s) -- ^ Input values.
+  -> SqlExpr (Value s) -- ^ Delimiter.
+  -> [OrderByClause] -- ^ ORDER BY clauses
+  -> SqlExpr (Value s) -- ^ Concatenation.
+stringAggWith mode expr delim os =
+  unsafeSqlAggregateFunction "string_agg" mode (expr, delim) os
+
 -- | (@string_agg@) Concatenate input values separated by a
 -- delimiter.
 --
 -- /Since: 2.2.8/
-stringAgg
+stringAgg ::
-  :: SqlString s
+     SqlString s
   => SqlExpr (Value s) -- ^ Input values.
   -> SqlExpr (Value s) -- ^ Delimiter.
   -> SqlExpr (Value s) -- ^ Concatenation.
-stringAgg expr delim = unsafeSqlFunction "string_agg" (expr, delim)
+stringAgg expr delim = stringAggWith AggModeAll expr delim []
-
 
 -- | (@chr@) Translate the given integer to a character. (Note the result will
 -- depend on the character set of your database.)

test/Common/Test.hs

@@ -1,4 +1,5 @@
 {-# OPTIONS_GHC -fno-warn-unused-binds  #-}
+{-# OPTIONS_GHC -fno-warn-deprecations  #-}
 {-# LANGUAGE ConstraintKinds
            , EmptyDataDecls
            , FlexibleContexts

test/PostgreSQL/Test.hs

@@ -4,23 +4,28 @@
            , RankNTypes
            , TypeFamilies
            , OverloadedStrings
+           , LambdaCase
  #-}
 module Main (main) where
 
+import Control.Arrow ((&&&))
 import Control.Monad (void)
 import Control.Monad.IO.Class (MonadIO(liftIO))
 import Control.Monad.Logger (runStderrLoggingT, runNoLoggingT)
-import Control.Monad.Trans.Reader (ReaderT)
+import Control.Monad.Trans.Reader (ReaderT, ask)
-import Database.Esqueleto hiding (random_)
-import Database.Esqueleto.PostgreSQL (random_)
-import Database.Persist.Postgresql (withPostgresqlConn)
-import Data.Ord (comparing)
-import Control.Arrow ((&&&))
-import qualified Database.Esqueleto.PostgreSQL as EP
-import Test.Hspec
 import qualified Control.Monad.Trans.Resource as R
+import qualified Data.Char as Char
 import qualified Data.List as L
+import Data.Ord (comparing)
+import qualified Data.Text as T
 import Data.Time.Clock (getCurrentTime, diffUTCTime)
+import Database.Esqueleto hiding (random_)
+import qualified Database.Esqueleto.Internal.Sql as ES
+import Database.Esqueleto.PostgreSQL (random_)
+import qualified Database.Esqueleto.PostgreSQL   as EP
+import Database.Persist.Postgresql (withPostgresqlConn)
+import System.Environment
+import Test.Hspec
 
 import Common.Test
 
@@ -256,18 +261,193 @@ testArrayRemoveNull = do
 
 
 
-testPostgresModule :: Spec
+testArrayAggWith :: Spec
-testPostgresModule = do
+testArrayAggWith = do
-  describe "PostgreSQL module" $ do
+  describe "ALL, no ORDER BY" $ do
-    it "arrayAgg looks sane" $
+    it "creates sane SQL" $ run $ do
-      run $ do
+      (query, args) <- showQuery ES.SELECT $ from $ \p ->
-        let people = [p1, p2, p3, p4, p5]
+            return (EP.arrayAggWith EP.AggModeAll (p ^. PersonAge) [])
-        mapM_ insert people
+      liftIO $ query `shouldBe`
-        [Value ret] <-
+        "SELECT array_agg(\"Person\".\"age\")\n\
-          select . from $ \p -> return (EP.arrayAgg (p ^. PersonName))
+        \FROM \"Person\"\n"
-        liftIO $ L.sort ret `shouldBe` L.sort (map personName people)
+      liftIO $ args `shouldBe` []
 
-    it "stringAgg looks sane" $
+    it "works on an example" $ run $ do
+      let people = [p1, p2, p3, p4, p5]
+      mapM_ insert people
+      [Value ret] <-
+        select . from $ \p ->
+          return (EP.arrayAggWith EP.AggModeAll (p ^. PersonName) [])
+      liftIO $ L.sort ret `shouldBe` L.sort (map personName people)
+
+  describe "DISTINCT, no ORDER BY" $ do
+    it "creates sane SQL" $ run $ do
+      (query, args) <- showQuery ES.SELECT $ from $ \p ->
+            return (EP.arrayAggWith EP.AggModeDistinct (p ^. PersonAge) [])
+      liftIO $ query `shouldBe`
+        "SELECT array_agg(DISTINCT \"Person\".\"age\")\n\
+        \FROM \"Person\"\n"
+      liftIO $ args `shouldBe` []
+
+    it "works on an example" $ run $ do
+      let people = [p1, p2, p3, p4, p5]
+      mapM_ insert people
+      [Value ret] <-
+        select . from $ \p ->
+          return (EP.arrayAggWith EP.AggModeDistinct (p ^. PersonAge) [])
+      liftIO $ L.sort ret `shouldBe` [Nothing, Just 17, Just 36]
+
+  describe "ALL, ORDER BY" $ do
+    it "creates sane SQL" $ run $ do
+      (query, args) <- showQuery ES.SELECT $ from $ \p ->
+            return (EP.arrayAggWith EP.AggModeAll (p ^. PersonAge)
+                    [ asc $ p ^. PersonName
+                    , desc $ p ^. PersonFavNum
+                    ])
+      liftIO $ query `shouldBe`
+        "SELECT array_agg(\"Person\".\"age\" \
+          \ORDER BY \"Person\".\"name\" ASC, \"Person\".\"favNum\" DESC)\n\
+        \FROM \"Person\"\n"
+      liftIO $ args `shouldBe` []
+
+    it "works on an example" $ run $ do
+      let people = [p1, p2, p3, p4, p5]
+      mapM_ insert people
+      [Value ret] <-
+        select . from $ \p ->
+          return (EP.arrayAggWith EP.AggModeAll (p ^. PersonName) [])
+      liftIO $ L.sort ret `shouldBe` L.sort (map personName people)
+
+  describe "DISTINCT, ORDER BY" $ do
+    it "creates sane SQL" $ run $ do
+      (query, args) <- showQuery ES.SELECT $ from $ \p ->
+          return (EP.arrayAggWith EP.AggModeDistinct (p ^. PersonAge)
+                   [asc $ p ^. PersonAge])
+      liftIO $ query `shouldBe`
+        "SELECT array_agg(DISTINCT \"Person\".\"age\" \
+          \ORDER BY \"Person\".\"age\" ASC)\n\
+        \FROM \"Person\"\n"
+      liftIO $ args `shouldBe` []
+
+    it "works on an example" $ run $ do
+      let people = [p1, p2, p3, p4, p5]
+      mapM_ insert people
+      [Value ret] <-
+        select . from $ \p ->
+          return (EP.arrayAggWith EP.AggModeDistinct (p ^. PersonAge)
+                   [asc $ p ^. PersonAge])
+      liftIO $ ret `shouldBe` [Just 17, Just 36, Nothing]
+
+
+
+
+
+testStringAggWith :: Spec
+testStringAggWith = do
+  describe "ALL, no ORDER BY" $ do
+    it "creates sane SQL" $ run $ do
+      (query, args) <- showQuery ES.SELECT $ from $ \p ->
+            return (EP.stringAggWith EP.AggModeAll (p ^. PersonName)
+                     (val " ") [])
+      liftIO $ query `shouldBe`
+        "SELECT string_agg(\"Person\".\"name\", ?)\n\
+        \FROM \"Person\"\n"
+      liftIO $ args `shouldBe` [PersistText " "]
+
+    it "works on an example" $ run $ do
+      let people = [p1, p2, p3, p4, p5]
+      mapM_ insert people
+      [Value ret] <-
+        select . from $ \p ->
+          return (EP.stringAggWith EP.AggModeAll (p ^. PersonName) (val " ")[])
+      liftIO $ (L.sort $ words ret) `shouldBe` L.sort (map personName people)
+
+    it "works with zero rows" $ run $ do
+      [Value ret] <-
+        select . from $ \p ->
+          return (EP.stringAggWith EP.AggModeAll (p ^. PersonName) (val " ")[])
+      liftIO $ ret `shouldBe` ""
+
+  describe "DISTINCT, no ORDER BY" $ do
+    it "creates sane SQL" $ run $ do
+      (query, args) <- showQuery ES.SELECT $ from $ \p ->
+            return $ EP.stringAggWith EP.AggModeDistinct (p ^. PersonName)
+                     (val " ") []
+      liftIO $ query `shouldBe`
+        "SELECT string_agg(DISTINCT \"Person\".\"name\", ?)\n\
+        \FROM \"Person\"\n"
+      liftIO $ args `shouldBe` [PersistText " "]
+
+    it "works on an example" $ run $ do
+      let people = [p1, p2, p3 {personName = "John"}, p4, p5]
+      mapM_ insert people
+      [Value ret] <-
+        select . from $ \p ->
+          return $ EP.stringAggWith EP.AggModeDistinct (p ^. PersonName) (val " ")
+                   []
+      liftIO $ (L.sort $ words ret) `shouldBe`
+        (L.sort . L.nub $ map personName people)
+
+  describe "ALL, ORDER BY" $ do
+    it "creates sane SQL" $ run $ do
+      (query, args) <- showQuery ES.SELECT $ from $ \p ->
+            return (EP.stringAggWith EP.AggModeAll (p ^. PersonName) (val " ")
+                    [ asc $ p ^. PersonName
+                    , desc $ p ^. PersonFavNum
+                    ])
+      liftIO $ query `shouldBe`
+        "SELECT string_agg(\"Person\".\"name\", ? \
+          \ORDER BY \"Person\".\"name\" ASC, \"Person\".\"favNum\" DESC)\n\
+        \FROM \"Person\"\n"
+      liftIO $ args `shouldBe` [PersistText " "]
+
+    it "works on an example" $ run $ do
+      let people = [p1, p2, p3, p4, p5]
+      mapM_ insert people
+      [Value ret] <-
+        select . from $ \p ->
+          return $ EP.stringAggWith EP.AggModeAll (p ^. PersonName) (val " ")
+                    [desc $ p ^. PersonName]
+      liftIO $ (words ret)
+        `shouldBe` (L.reverse . L.sort $ map personName people)
+
+  describe "DISTINCT, ORDER BY" $ do
+    it "creates sane SQL" $ run $ do
+      (query, args) <- showQuery ES.SELECT $ from $ \p ->
+            return $ EP.stringAggWith EP.AggModeDistinct (p ^. PersonName)
+                     (val " ") [desc $ p ^. PersonName]
+      liftIO $ query `shouldBe`
+        "SELECT string_agg(DISTINCT \"Person\".\"name\", ? \
+        \ORDER BY \"Person\".\"name\" DESC)\n\
+        \FROM \"Person\"\n"
+      liftIO $ args `shouldBe` [PersistText " "]
+
+    it "works on an example" $ run $ do
+      let people = [p1, p2, p3 {personName = "John"}, p4, p5]
+      mapM_ insert people
+      [Value ret] <-
+        select . from $ \p ->
+          return $ EP.stringAggWith EP.AggModeDistinct (p ^. PersonName) (val " ")
+                   [desc $ p ^. PersonName]
+      liftIO $ (words ret) `shouldBe`
+        (L.reverse . L.sort . L.nub $ map personName people)
+
+
+
+
+
+testAggregateFunctions :: Spec
+testAggregateFunctions = do
+  describe "arrayAgg" $ do
+    it "looks sane" $ run $ do
+      let people = [p1, p2, p3, p4, p5]
+      mapM_ insert people
+      [Value ret] <-
+        select . from $ \p -> return (EP.arrayAgg (p ^. PersonName))
+      liftIO $ L.sort ret `shouldBe` L.sort (map personName people)
+  describe "arrayAggWith" testArrayAggWith
+  describe "stringAgg" $ do
+    it "looks sane" $
       run $ do
         let people = [p1, p2, p3, p4, p5]
         mapM_ insert people
@@ -276,7 +456,16 @@ testPostgresModule = do
           from $ \p -> do
           return (EP.stringAgg (p ^. PersonName) (val " "))
         liftIO $ L.sort (words ret) `shouldBe` L.sort (map personName people)
+  describe "stringAggWith" testStringAggWith
 
+
+
+
+
+testPostgresModule :: Spec
+testPostgresModule = do
+  describe "PostgreSQL module" $ do
+    describe "Aggregate functions" testAggregateFunctions
     it "chr looks sane" $
       run $ do
         [Value (ret :: String)] <- select $ return (EP.chr (val 65))
@@ -329,14 +518,18 @@ main = do
 run, runSilent, runVerbose :: Run
 runSilent  act = runNoLoggingT     $ run_worker act
 runVerbose act = runStderrLoggingT $ run_worker act
-run =
+run f = do
-  if verbose
+  verbose' <- lookupEnv "VERBOSE" >>= \case
-  then runVerbose
+    Nothing -> return verbose
-  else runSilent
+    Just x | map Char.toLower x == "true" -> return True
-
+           | null x -> return True
+           | otherwise -> return False
+  if verbose'
+    then runVerbose f
+    else runSilent f
 
 verbose :: Bool
-verbose = True
+verbose = False
 
 migrateIt :: RunDbMonad m => SqlPersistT (R.ResourceT m) ()
 migrateIt = do
@@ -349,3 +542,11 @@ run_worker act = withConn $ runSqlConn (migrateIt >> act)
 withConn :: RunDbMonad m => (SqlBackend -> R.ResourceT m a) -> m a
 withConn =
   R.runResourceT . withPostgresqlConn "host=localhost port=5432 user=esqutest password=esqutest dbname=esqutest"
+
+-- | Show the SQL generated by a query
+showQuery :: (Monad m, ES.SqlSelect a r, BackendCompatible SqlBackend backend)
+          => ES.Mode -> SqlQuery a -> ReaderT backend m (T.Text, [PersistValue])
+showQuery mode query = do
+  backend <- ask
+  let (builder, values) = ES.toRawSql mode (backend, ES.initialIdentState) query
+  return (ES.builderToText builder, values)