From fbfdf9a04d9a164e8c0e9924796ec79cbf7f2a86 Mon Sep 17 00:00:00 2001
From: Rakesh Emmadi <12475069+rakeshkky@users.noreply.github.com>
Date: Tue, 18 Jan 2022 20:23:44 +0530
Subject: [PATCH] server/mssql: restrict "count" aggregate query on multiple
 columns

PR-URL: https://github.com/hasura/graphql-engine-mono/pull/3354
GitOrigin-RevId: d9890782ff8e3ea40ec52c0da82e43ecf5c36d2b
---
 CHANGELOG.md                                  | 27 ++++++++++++-----
 .../Backends/BigQuery/Instances/Schema.hs     | 30 +++++++++++--------
 .../src-lib/Hasura/Backends/MSSQL/FromIr.hs   |  4 +--
 .../Hasura/Backends/MSSQL/Instances/Schema.hs | 26 +++++++++-------
 .../src-lib/Hasura/Backends/MSSQL/ToQuery.hs  |  7 ++---
 .../Hasura/Backends/MSSQL/Types/Internal.hs   |  4 +--
 .../Hasura/Backends/MySQL/Instances/Schema.hs |  2 +-
 .../Backends/Postgres/Instances/Schema.hs     | 21 +++++++++----
 .../Hasura/Experimental/Adapter/Schema.hs     |  4 +--
 .../src-lib/Hasura/GraphQL/Schema/Backend.hs  |  7 ++++-
 .../src-lib/Hasura/GraphQL/Schema/Select.hs   | 13 +++++---
 server/src-lib/Hasura/RQL/IR/Select.hs        | 10 +++++++
 .../nodes_aggregates_conditions_mssql.yaml    |  4 +--
 13 files changed, 105 insertions(+), 54 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3740f3823fe..6038587b9ea 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,22 +3,35 @@
 ## Next release
 (Add highlights/major features below)
 
