refactor: configDbPreRequest to QualifiedIdentifier

src/PostgREST/Config.hs

@@ -51,9 +51,11 @@ import Numeric                 (readOct, showOct)
 import System.Environment      (getEnvironment)
 import System.Posix.Types      (FileMode)
 
-import PostgREST.Config.JSPath (JSPath, JSPathExp (..), pRoleClaimKey)
+import PostgREST.Config.JSPath           (JSPath, JSPathExp (..),
-import PostgREST.Config.Proxy  (Proxy (..), isMalformedProxyUri,
+                                          pRoleClaimKey)
-                                toURI)
+import PostgREST.Config.Proxy            (Proxy (..),
+                                          isMalformedProxyUri, toURI)
+import PostgREST.DbStructure.Identifiers (QualifiedIdentifier, toQi)
 
 import Protolude      hiding (Proxy, toList, toS)
 import Protolude.Conv (toS)
@@ -68,7 +70,7 @@ data AppConfig = AppConfig
   , configDbMaxRows             :: Maybe Integer
   , configDbPoolSize            :: Int
   , configDbPoolTimeout         :: NominalDiffTime
-  , configDbPreRequest          :: Maybe Text
+  , configDbPreRequest          :: Maybe QualifiedIdentifier
   , configDbPreparedStatements  :: Bool
   , configDbRootSpec            :: Maybe Text
   , configDbSchemas             :: NonEmpty Text
@@ -113,7 +115,7 @@ toText conf =
       ,("db-max-rows",                   maybe "\"\"" show . configDbMaxRows)
       ,("db-pool",                       show . configDbPoolSize)
       ,("db-pool-timeout",               show . floor . configDbPoolTimeout)
-      ,("db-pre-request",            q . fromMaybe mempty . configDbPreRequest)
+      ,("db-pre-request",            q . maybe mempty show . configDbPreRequest)
       ,("db-prepared-statements",        T.toLower . show . configDbPreparedStatements)
       ,("db-root-spec",              q . fromMaybe mempty . configDbRootSpec)
       ,("db-schemas",                q . T.intercalate "," . toList . configDbSchemas)
@@ -197,8 +199,8 @@ parser optPath env dbSettings =
                      (optInt "max-rows")
     <*> (fromMaybe 10 <$> optInt "db-pool")
     <*> (fromIntegral . fromMaybe 10 <$> optInt "db-pool-timeout")
-    <*> optWithAlias (optString "db-pre-request")
+    <*> (fmap toQi <$> optWithAlias (optString "db-pre-request")
-                     (optString "pre-request")
+                                    (optString "pre-request"))
     <*> (fromMaybe True <$> optBool "db-prepared-statements")
     <*> optWithAlias (optString "db-root-spec")
                      (optString "root-spec")

src/PostgREST/DbStructure/Identifiers.hs

@@ -6,9 +6,12 @@ module PostgREST.DbStructure.Identifiers
   , Schema
   , TableName
   , FieldName
+  , toQi
   ) where
 
 import qualified Data.Aeson as JSON
+import qualified Data.Text  as T
+import qualified GHC.Show
 
 import Protolude
 
@@ -23,6 +26,17 @@ data QualifiedIdentifier = QualifiedIdentifier
 
 instance Hashable QualifiedIdentifier
 
+instance Show QualifiedIdentifier where
+  show (QualifiedIdentifier s i) =
+    (if T.null s then mempty else toS s <> ".") <> toS i
+
+-- TODO: Handle a case where the QI comes like this: "my.fav.schema"."my.identifier"
+-- Right now it only handles the schema.identifier case
+toQi :: Text -> QualifiedIdentifier
+toQi txt = case T.drop 1 <$> T.breakOn "." txt of
+  (i, "") -> QualifiedIdentifier mempty i
+  (s, i)  -> QualifiedIdentifier s i
+
 type Schema = Text
 type TableName = Text
 type FieldName = Text

