Add support for aggregate functions

The aggregate functions SUM(), MAX(), MIN(), AVG(), and COUNT() are now supported.
PostgREST · Oct 17, 2023 · 78fe42a · 78fe42a
1 parent f10b4c3
commit 78fe42a
Show file tree

Hide file tree

Showing 10 changed files with 464 additions and 98 deletions.
diff --git a/postgrest.cabal b/postgrest.cabal
@@ -202,6 +202,7 @@ test-suite spec
                       Feature.OpenApi.RootSpec
                       Feature.OpenApi.SecurityOpenApiSpec
                       Feature.OptionsSpec
+                      Feature.Query.AggregateFunctionsSpec
                       Feature.Query.AndOrParamsSpec
                       Feature.Query.ComputedRelsSpec
                       Feature.Query.DeleteSpec

diff --git a/src/PostgREST/ApiRequest/QueryParams.hs b/src/PostgREST/ApiRequest/QueryParams.hs
@@ -31,8 +31,8 @@ import Data.Tree                     (Tree (..))
 import Text.Parsec.Error             (errorMessages,
                                       showErrorMessages)
 import Text.ParserCombinators.Parsec (GenParser, ParseError, Parser,
-                                      anyChar, between, char, digit,
-                                      eof, errorPos, letter,
+                                      anyChar, between, char, choice,
+                                      digit, eof, errorPos, letter,
                                       lookAhead, many1, noneOf,
                                       notFollowedBy, oneOf,
                                       optionMaybe, sepBy, sepBy1,
@@ -43,7 +43,8 @@ import PostgREST.RangeQuery              (NonnegRange, allRange,
                                           rangeOffset, restrictRange)
 import PostgREST.SchemaCache.Identifiers (FieldName)
 
-import PostgREST.ApiRequest.Types (EmbedParam (..), EmbedPath, Field,
+import PostgREST.ApiRequest.Types (AggregateFunction (..),
+                                   EmbedParam (..), EmbedPath, Field,
                                    Filter (..), FtsOperator (..),
                                    Hint, JoinType (..),
                                    JsonOperand (..),
@@ -58,7 +59,7 @@ import PostgREST.ApiRequest.Types (EmbedParam (..), EmbedPath, Field,
                                    SimpleOperator (..), SingleVal,
                                    TrileanVal (..))
 
-import Protolude hiding (try)
+import Protolude hiding (Sum, try)
 
 data QueryParams =
   QueryParams
@@ -452,10 +453,12 @@ pRelationSelect :: Parser SelectItem
 pRelationSelect = lexeme $ do
     alias <- optionMaybe ( try(pFieldName <* aliasSeparator) )
     name <- pFieldName
+    guard (name /= "count")
     (hint, jType) <- pEmbedParams
     try (void $ lookAhead (string "("))
     return $ SelectRelation name alias hint jType
 
+
 -- |
 -- Parse regular fields in select
 --
@@ -495,18 +498,36 @@ pFieldSelect :: Parser SelectItem
 pFieldSelect = lexeme $ try (do
     s <- pStar
     pEnd
-    return $ SelectField (s, []) Nothing Nothing)
+    return $ SelectField (s, []) Nothing Nothing Nothing Nothing)
+  <|> try (do
+    alias    <- optionMaybe ( try(pFieldName <* aliasSeparator) )
+    _        <- string "count()"
+    aggCast' <- optionMaybe (string "::" *> pIdentifier)
+    pEnd
+    return $ SelectField ("*", []) (Just Count) (toS <$> aggCast') Nothing alias)
   <|> do
-    alias <- optionMaybe ( try(pFieldName <* aliasSeparator) )
-    fld <- pField
-    cast' <- optionMaybe (string "::" *> pIdentifier)
+    alias    <- optionMaybe ( try(pFieldName <* aliasSeparator) )
+    fld      <- pField
+    cast'    <- optionMaybe (string "::" *> pIdentifier)
+    agg      <- optionMaybe (try (char '.' *> pAggregation <* string "()"))
+    aggCast' <- optionMaybe (string "::" *> pIdentifier)
     pEnd
-    return $ SelectField fld (toS <$> cast') alias
+    return $ SelectField fld agg (toS <$> aggCast') (toS <$> cast') alias
   where
     pEnd = try (void $ lookAhead (string ")")) <|>
            try (void $ lookAhead (string ",")) <|>
            try eof
     pStar = string "*" $> "*"
+    pAggregation = choice
+      [ string "sum"   $> Sum
+      , string "avg"   $> Avg
+      , string "count" $> Count
+      -- Using 'try' for "min" and "max" to allow backtracking.
+      -- This is necessary because both start with the same character 'm',
+      -- and without 'try', a partial match on "max" would prevent "min" from being tried.
+      , try (string "max") $> Max
+      , try (string "min") $> Min
+      ]
 
 
 -- |

diff --git a/src/PostgREST/ApiRequest/Types.hs b/src/PostgREST/ApiRequest/Types.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE DuplicateRecordFields #-}
 module PostgREST.ApiRequest.Types
-  ( Alias
+  ( AggregateFunction(..)
+  , Alias
   , Cast
   , Depth
   , EmbedParam(..)
@@ -42,12 +43,14 @@ import PostgREST.SchemaCache.Routine      (Routine (..))
 
 import Protolude
 
--- | The value in `/tbl?select=alias:field::cast`
+-- | The value in `/tbl?select=alias:field.aggregateFunction()::cast`
 data SelectItem
   = SelectField
-    { selField :: Field
-    , selCast  :: Maybe Cast
-    , selAlias :: Maybe Alias
+    { selField             :: Field
+    , selAggregateFunction :: Maybe AggregateFunction
+    , selAggregateCast     :: Maybe Cast
+    , selCast              :: Maybe Cast
+    , selAlias             :: Maybe Alias
     }
 -- | The value in `/tbl?select=alias:another_tbl(*)`
   | SelectRelation
@@ -135,6 +138,9 @@ type Cast = Text
 type Alias = Text
 type Hint = Text
 
+data AggregateFunction = Sum | Avg | Max | Min | Count
+  deriving (Show, Eq)
+
 data EmbedParam
   -- | Disambiguates an embedding operation when there's multiple relationships
   -- between two tables. Can be the name of a foreign key constraint, column