+### Breaking Changes
+- For any **MSSQL** backend, count aggregate query on multiple columns is restricted with a GraphQL
+  schema change as follows
+
+```diff
+count (
+---  columns: [table_select_column!]
++++  column: table_select_column
+  distinct: Boolean
+): Int!
+```
+  MSSQL doesn't support applying `COUNT()` on multiple columns.
+
+
+### Bug fixes and improvements
+(Add entries below in the order of server, console, cli, docs, others)
+
+- server: bigquery: implement `distinct_on`.
 - server: extend transactions to MSSQL GraphQL queries and `mssql_run_sql` /v2/query API
 - server: improve error messages in MSSQL database query exceptions
 - server: in mssql transactions, rollback only if the transaction is active
 - server: add request and response bodies to OpenAPI specification of REST endpoints
 - server: implement upsert mutations for MS SQL Server (close #7864)
-- server: bigquery: implement `distinct_on`.
-- console: action/event trigger transforms are now called REST connectors
-- console: fix list of tables (and schemas) being unsorted when creating a new trigger event (fix #6391)
-
-### Bug fixes and improvements
-(Add entries below in the order of server, console, cli, docs, others)
-
 - server: extend support for insert mutations to tables without primary key constraint in a MSSQL backend
 - server: fix parsing FLOAT64s in scientific notation and non-finite ones in BigQuery
 - server: extend support for the `min`/`max` aggregates to all comparable types in BigQuery
 - server: fix support for joins in aggregates nodes in BigQuery
+- console: action/event trigger transforms are now called REST connectors
+- console: fix list of tables (and schemas) being unsorted when creating a new trigger event (fix #6391)
 - cli: migrate and seed subcommands has an option in prompt to choose and apply operation on all available databases
 - cli: fix `metadata diff --type json | unified-json` behaving incorrectly and showing diff in YAML format.
 - cli: fix regression in `migrate create` command (#7971)
diff --git a/server/src-lib/Hasura/Backends/BigQuery/Instances/Schema.hs b/server/src-lib/Hasura/Backends/BigQuery/Instances/Schema.hs
index fcf7d78aca2..6cf2f478183 100644
--- a/server/src-lib/Hasura/Backends/BigQuery/Instances/Schema.hs
+++ b/server/src-lib/Hasura/Backends/BigQuery/Instances/Schema.hs
@@ -50,7 +50,7 @@ instance BackendSchema 'BigQuery where
   jsonPathArg = bqJsonPathArg
   orderByOperators = bqOrderByOperators
   comparisonExps = bqComparisonExps
-  mkCountType = bqMkCountType
+  countTypeInput = bqCountTypeInput
   aggregateOrderByCountType = BigQuery.IntegerScalarType
   computedField = bqComputedField
   node = bqNode
@@ -352,6 +352,23 @@ bqComparisonExps = P.memoize 'comparisonExps $ \columnType -> do
                    ]
             ]
 
+bqCountTypeInput ::
+  MonadParse n =>
+  Maybe (Parser 'Both n (Column 'BigQuery)) ->
+  InputFieldsParser n (IR.CountDistinct -> CountType 'BigQuery)
+bqCountTypeInput = \case
+  Just columnEnum -> do
+    columns <- P.fieldOptional $$(G.litName "columns") Nothing $ P.list columnEnum
+    pure $ flip mkCountType columns
+  Nothing -> pure $ flip mkCountType Nothing
+  where
+    mkCountType :: IR.CountDistinct -> Maybe [Column 'BigQuery] -> CountType 'BigQuery
+    mkCountType _ Nothing = BigQuery.StarCountable
+    mkCountType IR.SelectCountDistinct (Just cols) =
+      maybe BigQuery.StarCountable BigQuery.DistinctCountable $ nonEmpty cols
+    mkCountType IR.SelectCountNonDistinct (Just cols) =
+      maybe BigQuery.StarCountable BigQuery.NonNullFieldCountable $ nonEmpty cols
+
 geographyWithinDistanceInput ::
   forall m n r.
   (MonadSchema n m, MonadError QErr m, MonadReader r m, Has MkTypename r) =>
@@ -367,17 +384,6 @@ geographyWithinDistanceInput = do
         <*> (mkParameter <$> P.field $$(G.litName "from") Nothing geographyParser)
         <*> (mkParameter <$> P.fieldWithDefault $$(G.litName "use_spheroid") Nothing (G.VBoolean False) booleanParser)
 
-bqMkCountType ::
-  -- | distinct values
-  Maybe Bool ->
-  Maybe [Column 'BigQuery] ->
-  CountType 'BigQuery
-bqMkCountType _ Nothing = BigQuery.StarCountable
-bqMkCountType (Just True) (Just cols) =
-  maybe BigQuery.StarCountable BigQuery.DistinctCountable $ nonEmpty cols
-bqMkCountType _ (Just cols) =
-  maybe BigQuery.StarCountable BigQuery.NonNullFieldCountable $ nonEmpty cols
-
 -- | Computed field parser.
 -- Currently unsupported: returns Nothing for now.
 bqComputedField ::
diff --git a/server/src-lib/Hasura/Backends/MSSQL/FromIr.hs b/server/src-lib/Hasura/Backends/MSSQL/FromIr.hs
index 7a6ffe6eafe..1665bea5133 100644
--- a/server/src-lib/Hasura/Backends/MSSQL/FromIr.hs
+++ b/server/src-lib/Hasura/Backends/MSSQL/FromIr.hs
@@ -654,8 +654,8 @@ fromAggregateField alias aggregateField =
     IR.AFExp text -> AggregateProjection $ Aliased (TextAggregate text) alias
     IR.AFCount countType -> AggregateProjection . flip Aliased alias . CountAggregate $ case countType of
       StarCountable -> StarCountable
-      NonNullFieldCountable names -> NonNullFieldCountable $ fmap columnFieldAggEntity names
-      DistinctCountable names -> DistinctCountable $ fmap columnFieldAggEntity names
+      NonNullFieldCountable name -> NonNullFieldCountable $ columnFieldAggEntity name
+      DistinctCountable name -> DistinctCountable $ columnFieldAggEntity name
     IR.AFOp IR.AggregateOp {_aoOp = op, _aoFields = fields} ->
       let projections :: [Projection] =
             fields <&> \(fieldName, pgColFld) ->
diff --git a/server/src-lib/Hasura/Backends/MSSQL/Instances/Schema.hs b/server/src-lib/Hasura/Backends/MSSQL/Instances/Schema.hs
index a7b4c8e3fab..e605e1f57d5 100644
--- a/server/src-lib/Hasura/Backends/MSSQL/Instances/Schema.hs
+++ b/server/src-lib/Hasura/Backends/MSSQL/Instances/Schema.hs
@@ -63,7 +63,7 @@ instance BackendSchema 'MSSQL where
   jsonPathArg = msJsonPathArg
   orderByOperators = msOrderByOperators
   comparisonExps = msComparisonExps
-  mkCountType = msMkCountType
+  countTypeInput = msCountTypeInput
   aggregateOrderByCountType = MSSQL.IntegerType
   computedField = msComputedField
   node = msNode
@@ -410,16 +410,20 @@ msComparisonExps = P.memoize 'comparisonExps \columnType -> do
     mkListLiteral =
       P.UVLiteral . MSSQL.ListExpression . fmap (MSSQL.ValueExpression . cvValue)
 
-msMkCountType ::
-  -- | distinct values
-  Maybe Bool ->
-  Maybe [Column 'MSSQL] ->
-  CountType 'MSSQL
-msMkCountType _ Nothing = MSSQL.StarCountable
-msMkCountType (Just True) (Just cols) =
-  maybe MSSQL.StarCountable MSSQL.DistinctCountable $ nonEmpty cols
-msMkCountType _ (Just cols) =
-  maybe MSSQL.StarCountable MSSQL.NonNullFieldCountable $ nonEmpty cols
+msCountTypeInput ::
+  MonadParse n =>
+  Maybe (Parser 'Both n (Column 'MSSQL)) ->
+  InputFieldsParser n (IR.CountDistinct -> CountType 'MSSQL)
+msCountTypeInput = \case
+  Just columnEnum -> do
+    column <- P.fieldOptional $$(G.litName "column") Nothing columnEnum
+    pure $ flip mkCountType column
+  Nothing -> pure $ flip mkCountType Nothing
+  where
+    mkCountType :: IR.CountDistinct -> Maybe (Column 'MSSQL) -> CountType 'MSSQL
+    mkCountType _ Nothing = MSSQL.StarCountable
+    mkCountType IR.SelectCountDistinct (Just col) = MSSQL.DistinctCountable col
+    mkCountType IR.SelectCountNonDistinct (Just col) = MSSQL.NonNullFieldCountable col
 
 -- | Computed field parser.
 -- Currently unsupported: returns Nothing for now.
diff --git a/server/src-lib/Hasura/Backends/MSSQL/ToQuery.hs b/server/src-lib/Hasura/Backends/MSSQL/ToQuery.hs
index 6b05afd6ef0..39b4acb4c04 100644
--- a/server/src-lib/Hasura/Backends/MSSQL/ToQuery.hs
+++ b/server/src-lib/Hasura/Backends/MSSQL/ToQuery.hs
@@ -675,11 +675,8 @@ fromCountable :: Countable FieldName -> Printer
 fromCountable =
   \case
     StarCountable -> "*"
-    NonNullFieldCountable fields ->
-      SepByPrinter ", " (map fromFieldName (toList fields))
-    DistinctCountable fields ->
-      "DISTINCT "
-        <+> SepByPrinter ", " (map fromFieldName (toList fields))
+    NonNullFieldCountable field -> fromFieldName field
+    DistinctCountable field -> "DISTINCT " <+> fromFieldName field
 
 fromWhere :: Where -> Printer
 fromWhere =
diff --git a/server/src-lib/Hasura/Backends/MSSQL/Types/Internal.hs b/server/src-lib/Hasura/Backends/MSSQL/Types/Internal.hs
index 3dc045ca91c..eaad654d7cb 100644
--- a/server/src-lib/Hasura/Backends/MSSQL/Types/Internal.hs
+++ b/server/src-lib/Hasura/Backends/MSSQL/Types/Internal.hs
@@ -438,8 +438,8 @@ data Aggregate
 
 data Countable name
   = StarCountable
-  | NonNullFieldCountable (NonEmpty name)
-  | DistinctCountable (NonEmpty name)
+  | NonNullFieldCountable name
+  | DistinctCountable name
 
 deriving instance Functor Countable
 
diff --git a/server/src-lib/Hasura/Backends/MySQL/Instances/Schema.hs b/server/src-lib/Hasura/Backends/MySQL/Instances/Schema.hs
index b0aabb07ba9..0738176e816 100644
--- a/server/src-lib/Hasura/Backends/MySQL/Instances/Schema.hs
+++ b/server/src-lib/Hasura/Backends/MySQL/Instances/Schema.hs
@@ -41,7 +41,7 @@ instance BackendSchema 'MySQL where
   jsonPathArg = jsonPathArg'
   orderByOperators = orderByOperators'
   comparisonExps = comparisonExps'
-  mkCountType = error "mkCountType: MySQL backend does not support this operation yet."
+  countTypeInput = error "countTypeInput: MySQL backend does not support this operation yet."
   aggregateOrderByCountType = error "aggregateOrderByCountType: MySQL backend does not support this operation yet."
   computedField = error "computedField: MySQL backend does not support this operation yet."
   node = error "node: MySQL backend does not support this operation yet."
diff --git a/server/src-lib/Hasura/Backends/Postgres/Instances/Schema.hs b/server/src-lib/Hasura/Backends/Postgres/Instances/Schema.hs
index 0135e02e4d7..7c18a3aff48 100644
--- a/server/src-lib/Hasura/Backends/Postgres/Instances/Schema.hs
+++ b/server/src-lib/Hasura/Backends/Postgres/Instances/Schema.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE ApplicativeDo #-}
 {-# LANGUAGE UndecidableInstances #-}
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 
@@ -149,7 +150,7 @@ instance
   jsonPathArg = jsonPathArg
   orderByOperators = orderByOperators
   comparisonExps = comparisonExps
-  mkCountType = mkCountType
+  countTypeInput = countTypeInput
   aggregateOrderByCountType = PG.PGInteger
   computedField = computedFieldPG
   node = pgkNode
@@ -642,10 +643,20 @@ intersectsGeomNbandInput = do
         <$> (mkParameter <$> P.field $$(G.litName "geommin") Nothing geometryParser)
         <*> (fmap mkParameter <$> P.fieldOptional $$(G.litName "nband") Nothing integerParser)
 
-mkCountType :: Maybe Bool -> Maybe [Column ('Postgres pgKind)] -> CountType ('Postgres pgKind)
-mkCountType _ Nothing = PG.CTStar
-mkCountType (Just True) (Just cols) = PG.CTDistinct cols
-mkCountType _ (Just cols) = PG.CTSimple cols
+countTypeInput ::
+  MonadParse n =>
+  Maybe (Parser 'Both n (Column ('Postgres pgKind))) ->
+  InputFieldsParser n (IR.CountDistinct -> CountType ('Postgres pgKind))
+countTypeInput = \case
+  Just columnEnum -> do
+    columns <- P.fieldOptional $$(G.litName "columns") Nothing (P.list columnEnum)
+    pure $ flip mkCountType columns
+  Nothing -> pure $ flip mkCountType Nothing
+  where
+    mkCountType :: IR.CountDistinct -> Maybe [Column ('Postgres pgKind)] -> CountType ('Postgres pgKind)
+    mkCountType _ Nothing = PG.CTStar
+    mkCountType IR.SelectCountDistinct (Just cols) = PG.CTDistinct cols
+    mkCountType IR.SelectCountNonDistinct (Just cols) = PG.CTSimple cols
 
 -- | Update operator that prepends a value to a column containing jsonb arrays.
 --
diff --git a/server/src-lib/Hasura/Experimental/Adapter/Schema.hs b/server/src-lib/Hasura/Experimental/Adapter/Schema.hs
index db8377b6a4c..5a84918a562 100644
--- a/server/src-lib/Hasura/Experimental/Adapter/Schema.hs
+++ b/server/src-lib/Hasura/Experimental/Adapter/Schema.hs
@@ -49,8 +49,8 @@ instance BackendSchema 'Experimental where
   comparisonExps =
     error "comparisonExps: Unimplemented for Experimental backend."
 
-  mkCountType =
-    error "mkCountType: Unimplemented for Experimental backend."
+  countTypeInput =
+    error "countTypeInput: Unimplemented for Experimental backend."
   aggregateOrderByCountType =
     error "aggregateOrderByCountType: Unimplemented for Experimental backend."
   computedField =
diff --git a/server/src-lib/Hasura/GraphQL/Schema/Backend.hs b/server/src-lib/Hasura/GraphQL/Schema/Backend.hs
index bd2cec40a6a..38da4b25f4d 100644
--- a/server/src-lib/Hasura/GraphQL/Schema/Backend.hs
+++ b/server/src-lib/Hasura/GraphQL/Schema/Backend.hs
@@ -211,7 +211,12 @@ class Backend b => BackendSchema (b :: BackendType) where
     ColumnType b ->
     m (Parser 'Input n [ComparisonExp b])
 
-  mkCountType :: Maybe Bool -> Maybe [Column b] -> CountType b
+  -- | The input fields parser, for "count" aggregate field, yielding a function
+  -- which generates @'CountType b' from optional "distinct" field value
+  countTypeInput ::
+    MonadParse n =>
+    Maybe (Parser 'Both n (Column b)) ->
+    InputFieldsParser n (CountDistinct -> CountType b)
 
   aggregateOrderByCountType :: ScalarType b
 
diff --git a/server/src-lib/Hasura/GraphQL/Schema/Select.hs b/server/src-lib/Hasura/GraphQL/Schema/Select.hs
index ed0f486a06c..9f8c4804661 100644
--- a/server/src-lib/Hasura/GraphQL/Schema/Select.hs
+++ b/server/src-lib/Hasura/GraphQL/Schema/Select.hs
@@ -1016,12 +1016,17 @@ tableAggregationFields sourceName tableInfo selectPermissions = memoizeOn 'table
     countField :: m (FieldParser n (IR.AggregateField b))
     countField = do
       columnsEnum <- tableSelectColumnsEnum sourceName tableInfo selectPermissions
-      let columnsName = $$(G.litName "columns")
-          distinctName = $$(G.litName "distinct")
+      let distinctName = $$(G.litName "distinct")
           args = do
             distinct <- P.fieldOptional distinctName Nothing P.boolean
-            columns <- maybe (pure Nothing) (P.fieldOptional columnsName Nothing . P.list) columnsEnum
-            pure $ mkCountType @b distinct columns
+            mkCountType <- countTypeInput @b columnsEnum
+            pure $
+              mkCountType $
+                maybe
+                  IR.SelectCountNonDistinct -- If "distinct" is "null" or absent, we default to @'SelectCountNonDistinct'
+                  (bool IR.SelectCountNonDistinct IR.SelectCountDistinct)
+                  distinct
+
       pure $ IR.AFCount <$> P.selection $$(G.litName "count") Nothing args P.int
 
     parseAggOperator ::
diff --git a/server/src-lib/Hasura/RQL/IR/Select.hs b/server/src-lib/Hasura/RQL/IR/Select.hs
index 2c2d41d8680..357420d458d 100644
--- a/server/src-lib/Hasura/RQL/IR/Select.hs
+++ b/server/src-lib/Hasura/RQL/IR/Select.hs
@@ -93,6 +93,7 @@ module Hasura.RQL.IR.Select
     TableAggregateFieldsG,
     TablePerm,
     TablePermG (..),
+    CountDistinct (..),
     actionResponsePayloadColumn,
     asnArgs,
     asnFields,
@@ -930,6 +931,15 @@ type ActionFieldsG b r v = Fields (ActionFieldG b r v)
 
 type ActionFields b = ActionFieldsG b Void (SQLExpression b)
 
+-- | The "distinct" input field inside "count" aggregate field
+--
+-- count (
+--   distinct: Boolean
+-- ): Int!
+data CountDistinct
+  = SelectCountDistinct
+  | SelectCountNonDistinct
+
 -- Lenses
 
 $(makeLenses ''AnnSelectG)
diff --git a/server/tests-py/queries/graphql_query/basic/nodes_aggregates_conditions_mssql.yaml b/server/tests-py/queries/graphql_query/basic/nodes_aggregates_conditions_mssql.yaml
index e9c2f8153df..2b388fcdc11 100644
--- a/server/tests-py/queries/graphql_query/basic/nodes_aggregates_conditions_mssql.yaml
+++ b/server/tests-py/queries/graphql_query/basic/nodes_aggregates_conditions_mssql.yaml
@@ -53,7 +53,7 @@
         author {
           articles_aggregate {
             aggregate {
-              count(columns: author_id)
+              count(column: author_id)
               max {
                 id
               }
@@ -92,7 +92,7 @@
         author {
           articles_aggregate {
             aggregate {
-              count(columns: author_id)
+              count(column: author_id)
               max {
                 id
                 author_id