Changed RemoteSchemaIntrospection's internal representation from a list to a hashmap.

### Description

This PR changes the internal representation of a parsed remote schema. We were still using a list of type definitions, meaning every time we were doing a type lookup we had to iterate through a linked list! 🙀 It was very noticeable on large schemas, that need to do a lot of lookups. This PR consequently changes the internal representation to a HashMap. Building the OneGraph schema on my machine now takes **23 seconds**, compared to **367 seconds** before this patch.

Some important points:
- ~~this PR removes a check for type duplication in remote schemas; it's unclear to me whether that's something we need to add back or not~~ (no longer true)
- this PR makes it obvious that we do not distinguish between "this remote schema is missing type X" and "this remote schema expects type X to be an object, but it's a scalar"; this PR doesn't change anything about it, but adds a comment where we could surface that error (see [2991](https://github.com/hasura/graphql-engine-mono/issues/2991))

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

CHANGELOG.md

@@ -8,6 +8,7 @@
 - server: log locking DB queries during source catalog migration
 - server: fix to allow remote schema response to contain an "extensions" field (#7143)
 - server: support database-to-database joins with BigQuery
+- server: improved startup time when using large remote schemas
 - console: add comments to tracked functions
 - console: add select all columns option while selecting the columns in event triggers
 - console: add request transforms for events

server/src-lib/Hasura/GraphQL/RemoteServer.hs

@@ -19,7 +19,7 @@ import Data.Aeson.Types qualified as J
 import Data.ByteString.Lazy qualified as BL
 import Data.Environment qualified as Env
 import Data.FileEmbed (makeRelativeToProject)
-import Data.HashMap.Strict qualified as Map
+import Data.HashMap.Strict.Extended qualified as Map
 import Data.HashSet qualified as Set
 import Data.List.Extended (duplicates)
 import Data.Text qualified as T
@@ -113,7 +113,7 @@ validateSchemaCustomizationsDistinct remoteSchemaCustomizer (RemoteSchemaIntrosp
 
     validateTypeMappingsAreDistinct :: m ()
     validateTypeMappingsAreDistinct = do
-      let dups = duplicates $ (runMkTypename customizeTypeName . typeDefinitionName) <$> typeDefinitions
+      let dups = duplicates $ runMkTypename customizeTypeName <$> Map.keys typeDefinitions
       unless (Set.null dups) $
         throwRemoteSchema $
           "Type name mappings are not distinct; the following types appear more than once: "
@@ -392,7 +392,7 @@ instance J.FromJSON (FromIntrospection IntrospectionResult) where
             types
         r =
           IntrospectionResult
-            (RemoteSchemaIntrospection (fmap fromIntrospection types'))
+            (RemoteSchemaIntrospection $ Map.fromListOn getTypeName $ fromIntrospection <$> types')
             queryRoot
             mutationRoot
             subsRoot
@@ -456,15 +456,6 @@ execRemoteGQ env manager userInfo reqHdrs rsdef gqlReq@GQLReq {..} = do
 identityCustomizer :: RemoteSchemaCustomizer
 identityCustomizer = RemoteSchemaCustomizer Nothing mempty mempty
 
-typeDefinitionName :: G.TypeDefinition a b -> G.Name
-typeDefinitionName = \case
-  G.TypeDefinitionScalar G.ScalarTypeDefinition {..} -> _stdName
-  G.TypeDefinitionObject G.ObjectTypeDefinition {..} -> _otdName
-  G.TypeDefinitionInterface G.InterfaceTypeDefinition {..} -> _itdName
-  G.TypeDefinitionUnion G.UnionTypeDefinition {..} -> _utdName
-  G.TypeDefinitionEnum G.EnumTypeDefinition {..} -> _etdName
-  G.TypeDefinitionInputObject G.InputObjectTypeDefinition {..} -> _iotdName
-
 getCustomizer :: IntrospectionResult -> Maybe RemoteSchemaCustomization -> RemoteSchemaCustomizer
 getCustomizer _ Nothing = identityCustomizer
 getCustomizer IntrospectionResult {..} (Just RemoteSchemaCustomization {..}) = RemoteSchemaCustomizer {..}
@@ -487,7 +478,7 @@ getCustomizer IntrospectionResult {..} (Just RemoteSchemaCustomization {..}) = R
       (Just prefix, Just suffix) -> map (\name -> (name, prefix <> name <> suffix)) names
 
     RemoteSchemaIntrospection typeDefinitions = irDoc
-    typesToRename = filter nameFilter $ typeDefinitionName <$> typeDefinitions
+    typesToRename = filter nameFilter $ Map.keys typeDefinitions
     typeRenameMap =
       case _rscTypeNames of
         Nothing -> Map.empty
@@ -496,11 +487,12 @@ getCustomizer IntrospectionResult {..} (Just RemoteSchemaCustomization {..}) = R
 
     typeFieldMap :: HashMap G.Name [G.Name] -- typeName -> fieldNames
     typeFieldMap =
-      Map.fromList $
-        typeDefinitions >>= \case
-          G.TypeDefinitionObject G.ObjectTypeDefinition {..} -> pure (_otdName, G._fldName <$> _otdFieldsDefinition)
-          G.TypeDefinitionInterface G.InterfaceTypeDefinition {..} -> pure (_itdName, G._fldName <$> _itdFieldsDefinition)
-          _ -> []
+      Map.mapMaybe getFieldsNames typeDefinitions
+      where
+        getFieldsNames = \case
+          G.TypeDefinitionObject G.ObjectTypeDefinition {..} -> Just $ G._fldName <$> _otdFieldsDefinition
+          G.TypeDefinitionInterface G.InterfaceTypeDefinition {..} -> Just $ G._fldName <$> _itdFieldsDefinition
+          _ -> Nothing
 
     mkFieldRenameMap RemoteFieldCustomization {..} fieldNames =
       _rfcMapping <> mkPrefixSuffixMap _rfcPrefix _rfcSuffix fieldNames

server/src-lib/Hasura/GraphQL/Schema/Select.hs

@@ -35,7 +35,7 @@ import Data.Aeson.Internal qualified as J
 import Data.Align (align)
 import Data.ByteString.Lazy qualified as BL
 import Data.Has
-import Data.HashMap.Strict qualified as Map
+import Data.HashMap.Strict.Extended qualified as Map
 import Data.HashSet qualified as Set
 import Data.Int (Int64)
 import Data.List.NonEmpty qualified as NE
@@ -1357,7 +1357,7 @@ remoteRelationshipField remoteFieldInfo = runMaybeT do
       let roleIntrospection@(RemoteSchemaIntrospection typeDefns) = irDoc roleIntrospectionResultOriginal
           -- add the new input value definitions created by the remote relationship
           -- to the existing schema introspection of the role
-          remoteRelationshipIntrospection = RemoteSchemaIntrospection $ typeDefns <> newInpValDefns
+          remoteRelationshipIntrospection = RemoteSchemaIntrospection $ typeDefns <> Map.fromListOn getTypeName newInpValDefns
       fieldName <- textToName $ remoteRelationshipNameToText name
 
       -- This selection set parser, should be of the remote node's selection set parser, which comes

server/src-lib/Hasura/RQL/DDL/RemoteSchema/Permission.hs

@@ -38,7 +38,7 @@ module Hasura.RQL.DDL.RemoteSchema.Permission
 where
 
 import Control.Monad.Validate
-import Data.HashMap.Strict qualified as Map
+import Data.HashMap.Strict.Extended qualified as Map
 import Data.HashSet qualified as S
 import Data.List.Extended (duplicates, getDifference)
 import Data.List.NonEmpty qualified as NE
@@ -900,7 +900,7 @@ getSchemaDocIntrospection providedTypeDefns (queryRoot, mutationRoot, subscripti
           G.TypeDefinitionObject obj -> pure $ G.TypeDefinitionObject obj
           G.TypeDefinitionUnion union' -> pure $ G.TypeDefinitionUnion union'
           G.TypeDefinitionInputObject inpObj -> pure $ G.TypeDefinitionInputObject inpObj
-      remoteSchemaIntrospection = RemoteSchemaIntrospection modifiedTypeDefns
+      remoteSchemaIntrospection = RemoteSchemaIntrospection $ Map.fromListOn getTypeName modifiedTypeDefns
    in IntrospectionResult remoteSchemaIntrospection (fromMaybe $$(G.litName "Query") queryRoot) mutationRoot subscriptionRoot
 
 -- | validateRemoteSchema accepts two arguments, the `SchemaDocument` of
@@ -966,14 +966,6 @@ validateRemoteSchema upstreamRemoteSchemaIntrospection = do
           <$> validateInputObjectTypeDefinition providedInputObjectTypeDefn upstreamInputObjectTypeDefn
   pure $ getSchemaDocIntrospection validatedTypeDefinitions rootTypeNames
   where
-    getTypeName = \case
-      G.TypeDefinitionScalar scalar -> G._stdName scalar
-      G.TypeDefinitionEnum enum -> G._etdName enum
-      G.TypeDefinitionObject obj -> G._otdName obj
-      G.TypeDefinitionUnion union' -> G._utdName union'
-      G.TypeDefinitionInterface iface -> G._itdName iface
-      G.TypeDefinitionInputObject inpObj -> G._iotdName inpObj
-
     typeNotFound gType name = refute (pure $ TypeDoesNotExist gType name)
 
 resolveRoleBasedRemoteSchema ::

server/src-lib/Hasura/RQL/Types/RemoteSchema.hs

@@ -34,6 +34,7 @@ module Hasura.RQL.Types.RemoteSchema
     lookupUnion,
     modifyFieldByName,
     remoteSchemaCustomizeFieldName,
+    getTypeName,
     remoteSchemaCustomizeTypeName,
     rfField,
     rfRemoteSchemaInfo,
@@ -408,7 +409,7 @@ instance Hashable RemoteSchemaInputValueDefinition
 instance Cacheable RemoteSchemaInputValueDefinition
 
 newtype RemoteSchemaIntrospection
-  = RemoteSchemaIntrospection [(G.TypeDefinition [G.Name] RemoteSchemaInputValueDefinition)]
+  = RemoteSchemaIntrospection (HashMap G.Name (G.TypeDefinition [G.Name] RemoteSchemaInputValueDefinition))
   deriving (Show, Eq, Generic, Hashable, Cacheable, Ord)
 
 data RemoteFieldG var = RemoteFieldG
@@ -437,14 +438,10 @@ instance J.ToJSON RemoteSchemaPermsCtx where
     RemoteSchemaPermsEnabled -> J.Bool True
     RemoteSchemaPermsDisabled -> J.Bool False
 
-lookupType ::
-  RemoteSchemaIntrospection ->
-  G.Name ->
-  Maybe (G.TypeDefinition [G.Name] RemoteSchemaInputValueDefinition)
-lookupType (RemoteSchemaIntrospection types) name = find (\tp -> getNamedTyp tp == name) types
-  where
-    getNamedTyp :: G.TypeDefinition possibleTypes RemoteSchemaInputValueDefinition -> G.Name
-    getNamedTyp ty = case ty of
+-- | Extracts the name of a given type from its definition.
+-- TODO: move this to Language.GraphQL.Draft.Syntax.
+getTypeName :: G.TypeDefinition possibleTypes inputType -> G.Name
+getTypeName = \case
   G.TypeDefinitionScalar t -> G._stdName t
   G.TypeDefinitionObject t -> G._otdName t
   G.TypeDefinitionInterface t -> G._itdName t
@@ -452,23 +449,33 @@ lookupType (RemoteSchemaIntrospection types) name = find (\tp -> getNamedTyp tp
   G.TypeDefinitionEnum t -> G._etdName t
   G.TypeDefinitionInputObject t -> G._iotdName t
 
+lookupType ::
+  RemoteSchemaIntrospection ->
+  G.Name ->
+  Maybe (G.TypeDefinition [G.Name] RemoteSchemaInputValueDefinition)
+lookupType (RemoteSchemaIntrospection types) name = Map.lookup name types
+
 lookupObject ::
   RemoteSchemaIntrospection ->
   G.Name ->
   Maybe (G.ObjectTypeDefinition RemoteSchemaInputValueDefinition)
-lookupObject (RemoteSchemaIntrospection types) name =
-  choice $
-    types <&> \case
+lookupObject introspection name =
+  lookupType introspection name >>= \case
     G.TypeDefinitionObject t | G._otdName t == name -> Just t
+    -- if this happens, it means the schema is inconsistent: we expected to
+    -- find an object with that name, but instead found something that wasn't
+    -- an object; we might want to indicate this with a proper failure, so we
+    -- can show better diagnostics to the user?
+    -- This also applies to all following functions.
+    -- See: https://github.com/hasura/graphql-engine-mono/issues/2991
     _ -> Nothing
 
 lookupInterface ::
   RemoteSchemaIntrospection ->
   G.Name ->
   Maybe (G.InterfaceTypeDefinition [G.Name] RemoteSchemaInputValueDefinition)
-lookupInterface (RemoteSchemaIntrospection types) name =
-  choice $
-    types <&> \case
+lookupInterface introspection name =
+  lookupType introspection name >>= \case
     G.TypeDefinitionInterface t | G._itdName t == name -> Just t
     _ -> Nothing
 
@@ -476,9 +483,8 @@ lookupScalar ::
   RemoteSchemaIntrospection ->
   G.Name ->
   Maybe G.ScalarTypeDefinition
-lookupScalar (RemoteSchemaIntrospection types) name =
-  choice $
-    types <&> \case
+lookupScalar introspection name =
+  lookupType introspection name >>= \case
     G.TypeDefinitionScalar t | G._stdName t == name -> Just t
     _ -> Nothing
 
@@ -486,9 +492,8 @@ lookupUnion ::
   RemoteSchemaIntrospection ->
   G.Name ->
   Maybe G.UnionTypeDefinition
-lookupUnion (RemoteSchemaIntrospection types) name =
-  choice $
-    types <&> \case
+lookupUnion introspection name =
+  lookupType introspection name >>= \case
     G.TypeDefinitionUnion t | G._utdName t == name -> Just t
     _ -> Nothing
 
@@ -496,9 +501,8 @@ lookupEnum ::
   RemoteSchemaIntrospection ->
   G.Name ->
   Maybe G.EnumTypeDefinition
-lookupEnum (RemoteSchemaIntrospection types) name =
-  choice $
-    types <&> \case
+lookupEnum introspection name =
+  lookupType introspection name >>= \case
     G.TypeDefinitionEnum t | G._etdName t == name -> Just t
     _ -> Nothing
 
@@ -506,8 +510,7 @@ lookupInputObject ::
   RemoteSchemaIntrospection ->
   G.Name ->
   Maybe (G.InputObjectTypeDefinition RemoteSchemaInputValueDefinition)
-lookupInputObject (RemoteSchemaIntrospection types) name =
-  choice $
-    types <&> \case
+lookupInputObject introspection name =
+  lookupType introspection name >>= \case
     G.TypeDefinitionInputObject t | G._iotdName t == name -> Just t
     _ -> Nothing

server/src-lib/Hasura/RQL/Types/Roles/Internal.hs

@@ -8,6 +8,7 @@ module Hasura.RQL.Types.Roles.Internal
   )
 where
 
+import Data.HashMap.Strict qualified as Map
 import Data.HashSet qualified as Set
 import Data.Semigroup (Any (..), Max (..))
 import Hasura.Prelude
@@ -197,12 +198,12 @@ instance OnlyRelevantEq RemoteSchemaInputValueDefinition where
 instance OnlyRelevantEq RemoteSchemaIntrospection where
   RemoteSchemaIntrospection typeDefinitionsL
     `relevantEq` RemoteSchemaIntrospection typeDefinitionsR =
-      (sort typeDefinitionsL) `relevantEq` (sort typeDefinitionsR)
+      sort (Map.elems typeDefinitionsL) `relevantEq` sort (Map.elems typeDefinitionsR)
 
 instance OnlyRelevantEq IntrospectionResult where
   IntrospectionResult (RemoteSchemaIntrospection typeDefnsL) queryRootL mutationRootL subsRootL
     `relevantEq` IntrospectionResult (RemoteSchemaIntrospection typeDefnsR) queryRootR mutationRootR subsRootR =
-      (sort typeDefnsL) `relevantEq` (sort typeDefnsR)
+      sort (Map.elems typeDefnsL) `relevantEq` sort (Map.elems typeDefnsR)
         && queryRootL == queryRootR
         && mutationRootL == mutationRootR
         && subsRootL == subsRootR

server/src-test/Hasura/Generator.hs

@@ -3,7 +3,7 @@
 module Hasura.Generator () where
 
 import Data.Containers.ListUtils (nubOrd)
-import Data.HashMap.Strict qualified as Map
+import Data.HashMap.Strict.Extended qualified as Map
 import Data.HashMap.Strict.InsOrd qualified as OMap
 import Data.Ratio ((%))
 import Data.Text qualified as T
@@ -136,6 +136,7 @@ instance Arbitrary IntrospectionResult where
     -- finally, create an IntrospectionResult from the aggregated definitions
     let irDoc =
           RemoteSchemaIntrospection $
+            Map.fromListOn getTypeName $
               concat
                 [ G.TypeDefinitionScalar <$> scalarTypeDefinitions,
                   G.TypeDefinitionObject <$> objectTypeDefinitions,

server/src-test/Hasura/GraphQL/Schema/RemoteTest.hs

@@ -3,7 +3,7 @@ module Hasura.GraphQL.Schema.RemoteTest (spec) where
 import Control.Lens (Prism', prism', to, (^..), _Right)
 import Data.Aeson qualified as J
 import Data.ByteString.Lazy qualified as LBS
-import Data.HashMap.Strict qualified as M
+import Data.HashMap.Strict.Extended qualified as M
 import Data.Text qualified as T
 import Data.Text.Extended
 import Data.Text.RawString
@@ -33,6 +33,7 @@ runError = runExceptT >=> (`onLeft` (error . T.unpack . qeError))
 
 mkTestRemoteSchema :: Text -> RemoteSchemaIntrospection
 mkTestRemoteSchema schema = RemoteSchemaIntrospection $
+  M.fromListOn getTypeName $
     runIdentity $
       runError $ do
         G.SchemaDocument types <- G.parseSchemaDocument schema `onLeft` throw500