Skip to content

Commit

Permalink
refactor: stricter plan media type
Browse files Browse the repository at this point in the history
  • Loading branch information
steve-chavez committed Jun 20, 2023
1 parent c1a8661 commit f8e8e36
Show file tree
Hide file tree
Showing 6 changed files with 53 additions and 51 deletions.
7 changes: 4 additions & 3 deletions src/PostgREST/ApiRequest.hs
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,8 @@ import PostgREST.ApiRequest.Types (ApiRequestError (..),
RangeError (..))
import PostgREST.Config (AppConfig (..),
OpenAPIMode (..))
import PostgREST.MediaType (MediaType (..))
import PostgREST.MediaType (MTPlanFormat (..),
MediaType (..))
import PostgREST.RangeQuery (NonnegRange, allRange,
convertToLimitZeroRange,
hasLimitZero,
Expand Down Expand Up @@ -363,15 +364,15 @@ producedMediaTypes conf action path =
case action of
ActionRead _ -> defaultMediaTypes ++ rawMediaTypes
ActionInvoke _ -> invokeMediaTypes
ActionInspect _ -> [MTOpenAPI, MTApplicationJSON, MTAny]
ActionInfo -> defaultMediaTypes
ActionMutate _ -> defaultMediaTypes
ActionInspect _ -> [MTOpenAPI, MTApplicationJSON, MTAny]
where
invokeMediaTypes =
defaultMediaTypes
++ rawMediaTypes
++ [MTOpenAPI | pathIsRootSpec path]
defaultMediaTypes =
[MTApplicationJSON, MTSingularJSON, MTGeoJSON, MTTextCSV] ++
[MTPlan Nothing Nothing mempty | configDbPlanEnabled conf] ++ [MTAny]
[MTPlan MTApplicationJSON PlanText mempty | configDbPlanEnabled conf] ++ [MTAny]
rawMediaTypes = configRawMediaTypes conf `union` [MTOctetStream, MTTextPlain, MTTextXML]
52 changes: 27 additions & 25 deletions src/PostgREST/MediaType.hs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ module PostgREST.MediaType

import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BS (c2w)
import Data.Maybe (fromJust)

import Network.HTTP.Types.Header (Header, hContentType)

Expand All @@ -39,7 +38,7 @@ data MediaType
| MTOctetStream
| MTAny
| MTOther ByteString
| MTPlan (Maybe MediaType) (Maybe MTPlanFormat) [MTPlanOption]
| MTPlan MediaType MTPlanFormat [MTPlanOption]
instance Eq MediaType where
MTApplicationJSON == MTApplicationJSON = True
MTSingularJSON == MTSingularJSON = True
Expand Down Expand Up @@ -84,8 +83,8 @@ toMime MTOctetStream = "application/octet-stream"
toMime MTAny = "*/*"
toMime (MTOther ct) = ct
toMime (MTPlan mt fmt opts) =
"application/vnd.pgrst.plan" <> maybe mempty (\x -> "+" <> toMimePlanFormat x) fmt <>
(if isNothing mt then mempty else "; for=\"" <> toMime (fromJust mt) <> "\"") <>
"application/vnd.pgrst.plan+" <> toMimePlanFormat fmt <>
("; for=\"" <> toMime mt <> "\"") <>
(if null opts then mempty else "; options=" <> BS.intercalate "|" (toMimePlanOption <$> opts))

toMimePlanOption :: MTPlanOption -> ByteString
Expand All @@ -105,13 +104,13 @@ toMimePlanFormat PlanText = "text"
-- MTApplicationJSON
--
-- >>> decodeMediaType "application/vnd.pgrst.plan;"
-- MTPlan Nothing Nothing []
-- MTPlan MTApplicationJSON PlanText []
--
-- >>> decodeMediaType "application/vnd.pgrst.plan;for=\"application/json\""
-- MTPlan (Just MTApplicationJSON) Nothing []
-- MTPlan MTApplicationJSON PlanText []
--
-- >>> decodeMediaType "application/vnd.pgrst.plan+text;for=\"text/csv\""
-- MTPlan (Just MTTextCSV) (Just PlanText) []
-- >>> decodeMediaType "application/vnd.pgrst.plan+json;for=\"text/csv\""
-- MTPlan MTTextCSV PlanJSON []
decodeMediaType :: BS.ByteString -> MediaType
decodeMediaType mt =
case BS.split (BS.c2w ';') mt of
Expand All @@ -125,28 +124,31 @@ decodeMediaType mt =
"application/vnd.pgrst.object":_ -> MTSingularJSON
"application/x-www-form-urlencoded":_ -> MTUrlEncoded
"application/octet-stream":_ -> MTOctetStream
"application/vnd.pgrst.plan":rest -> getPlan Nothing rest
"application/vnd.pgrst.plan+text":rest -> getPlan (Just PlanText) rest
"application/vnd.pgrst.plan+json":rest -> getPlan (Just PlanJSON) rest
"application/vnd.pgrst.plan":rest -> getPlan PlanText rest
"application/vnd.pgrst.plan+text":rest -> getPlan PlanText rest
"application/vnd.pgrst.plan+json":rest -> getPlan PlanJSON rest
"*/*":_ -> MTAny
other:_ -> MTOther other
_ -> MTAny
where
getPlan fmt rest =
let
opts = BS.split (BS.c2w '|') $ fromMaybe mempty (BS.stripPrefix "options=" =<< find (BS.isPrefixOf "options=") rest)
inOpts str = str `elem` opts
mtFor = decodeMediaType . dropAround (== BS.c2w '"') <$> (BS.stripPrefix "for=" =<< find (BS.isPrefixOf "for=") rest)
dropAround p = BS.dropWhile p . BS.dropWhileEnd p in
MTPlan mtFor fmt $
[PlanAnalyze | inOpts "analyze" ] ++
[PlanVerbose | inOpts "verbose" ] ++
[PlanSettings | inOpts "settings"] ++
[PlanBuffers | inOpts "buffers" ] ++
[PlanWAL | inOpts "wal" ]
let
opts = BS.split (BS.c2w '|') $ fromMaybe mempty (BS.stripPrefix "options=" =<< find (BS.isPrefixOf "options=") rest)
inOpts str = str `elem` opts
dropAround p = BS.dropWhile p . BS.dropWhileEnd p
mtFor = fromMaybe MTApplicationJSON $ do
foundFor <- find (BS.isPrefixOf "for=") rest
strippedFor <- BS.stripPrefix "for=" foundFor
pure . decodeMediaType $ dropAround (== BS.c2w '"') strippedFor
in
MTPlan mtFor fmt $
[PlanAnalyze | inOpts "analyze" ] ++
[PlanVerbose | inOpts "verbose" ] ++
[PlanSettings | inOpts "settings"] ++
[PlanBuffers | inOpts "buffers" ] ++
[PlanWAL | inOpts "wal" ]

getMediaType :: MediaType -> MediaType
getMediaType mt = case mt of
MTPlan (Just mType) _ _ -> mType
MTPlan Nothing _ _ -> MTApplicationJSON
other -> other
MTPlan mType _ _ -> mType
other -> other
8 changes: 4 additions & 4 deletions src/PostgREST/Plan.hs
Original file line number Diff line number Diff line change
Expand Up @@ -633,10 +633,10 @@ binaryField AppConfig{configRawMediaTypes} acceptMediaType proc rpTree
where
isRawMediaType = acceptMediaType `elem` configRawMediaTypes `L.union` [MTOctetStream, MTTextPlain, MTTextXML] || isRawPlan acceptMediaType
isRawPlan mt = case mt of
MTPlan (Just MTOctetStream) _ _ -> True
MTPlan (Just MTTextPlain) _ _ -> True
MTPlan (Just MTTextXML) _ _ -> True
_ -> False
MTPlan MTOctetStream _ _ -> True
MTPlan MTTextPlain _ _ -> True
MTPlan MTTextXML _ _ -> True
_ -> False

fstFieldName :: ReadPlanTree -> Maybe FieldName
fstFieldName (Node ReadPlan{select=(("*", []), _, _):_} []) = Nothing
Expand Down
7 changes: 3 additions & 4 deletions src/PostgREST/Query/SqlFragment.hs
Original file line number Diff line number Diff line change
Expand Up @@ -431,7 +431,7 @@ intercalateSnippet :: ByteString -> [SQL.Snippet] -> SQL.Snippet
intercalateSnippet _ [] = mempty
intercalateSnippet frag snippets = foldr1 (\a b -> a <> SQL.sql frag <> b) snippets

explainF :: Maybe MTPlanFormat -> [MTPlanOption] -> SQL.Snippet -> SQL.Snippet
explainF :: MTPlanFormat -> [MTPlanOption] -> SQL.Snippet -> SQL.Snippet
explainF fmt opts snip =
"EXPLAIN (" <>
SQL.sql (BS.intercalate ", " (fmtPlanFmt fmt : (fmtPlanOpt <$> opts))) <>
Expand All @@ -444,9 +444,8 @@ explainF fmt opts snip =
fmtPlanOpt PlanBuffers = "BUFFERS"
fmtPlanOpt PlanWAL = "WAL"

fmtPlanFmt Nothing = "FORMAT TEXT"
fmtPlanFmt (Just PlanJSON) = "FORMAT JSON"
fmtPlanFmt (Just PlanText) = "FORMAT TEXT"
fmtPlanFmt PlanText = "FORMAT TEXT"
fmtPlanFmt PlanJSON = "FORMAT JSON"

-- | Do a pg set_config(setting, value, true) call. This is equivalent to a SET LOCAL.
setConfigLocal :: ByteString -> (ByteString, ByteString) -> SQL.Snippet
Expand Down
2 changes: 1 addition & 1 deletion src/PostgREST/Query/Statements.hs
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,7 @@ preparePlanRows :: SQL.Snippet -> Bool -> SQL.Statement () (Maybe Int64)
preparePlanRows countQuery =
SQL.dynamicallyParameterized snippet decodeIt
where
snippet = explainF (Just PlanJSON) mempty countQuery
snippet = explainF PlanJSON mempty countQuery
decodeIt :: HD.Result (Maybe Int64)
decodeIt =
let row = HD.singleRow $ column HD.bytea in
Expand Down
28 changes: 14 additions & 14 deletions test/spec/Feature/Query/PlanSpec.hs
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ spec actualPgVersion = do
resStatus = simpleStatus r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; charset=utf-8")
resStatus `shouldBe` Status { statusCode = 200, statusMessage="OK" }
totalCost `shouldBe`
if actualPgVersion > pgVersion120
Expand All @@ -49,7 +49,7 @@ spec actualPgVersion = do
resStatus = simpleStatus r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; charset=utf-8")
resStatus `shouldBe` Status { statusCode = 200, statusMessage="OK" }
totalCost `shouldBe`
if actualPgVersion > pgVersion120
Expand All @@ -65,7 +65,7 @@ spec actualPgVersion = do
resHeaders = simpleHeaders r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; options=buffers; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; options=buffers; charset=utf-8")
resBody `shouldSatisfy` (\t -> T.isInfixOf "Shared Hit Blocks" (decodeUtf8 $ BS.toStrict t))
else do
-- analyze is required for buffers on pg < 13
Expand All @@ -75,7 +75,7 @@ spec actualPgVersion = do
resHeaders = simpleHeaders r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; options=analyze|buffers; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; options=analyze|buffers; charset=utf-8")
blocks `shouldBe` Just [aesonQQ| 1.0 |]

when (actualPgVersion >= pgVersion120) $
Expand All @@ -86,7 +86,7 @@ spec actualPgVersion = do
resHeaders = simpleHeaders r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; options=settings; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; options=settings; charset=utf-8")
searchPath `shouldBe`
Just [aesonQQ|
{
Expand All @@ -102,7 +102,7 @@ spec actualPgVersion = do
resHeaders = simpleHeaders r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; options=analyze|wal; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; options=analyze|wal; charset=utf-8")
walRecords `shouldBe` Just [aesonQQ|0|]

it "outputs columns info when using the verbose option" $ do
Expand All @@ -112,7 +112,7 @@ spec actualPgVersion = do
resHeaders = simpleHeaders r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; options=verbose; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; options=verbose; charset=utf-8")
cols `shouldBe` Just [aesonQQ| ["projects.id", "projects.name", "projects.client_id"] |]

it "outputs the plan for application/json " $ do
Expand Down Expand Up @@ -151,7 +151,7 @@ spec actualPgVersion = do
resStatus = simpleStatus r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; charset=utf-8")
resStatus `shouldBe` Status { statusCode = 200, statusMessage="OK" }
totalCost `shouldBe` 3.27

Expand All @@ -164,7 +164,7 @@ spec actualPgVersion = do
resStatus = simpleStatus r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; charset=utf-8")
resStatus `shouldBe` Status { statusCode = 200, statusMessage="OK" }
totalCost `shouldBe` 12.45

Expand All @@ -177,7 +177,7 @@ spec actualPgVersion = do
resStatus = simpleStatus r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; charset=utf-8")
resStatus `shouldBe` Status { statusCode = 200, statusMessage="OK" }
totalCost `shouldBe` 15.68

Expand All @@ -191,7 +191,7 @@ spec actualPgVersion = do
resStatus = simpleStatus r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; charset=utf-8")
resStatus `shouldBe` Status { statusCode = 200, statusMessage="OK" }
totalCost `shouldBe` 1.29

Expand All @@ -216,7 +216,7 @@ spec actualPgVersion = do
resStatus = simpleStatus r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; charset=utf-8")
resStatus `shouldBe` Status { statusCode = 200, statusMessage="OK" }
totalCost `shouldBe` 68.56

Expand All @@ -241,7 +241,7 @@ spec actualPgVersion = do
resStatus = simpleStatus r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+text; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+text; for=\"application/json\"; charset=utf-8")
resStatus `shouldBe` Status { statusCode = 200, statusMessage="OK" }
resBody `shouldSatisfy` (\t -> LBS.take 9 t == "Aggregate")

Expand All @@ -254,7 +254,7 @@ spec actualPgVersion = do
resStatus = simpleStatus r

liftIO $ do
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan; charset=utf-8")
resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+text; for=\"application/json\"; charset=utf-8")
resStatus `shouldBe` Status { statusCode = 200, statusMessage="OK" }
resBody `shouldSatisfy` (\t -> LBS.take 9 t == "Aggregate")

Expand Down

0 comments on commit f8e8e36

Please sign in to comment.