server: add connection_template as an argument in pg_test_connection_template API

PR-URL: https://github.com/hasura/graphql-engine-mono/pull/8520
GitOrigin-RevId: 525bba9015ad4e143e94124e42ec4518252932cb

docs/docs/api-reference/metadata-api/table-view.mdx

@@ -301,6 +301,9 @@ X-Hasura-Role: admin
         "operation_type": "query",
         "operation_name": "users_by_pk"
       }
+    },
+    "connection_template": {
+        "template": "{{ if $.request.session?[\"x-hasura-role\"] == \"user\" }} {{$.connection_set.connection_set_member_name}} {{else}} {{$.primary}} {{ end }}"
       }
   }
 }
@@ -309,9 +312,10 @@ X-Hasura-Role: admin
 ### Args syntax {#metadata-pg-test-connection-template-syntax}
 
 | Key                 | Required  | Schema                                                                      | Description                                                             |
-| --------------- | -------- | --------------------------------------------------------------- | ------------------------------------------------------------- |
+| ------------------- | --------- | --------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
 | source_name         | false     | [SourceName](/api-reference/syntax-defs.mdx#sourcename)                     | Name of the source database of the table (default: `default`)           |
 | request_context     | true      | [RequestContext](/api-reference/syntax-defs.mdx#requestcontext)             | Request context                                                         |
+| connection_template | false     | [PGConnectionTemplate](/api-reference/syntax-defs.mdx#pgconnectiontemplate) | DB connection template                                                  |
 
 :::info Enterprise only
 

docs/docs/databases/database-config/dynamic-db-connection.mdx

@@ -354,7 +354,7 @@ X-Hasura-Role: admin
 {
   "type": "pg_test_connection_template",
   "args": {
-    "source": "source_name",
+    "source_name": "source_name",
     "request_context": {
       "headers": {
         "header_name": "header_value"
@@ -366,6 +366,9 @@ X-Hasura-Role: admin
         "operation_type": "query",
         "operation_name": "op_name"
       }
+    },
+    "connection_template": {
+        "template": "{{ if $.request.session?[\"x-hasura-role\"] == \"user\" }} {{$.primary}} {{else}} {{$.connection_set.db_1}} {{ end }}"
       }
   }
 }
@@ -382,6 +385,14 @@ Success Response:
 }
 ```
 
+:::info Note
+
+`connection_template` is an optional argument which has precedence over the connection template present in the source. If
+`connection_template` is provided in the API, then the source's connection template will be ignored.
+
+:::
+
+
 ## Limitations
 
 ### Postgres schema of connection set

server/src-lib/Hasura/Backends/Postgres/Execute/Types.hs

@@ -23,6 +23,7 @@ module Hasura.Backends.Postgres.Execute.Types
     runPgSourceWriteTx,
     applyConnectionTemplateResolverNonAdmin,
     pgResolveConnectionTemplate,
+    resolvePostgresConnectionTemplate,
   )
 where
 
@@ -33,7 +34,7 @@ import Data.CaseInsensitive qualified as CI
 import Data.HashMap.Internal.Strict qualified as Map
 import Database.PG.Query qualified as PG
 import Database.PG.Query.Connection qualified as PG
-import Hasura.Backends.Postgres.Connection.Settings (PostgresConnectionSetMemberName)
+import Hasura.Backends.Postgres.Connection.Settings (ConnectionTemplate (..), PostgresConnectionSetMemberName)
 import Hasura.Backends.Postgres.Execute.ConnectionTemplate
 import Hasura.Backends.Postgres.SQL.Error
 import Hasura.Base.Error
@@ -42,6 +43,7 @@ import Hasura.Prelude
 import Hasura.RQL.Types.ResizePool
 import Hasura.SQL.Types (ExtensionsSchema)
 import Hasura.Session
+import Kriti.Error qualified as Kriti
 import Network.HTTP.Types qualified as HTTP
 
 -- See Note [Existentially Quantified Types]
@@ -270,20 +272,48 @@ applyConnectionTemplateResolver connectionTemplateResolver sessionVariables requ
   for connectionTemplateResolver $ \resolver ->
     _runResolver resolver sessionVariables requestHeaders queryContext
 
-pgResolveConnectionTemplate :: (MonadError QErr m) => PGSourceConfig -> RequestContext -> m EncJSON
+pgResolveConnectionTemplate :: (MonadError QErr m) => PGSourceConfig -> RequestContext -> Maybe ConnectionTemplate -> m EncJSON
-pgResolveConnectionTemplate sourceConfig (RequestContext (RequestContextHeaders headersMap) sessionVariables queryContext) = do
+pgResolveConnectionTemplate sourceConfig (RequestContext (RequestContextHeaders headersMap) sessionVariables queryContext) connectionTemplateMaybe = do
-  let headers = map (\(hName, hVal) -> (CI.mk (txtToBs hName), txtToBs hVal)) $ Map.toList headersMap
   connectionTemplateResolver <-
+    case connectionTemplateMaybe of
+      Nothing ->
         case _pscConnectionTemplateConfig sourceConfig of
-      ConnTemplate_NotApplicable ->
+          ConnTemplate_NotApplicable -> connectionTemplateNotApplicableError
-        throw400 NotSupported "Connection templating feature is enterprise edition only"
           ConnTemplate_NotConfigured ->
             throw400 TemplateResolutionFailed "Connection template not defined for the source"
           ConnTemplate_Resolver resolver ->
             pure resolver
+      Just connectionTemplate ->
+        case _pscConnectionTemplateConfig sourceConfig of
+          -- connection template is an enterprise edition only feature. `ConnTemplate_NotApplicable` error is thrown
+          -- when community edition engine is used to test the connection template
+          ConnTemplate_NotApplicable -> connectionTemplateNotApplicableError
+          _ -> pure $ ConnectionTemplateResolver $ \sessionVariables' reqHeaders queryContext' ->
+            resolvePostgresConnectionTemplate connectionTemplate (Map.keys (_pscConnectionSet sourceConfig)) sessionVariables' reqHeaders queryContext'
+  let headers = map (\(hName, hVal) -> (CI.mk (txtToBs hName), txtToBs hVal)) $ Map.toList headersMap
   case maybeRoleFromSessionVariables sessionVariables of
     Nothing -> throw400 InvalidParams "No `x-hasura-role` found in session variables. Please try again with non-admin 'x-hasura-role' in the session context."
     Just roleName ->
       when (roleName == adminRoleName) $ throw400 InvalidParams "Only requests made with a non-admin context can resolve the connection template. Please try again with non-admin 'x-hasura-role' in the session context."
   resolvedTemplate <- _runResolver connectionTemplateResolver sessionVariables headers queryContext
   pure . encJFromJValue $ J.object ["result" J..= resolvedTemplate]
+  where
+    connectionTemplateNotApplicableError = throw400 NotSupported "Connection templating feature is enterprise edition only"
+
+resolvePostgresConnectionTemplate ::
+  (MonadError QErr m) =>
+  ConnectionTemplate ->
+  [PostgresConnectionSetMemberName] ->
+  SessionVariables ->
+  [HTTP.Header] ->
+  Maybe QueryContext ->
+  m (PostgresResolvedConnectionTemplate)
+resolvePostgresConnectionTemplate (ConnectionTemplate _templateSrc connectionTemplate) connectionSetMembers sessionVariables reqHeaders queryContext = do
+  let requestContext = makeRequestContext queryContext reqHeaders sessionVariables
+      connectionTemplateCtx = makeConnectionTemplateContext requestContext connectionSetMembers
+
+  case runKritiEval connectionTemplateCtx connectionTemplate of
+    Left err ->
+      let serializedErr = Kriti.serialize err
+       in throw400WithDetail TemplateResolutionFailed ("Connection template evaluation failed: " <> Kriti._message serializedErr) (J.toJSON $ serializedErr)
+    Right val -> runAesonParser (J.parseJSON @PostgresResolvedConnectionTemplate) val

server/src-lib/Hasura/RQL/DDL/ConnectionTemplate.hs

@@ -12,6 +12,7 @@ where
 
 import Data.Aeson (FromJSON (parseJSON), ToJSON (toJSON))
 import Data.Aeson qualified as J
+import Hasura.Backends.Postgres.Connection.Settings (ConnectionTemplate (..))
 import Hasura.Base.Error
 import Hasura.EncJSON
 import Hasura.Prelude
@@ -24,7 +25,8 @@ import Hasura.SQL.AnyBackend qualified as AB
 -- | The input type for the metadata API `<backend>_test_connection_template`
 data TestConnectionTemplate b = TestConnectionTemplate
   { _tctSourceName :: SourceName,
-    _tctRequestContext :: ConnectionTemplateRequestContext b
+    _tctRequestContext :: ConnectionTemplateRequestContext b,
+    _tctConnectionTemplate :: (Maybe ConnectionTemplate)
   }
 
 instance (Backend b) => FromJSON (TestConnectionTemplate b) where
@@ -33,6 +35,7 @@ instance (Backend b) => FromJSON (TestConnectionTemplate b) where
       TestConnectionTemplate
         <$> o J..:? "source_name" J..!= defaultSource
         <*> o J..: "request_context"
+        <*> o J..:? "connection_template"
 
 -- | Resolver for the metadata API `<backend>_test_connection_template`
 runTestConnectionTemplate ::
@@ -40,9 +43,9 @@ runTestConnectionTemplate ::
   (MonadError QErr m, CacheRM m, Backend b, MetadataM m) =>
   TestConnectionTemplate b ->
   m EncJSON
-runTestConnectionTemplate (TestConnectionTemplate sourceName requestContext) = do
+runTestConnectionTemplate (TestConnectionTemplate sourceName requestContext connectionTemplateMaybe) = do
   sourceConfig <- askSourceConfig @b sourceName
-  liftEither $ resolveConnectionTemplate @b sourceConfig requestContext
+  liftEither $ resolveConnectionTemplate @b sourceConfig requestContext connectionTemplateMaybe
 
 -- A wrapper around the `ResolvedConnectionTemplate` for adding this to `query-log`
 newtype ResolvedConnectionTemplateWrapper b = ResolvedConnectionTemplateWrapper

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

@@ -21,6 +21,7 @@ import Data.Kind (Type)
 import Data.Text.Casing (GQLNameIdentifier)
 import Data.Text.Extended
 import Data.Typeable (Typeable)
+import Hasura.Backends.Postgres.Connection.Settings (ConnectionTemplate (..))
 import Hasura.Base.Error
 import Hasura.Base.ToErrorValue
 import Hasura.EncJSON (EncJSON)
@@ -310,8 +311,8 @@ class
   type ConnectionTemplateRequestContext b :: Type
   type ConnectionTemplateRequestContext b = () -- Uninmplemented value
 
-  resolveConnectionTemplate :: SourceConfig b -> ConnectionTemplateRequestContext b -> Either QErr EncJSON
+  resolveConnectionTemplate :: SourceConfig b -> ConnectionTemplateRequestContext b -> Maybe ConnectionTemplate -> Either QErr EncJSON
-  resolveConnectionTemplate _ _ = Left (err400 (NotSupported) "connection template is not implemented")
+  resolveConnectionTemplate _ _ _ = Left (err400 (NotSupported) "connection template is not implemented")
 
   -- | Information about the query execution that may be useful for debugging
   -- or reporting.