From 381e50494a6cb298568e4c78216eae30387f5d88 Mon Sep 17 00:00:00 2001 From: Philipp Balzarek Date: Thu, 8 Mar 2018 14:08:43 +0100 Subject: [PATCH] Generalize postgresql aggregate functions --- src/Database/Esqueleto/Internal/Sql.hs | 23 ++- src/Database/Esqueleto/PostgreSQL.hs | 101 +++++++--- test/Common/Test.hs | 1 + test/PostgreSQL/Test.hs | 251 ++++++++++++++++++++++--- 4 files changed, 325 insertions(+), 51 deletions(-) diff --git a/src/Database/Esqueleto/Internal/Sql.hs b/src/Database/Esqueleto/Internal/Sql.hs index bd8a5b9..c68e206 100644 --- a/src/Database/Esqueleto/Internal/Sql.hs +++ b/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]) -makeOrderBy _ [] = mempty -makeOrderBy info os = first ("\nORDER BY " <>) . uncommas' $ concatMap mk os +makeOrderByNoNewline :: + IdentInfo -> [OrderByClause] -> (TLB.Builder, [PersistValue]) +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]) diff --git a/src/Database/Esqueleto/PostgreSQL.hs b/src/Database/Esqueleto/PostgreSQL.hs index b45be5b..67d28da 100644 --- a/src/Database/Esqueleto/PostgreSQL.hs +++ b/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 Database.Esqueleto.Internal.PersistentImport -import Database.Esqueleto.Internal.Sql -import Data.Time.Clock (UTCTime) - --- | (@array_agg@) Concatenate distinct input values, including @NULL@s, into --- an array. --- --- /Since: 2.5.3/ -arrayAggDistinct :: SqlExpr (Value a) -> SqlExpr (Value [a]) -arrayAggDistinct = arrayAgg . distinct' - where - distinct' = unsafeSqlBinOp " " (unsafeSqlValue "DISTINCT") +import Data.Monoid +import qualified Data.Text.Internal.Builder as TLB +import Data.Time.Clock (UTCTime) +import Database.Esqueleto.Internal.Language hiding (random_) +import Database.Esqueleto.Internal.PersistentImport +import Database.Esqueleto.Internal.Sql -- | (@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, --- into an array. +-- | Aggregate mode +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 - :: SqlString s +stringAgg :: + 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.) diff --git a/test/Common/Test.hs b/test/Common/Test.hs index 722f373..06e8426 100644 --- a/test/Common/Test.hs +++ b/test/Common/Test.hs @@ -1,4 +1,5 @@ {-# OPTIONS_GHC -fno-warn-unused-binds #-} +{-# OPTIONS_GHC -fno-warn-deprecations #-} {-# LANGUAGE ConstraintKinds , EmptyDataDecls , FlexibleContexts diff --git a/test/PostgreSQL/Test.hs b/test/PostgreSQL/Test.hs index e9a3c72..294a3ac 100644 --- a/test/PostgreSQL/Test.hs +++ b/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 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 Control.Monad.Trans.Reader (ReaderT, ask) 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 -testPostgresModule = do - describe "PostgreSQL module" $ do - it "arrayAgg 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) +testArrayAggWith :: Spec +testArrayAggWith = do + describe "ALL, no ORDER BY" $ do + it "creates sane SQL" $ run $ do + (query, args) <- showQuery ES.SELECT $ from $ \p -> + return (EP.arrayAggWith EP.AggModeAll (p ^. PersonAge) []) + liftIO $ query `shouldBe` + "SELECT array_agg(\"Person\".\"age\")\n\ + \FROM \"Person\"\n" + 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 = - if verbose - then runVerbose - else runSilent - +run f = do + verbose' <- lookupEnv "VERBOSE" >>= \case + Nothing -> return verbose + 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)