Prevent impure errors with monadic map union

### Description

This very small PR introduces `unionWithM`, to allow hashmap union that might fail, and uses it to transform an `error` into a `throw500`. It also reorders `HashMap.Strict.Extended` to group all "union" functions together.

There is, however, a broader question of whether we should encourage the proliferation of such functions. If so, we might also want to consider:
- `mapWithKeyM`, to remove the `unsafeMkName` of `RemoteJoin.Collect`
- `forWithKey`, as a flipped version of `traverseWithKey`

PR-URL: https://github.com/hasura/graphql-engine-mono/pull/3917
GitOrigin-RevId: a488d5bf04a73878b9e42f27ede36199bb4c920a

server/src-lib/Data/HashMap/Strict/Extended.hs

@@ -2,16 +2,18 @@ module Data.HashMap.Strict.Extended
   ( module M,
     catMaybes,
     fromListOn,
-    unionsAll,
     groupOn,
     groupOnNE,
     differenceOn,
     lpadZip,
-    unionsWith,
     isInverseOf,
+    unionWithM,
+    unionsAll,
+    unionsWith,
   )
 where
 
+import Control.Monad (foldM)
 import Data.Align qualified as A
 import Data.Foldable qualified as F
 import Data.Function (on)
@@ -27,11 +29,6 @@ catMaybes = M.mapMaybe id
 fromListOn :: (Eq k, Hashable k) => (v -> k) -> [v] -> HashMap k v
 fromListOn f = fromList . Prelude.map (\v -> (f v, v))
 
--- | Like 'M.unions', but keeping all elements in the result.
-unionsAll ::
-  (Eq k, Hashable k, Foldable t) => t (HashMap k v) -> HashMap k (NonEmpty v)
-unionsAll = F.foldl' (\a b -> unionWith (<>) a (fmap (:| []) b)) M.empty
-
 -- | Given a 'Foldable' sequence of values and a function that extracts a key from each value,
 -- returns a 'HashMap' that maps each key to a list of all values in the sequence for which the
 -- given function produced it.
@@ -64,20 +61,6 @@ lpadZip left =
       That b -> Just (Nothing, b)
       These a b -> Just (Just a, b)
 
--- | The union of a list of maps, with a combining operation:
---   (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@).
---
--- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
--- >     == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")]
---
--- copied from https://hackage.haskell.org/package/containers-0.6.4.1/docs/src/Data.Map.Internal.html#unionsWith
-unionsWith ::
-  (Foldable f, Hashable k, Ord k) =>
-  (a -> a -> a) ->
-  f (HashMap k a) ->
-  HashMap k a
-unionsWith f ts = F.foldl' (unionWith f) empty ts
-
 -- | Determines whether the left-hand-side and the right-hand-side are inverses of each other.
 --
 -- More specifically, for two maps @A@ and @B@, 'isInverseOf' is satisfied when both of the
@@ -97,3 +80,41 @@ lhs `isInverseOf` rhs = lhs `invertedBy` rhs && rhs `invertedBy` lhs
     a `invertedBy` b = and $ do
       (k, v) <- M.toList a
       pure $ M.lookup v b == Just k
+
+-- | The union of two maps.
+--
+-- If a key occurs in both maps, the provided function (first argument) will be
+-- used to compute the result. Unlike 'unionWith', 'unionWithA' performs the
+-- computation in an arbitratry monad.
+unionWithM ::
+  (Monad m, Eq k, Hashable k) =>
+  (v -> v -> m v) ->
+  HashMap k v ->
+  HashMap k v ->
+  m (HashMap k v)
+unionWithM f m1 m2 = foldM step m1 (toList m2)
+  where
+    step m (k, new) = case M.lookup k m of
+      Nothing -> pure $ insert k new m
+      Just old -> do
+        combined <- f new old
+        pure $ insert k combined m
+
+-- | Like 'M.unions', but keeping all elements in the result.
+unionsAll ::
+  (Eq k, Hashable k, Foldable t) => t (HashMap k v) -> HashMap k (NonEmpty v)
+unionsAll = F.foldl' (\a b -> unionWith (<>) a (fmap (:| []) b)) M.empty
+
+-- | The union of a list of maps, with a combining operation:
+--   (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@).
+--
+-- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
+-- >     == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")]
+--
+-- copied from https://hackage.haskell.org/package/containers-0.6.4.1/docs/src/Data.Map.Internal.html#unionsWith
+unionsWith ::
+  (Foldable f, Hashable k, Ord k) =>
+  (a -> a -> a) ->
+  f (HashMap k a) ->
+  HashMap k a
+unionsWith f ts = F.foldl' (unionWith f) empty ts

