refactor: add PreQuery module

Small step towards solving #3934.

Adds a module that builds SQL snippets for the transaction variables
and the pre-request function.

postgrest.cabal

@@ -71,6 +71,7 @@ library
                       PostgREST.Network
                       PostgREST.Observation
                       PostgREST.Query
+                      PostgREST.Query.PreQuery
                       PostgREST.Query.QueryBuilder
                       PostgREST.Query.SqlFragment
                       PostgREST.Query.Statements

src/PostgREST/Query.hs

@@ -11,7 +11,6 @@ module PostgREST.Query
 import qualified Data.Aeson                        as JSON
 import qualified Data.Aeson.KeyMap                 as KM
 import qualified Data.ByteString                   as BS
-import qualified Data.ByteString.Lazy.Char8        as LBS
 import qualified Data.HashMap.Strict               as HM
 import qualified Data.Set                          as S
 import qualified Hasql.Decoders                    as HD
@@ -22,6 +21,7 @@ import qualified Hasql.Transaction                 as SQL
 import qualified Hasql.Transaction.Sessions        as SQL
 
 import qualified PostgREST.Error              as Error
+import qualified PostgREST.Query.PreQuery     as PreQuery
 import qualified PostgREST.Query.QueryBuilder as QueryBuilder
 import qualified PostgREST.Query.Statements   as Statements
 import qualified PostgREST.SchemaCache        as SchemaCache