src/PostgREST/Middleware.hs

@@ -42,7 +42,7 @@ import System.Log.FastLogger     (toLogStr)
 import PostgREST.Config             (AppConfig (..), LogLevel (..))
 import PostgREST.Error              (Error, errorResponseFor)
 import PostgREST.GucHeader          (addHeadersIfNotIncluded)
-import PostgREST.Query.SqlFragment  (intercalateSnippet,
+import PostgREST.Query.SqlFragment  (fromQi, intercalateSnippet,
                                      unknownLiteral)
 import PostgREST.Request.ApiRequest (ApiRequest (..))
 
@@ -75,7 +75,7 @@ runPgLocals conf claims app req = do
     searchPathSql =
       let schemas = T.intercalate ", " (iSchema req : configDbExtraSearchPath conf) in
       setConfigLocal mempty ("search_path", schemas)
-    preReqSql = (\f -> "select " <> toS f <> "();") <$> configDbPreRequest conf
+    preReqSql = (\f -> "select " <> fromQi f <> "();") <$> configDbPreRequest conf
 
     -- | Do a pg set_config(setting, value, true) call. This is equivalent to a SET LOCAL.
     setConfigLocal :: Text -> (Text, Text) -> H.Snippet

src/PostgREST/Query/SqlFragment.hs

@@ -135,6 +135,7 @@ pgFmtLit x =
    then "E" <> slashed
    else slashed
 
+-- TODO: refactor by following https://github.com/PostgREST/postgrest/pull/1631#issuecomment-711070833
 pgFmtIdent :: Text -> SqlFragment
 pgFmtIdent x = encodeUtf8 $ "\"" <> T.replace "\"" "\"\"" (trimNullChars x) <> "\""
 

test/SpecHelper.hs

@@ -22,10 +22,12 @@ import Test.Hspec
 import Test.Hspec.Wai
 import Text.Heredoc
 
-import PostgREST.Config (AppConfig (..), JSPathExp (..),
+import PostgREST.Config                  (AppConfig (..),
-                         LogLevel (..), parseSecret)
+                                          JSPathExp (..),
-import Protolude        hiding (toS)
+                                          LogLevel (..), parseSecret)
-import Protolude.Conv   (toS)
+import PostgREST.DbStructure.Identifiers (QualifiedIdentifier (..))
+import Protolude                         hiding (toS)
+import Protolude.Conv                    (toS)
 
 matchContentTypeJson :: MatchHeader
 matchContentTypeJson = "Content-Type" <:> "application/json; charset=utf-8"
@@ -79,7 +81,7 @@ _baseCfg = let secret = Just $ encodeUtf8 "reallyreallyreallyreallyverysafe" in
   , configDbMaxRows             = Nothing
   , configDbPoolSize            = 10
   , configDbPoolTimeout         = 10
-  , configDbPreRequest          = Just "test.switch_role"
+  , configDbPreRequest          = Just $ QualifiedIdentifier "test" "switch_role"
   , configDbPreparedStatements  = True
   , configDbRootSpec            = Nothing
   , configDbSchemas             = fromList ["test"]
@@ -169,7 +171,7 @@ testCfgHtmlRawOutput :: Text -> AppConfig
 testCfgHtmlRawOutput testDbConn = (testCfg testDbConn) { configRawMediaTypes = ["text/html"] }
 
 testCfgResponseHeaders :: Text -> AppConfig
-testCfgResponseHeaders testDbConn = (testCfg testDbConn) { configDbPreRequest = Just "custom_headers" }
+testCfgResponseHeaders testDbConn = (testCfg testDbConn) { configDbPreRequest = Just $ QualifiedIdentifier mempty "custom_headers" }
 
 testMultipleSchemaCfg :: Text -> AppConfig
 testMultipleSchemaCfg testDbConn = (testCfg testDbConn) { configDbSchemas = fromList ["v1", "v2"] }