server/src-lib/Hasura/GraphQL/Execute/RemoteJoin/RemoteSchema.hs

@@ -25,13 +25,12 @@ import Control.Lens (view, _2, _3)
 import Data.Aeson qualified as A
 import Data.Aeson.Ordered qualified as AO
 import Data.ByteString.Lazy qualified as BL
-import Data.HashMap.Strict qualified as Map
+import Data.HashMap.Strict.Extended qualified as Map
 import Data.IntMap.Strict qualified as IntMap
 import Data.List.NonEmpty qualified as NE
 import Data.Text qualified as T
 import Data.Text.Extended (commaSeparated, toTxt, (<<>))
 import Data.Validation (Validation (..), toEither)
-import GHC.Stack (HasCallStack)
 import Hasura.Base.Error
 import Hasura.GraphQL.Execute.Remote
   ( getVariableDefinitionAndValue,
@@ -159,15 +158,15 @@ fieldCallsToField rrArguments variables finalSelSet topAlias =
         Just f -> do
           s <- nest f
           pure (templatedArguments, [G.SelectionField s])
-        Nothing ->
-          let arguments =
-                Map.unionWith
-                  combineValues
-                  graphQLarguments
-                  -- converting (G.Value Void) -> (G.Value Variable) to merge the
-                  -- 'rrArguments' with the 'variables'
-                  templatedArguments
-           in pure (arguments, finalSelSet)
+        Nothing -> do
+          arguments <-
+            Map.unionWithM
+              combineValues
+              graphQLarguments
+              -- converting (G.Value Void) -> (G.Value Variable) to merge the
+              -- 'rrArguments' with the 'variables'
+              templatedArguments
+          pure (arguments, finalSelSet)
       pure $ G.Field Nothing name args [] selSet
 
     convert :: Map.HashMap G.Name (G.Value Void) -> Map.HashMap G.Name (G.Value RemoteSchemaVariable)
@@ -216,15 +215,13 @@ createArguments variables (RemoteArguments arguments) =
 --
 -- >>> combineValues (Object (fromList [("id", Number 1)]) (Object (fromList [("name", String "foo")])
 -- Object (fromList [("id", Number 1), ("name", String "foo")])
---
--- NOTE: this function *panics* if it fails.
 combineValues ::
-  HasCallStack => G.Value RemoteSchemaVariable -> G.Value RemoteSchemaVariable -> G.Value RemoteSchemaVariable
-combineValues (G.VList l) (G.VList r) = G.VList $ l <> r
-combineValues (G.VObject l) (G.VObject r) = G.VObject $ Map.unionWith combineValues l r
+  MonadError QErr m => G.Value RemoteSchemaVariable -> G.Value RemoteSchemaVariable -> m (G.Value RemoteSchemaVariable)
+combineValues (G.VObject l) (G.VObject r) = G.VObject <$> Map.unionWithM combineValues l r
+combineValues (G.VList l) (G.VList r) = pure $ G.VList $ l <> r
 combineValues l r =
-  error $
-    "combineValues: cannot combine values (" <> show l <> ") and (" <> show r
+  throw500 $
+    "combineValues: cannot combine values (" <> tshow l <> ") and (" <> tshow r
       <> "); \
          \lists can only be merged with lists, objects can only be merged with objects"