@@ -32,7 +32,6 @@ import PostgREST.ApiRequest              (ApiRequest (..),
 import PostgREST.ApiRequest.Preferences  (PreferCount (..),
                                           PreferHandling (..),
                                           PreferMaxAffected (..),
-                                          PreferTimezone (..),
                                           PreferTransaction (..),
                                           Preferences (..),
                                           shouldCount)
@@ -48,11 +47,6 @@ import PostgREST.Plan                    (ActionPlan (..),
                                           InfoPlan (..),
                                           InspectPlan (..))
 import PostgREST.Plan.MutatePlan         (MutatePlan (..))
-import PostgREST.Query.SqlFragment       (escapeIdentList, fromQi,
-                                          intercalateSnippet,
-                                          setConfigWithConstantName,
-                                          setConfigWithConstantNameJSON,
-                                          setConfigWithDynamicName)
 import PostgREST.Query.Statements        (ResultSet (..))
 import PostgREST.SchemaCache             (SchemaCache (..))
 import PostgREST.SchemaCache.Identifiers (QualifiedIdentifier (..))
@@ -90,8 +84,8 @@ query config AuthResult{..} apiReq (Db plan) sCache =
     txMode = planTxMode plan
     (mainActionQuery, mainSQLQuery) = actionQuery plan config apiReq sCache
     dbHandler = do
-      setPgLocals plan config authClaims authRole apiReq
+      runTxVarQuery plan config authClaims authRole apiReq
-      runPreReq config
+      runPreReqQuery config
       mainActionQuery
 
 planTxMode :: DbActionPlan -> SQL.Mode
@@ -263,38 +257,17 @@ optionalRollback AppConfig{..} ApiRequest{iPreferences=Preferences{..}} = do
     shouldRollback =
       preferTransaction == Just Rollback
 
--- | Set transaction scoped settings
+runTxVarQuery :: DbActionPlan -> AppConfig -> KM.KeyMap JSON.Value -> BS.ByteString -> ApiRequest -> DbHandler ()
-setPgLocals :: DbActionPlan -> AppConfig -> KM.KeyMap JSON.Value -> BS.ByteString -> ApiRequest -> DbHandler ()
+runTxVarQuery dbActPlan conf@AppConfig{..} claims role apireq = lift $
-setPgLocals dbActPlan AppConfig{..} claims role ApiRequest{..} = lift $
   SQL.statement mempty $ SQL.dynamicallyParameterized
-    -- To ensure `GRANT SET ON PARAMETER <superuser_setting> TO authenticator` works, the role settings must be set before the impersonated role.
+    (PreQuery.txVarQuery dbActPlan conf claims role apireq)
-    -- Otherwise the GRANT SET would have to be applied to the impersonated role. See https://github.com/PostgREST/postgrest/issues/3045
-    ("select " <> intercalateSnippet ", " (searchPathSql : roleSettingsSql ++ roleSql ++ claimsSql ++ [methodSql, pathSql] ++ headersSql ++ cookiesSql ++ timezoneSql ++ funcSettingsSql ++ appSettingsSql))
     HD.noResult configDbPreparedStatements
-  where
-    methodSql = setConfigWithConstantName ("request.method", iMethod)
-    pathSql = setConfigWithConstantName ("request.path", iPath)
-    headersSql = setConfigWithConstantNameJSON "request.headers" iHeaders
-    cookiesSql = setConfigWithConstantNameJSON "request.cookies" iCookies
-    claimsSql = [setConfigWithConstantName ("request.jwt.claims", LBS.toStrict $ JSON.encode claims)]
-    roleSql = [setConfigWithConstantName ("role", role)]
-    roleSettingsSql = setConfigWithDynamicName <$> HM.toList (fromMaybe mempty $ HM.lookup role configRoleSettings)
-    appSettingsSql = setConfigWithDynamicName . join bimap toUtf8 <$> configAppSettings
-    timezoneSql = maybe mempty (\(PreferTimezone tz) -> [setConfigWithConstantName ("timezone", tz)]) $ preferTimezone iPreferences
-    funcSettingsSql = setConfigWithDynamicName . join bimap toUtf8 <$> funcSettings
-    searchPathSql =
-      let schemas = escapeIdentList (iSchema : configDbExtraSearchPath) in
-      setConfigWithConstantName ("search_path", schemas)
-    funcSettings = case dbActPlan of
-      DbCall CallReadPlan{crProc} -> pdFuncSettings crProc
-      _                           -> mempty
 
--- | Runs the pre-request function.
+runPreReqQuery :: AppConfig -> DbHandler ()
-runPreReq :: AppConfig -> DbHandler ()
+runPreReqQuery conf = lift $ traverse_ (SQL.statement mempty . stmt) (configDbPreRequest conf)
-runPreReq conf = lift $ traverse_ (SQL.statement mempty . stmt) (configDbPreRequest conf)
   where
     stmt req = SQL.dynamicallyParameterized
-      ("select " <> fromQi req <> "()")
+      (PreQuery.preReqQuery req)
       HD.noResult
       (configDbPreparedStatements conf)
 

src/PostgREST/Query/PreQuery.hs

@@ -0,0 +1,65 @@
+{-# LANGUAGE NamedFieldPuns  #-}
+{-# LANGUAGE RecordWildCards #-}
+{-|
+Module      : PostgREST.Query.PreQuery
+Description : Builds queries that run prior to the main query
+-}
+module PostgREST.Query.PreQuery
+  ( txVarQuery
+  , preReqQuery
+  ) where
+
+import qualified Data.Aeson                      as JSON
+import qualified Data.Aeson.KeyMap               as KM
+import qualified Data.ByteString                 as BS
+import qualified Data.ByteString.Lazy.Char8      as LBS
+import qualified Data.HashMap.Strict             as HM
+import qualified Hasql.DynamicStatements.Snippet as SQL hiding (sql)
+
+
+
+import PostgREST.ApiRequest              (ApiRequest (..))
+import PostgREST.ApiRequest.Preferences  (PreferTimezone (..),
+                                          Preferences (..))
+import PostgREST.Config                  (AppConfig (..))
+import PostgREST.Plan                    (CallReadPlan (..),
+                                          DbActionPlan (..))
+import PostgREST.Query.SqlFragment       (escapeIdentList, fromQi,
+                                          intercalateSnippet,
+                                          setConfigWithConstantName,
+                                          setConfigWithConstantNameJSON,
+                                          setConfigWithDynamicName)
+import PostgREST.SchemaCache.Identifiers (QualifiedIdentifier (..))
+import PostgREST.SchemaCache.Routine     (Routine (..))
+
+import Protolude hiding (Handler)
+
+-- sets transaction variables
+txVarQuery :: DbActionPlan -> AppConfig -> KM.KeyMap JSON.Value -> BS.ByteString -> ApiRequest -> SQL.Snippet
+txVarQuery dbActPlan AppConfig{..} claims role ApiRequest{..} =
+    -- To ensure `GRANT SET ON PARAMETER <superuser_setting> TO authenticator` works, the role settings must be set before the impersonated role.
+    -- Otherwise the GRANT SET would have to be applied to the impersonated role. See https://github.com/PostgREST/postgrest/issues/3045
+    "select " <> intercalateSnippet ", " (
+      searchPathSql : roleSettingsSql ++ roleSql ++ claimsSql ++ [methodSql, pathSql] ++ headersSql ++ cookiesSql ++ timezoneSql ++ funcSettingsSql ++ appSettingsSql
+    )
+  where
+    methodSql = setConfigWithConstantName ("request.method", iMethod)
+    pathSql = setConfigWithConstantName ("request.path", iPath)
+    headersSql = setConfigWithConstantNameJSON "request.headers" iHeaders
+    cookiesSql = setConfigWithConstantNameJSON "request.cookies" iCookies
+    claimsSql = [setConfigWithConstantName ("request.jwt.claims", LBS.toStrict $ JSON.encode claims)]
+    roleSql = [setConfigWithConstantName ("role", role)]
+    roleSettingsSql = setConfigWithDynamicName <$> HM.toList (fromMaybe mempty $ HM.lookup role configRoleSettings)
+    appSettingsSql = setConfigWithDynamicName . join bimap toUtf8 <$> configAppSettings
+    timezoneSql = maybe mempty (\(PreferTimezone tz) -> [setConfigWithConstantName ("timezone", tz)]) $ preferTimezone iPreferences
+    funcSettingsSql = setConfigWithDynamicName . join bimap toUtf8 <$> funcSettings
+    searchPathSql =
+      let schemas = escapeIdentList (iSchema : configDbExtraSearchPath) in
+      setConfigWithConstantName ("search_path", schemas)
+    funcSettings = case dbActPlan of
+      DbCall CallReadPlan{crProc} -> pdFuncSettings crProc
+      _                           -> mempty
+
+-- runs the pre-request function
+preReqQuery :: QualifiedIdentifier -> SQL.Snippet
+preReqQuery preRequest = "select " <> fromQi preRequest <> "()"