refactor: make more of Response pure

PostgREST · Sep 25, 2023 · f1f53ee · f1f53ee
1 parent edbc4f9
commit f1f53ee
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 17 deletions.
diff --git a/src/PostgREST/App.hs b/src/PostgREST/App.hs
@@ -178,12 +178,14 @@ handleRequest AuthResult{..} conf appState authenticated prepared pgVer apiReq@A
     (ActionMutate MutationCreate, TargetIdent identifier) -> do
       mrPlan <- liftEither $ Plan.mutateReadPlan MutationCreate apiReq identifier conf sCache
       resultSet <- runQuery roleIsoLvl (Plan.mrTxMode mrPlan) $ Query.createQuery mrPlan apiReq conf
-      return $ Response.createResponse identifier mrPlan apiReq resultSet serverTimingParams
+      pgrst <- liftEither $ Response.createResponse identifier mrPlan apiReq resultSet serverTimingParams
+      return $ pgrstResponse pgrst
 
     (ActionMutate MutationUpdate, TargetIdent identifier) -> do
       mrPlan <- liftEither $ Plan.mutateReadPlan MutationUpdate apiReq identifier conf sCache
       resultSet <- runQuery roleIsoLvl (Plan.mrTxMode mrPlan) $ Query.updateQuery mrPlan apiReq conf
-      return $ Response.updateResponse mrPlan apiReq resultSet serverTimingParams
+      pgrst <- liftEither $ Response.updateResponse mrPlan apiReq resultSet serverTimingParams
+      return $ pgrstResponse pgrst
 
     (ActionMutate MutationSingleUpsert, TargetIdent identifier) -> do
       mrPlan <- liftEither $ Plan.mutateReadPlan MutationSingleUpsert apiReq identifier conf sCache

diff --git a/src/PostgREST/Response.hs b/src/PostgREST/Response.hs
@@ -111,12 +111,11 @@ readResponse WrappedReadPlan{wrMedia} headersOnly identifier ctxApiRequest@ApiRe
     RSPlan plan ->
       Right $ PgrstResponse HTTP.status200 (contentTypeHeaders wrMedia ctxApiRequest) $ LBS.fromStrict plan
 
-createResponse :: QualifiedIdentifier -> MutateReadPlan -> ApiRequest -> ResultSet -> Maybe ServerTimingParams -> Wai.Response
+createResponse :: QualifiedIdentifier -> MutateReadPlan -> ApiRequest -> ResultSet -> Maybe ServerTimingParams -> Either Error.Error PgrstResponse
 createResponse QualifiedIdentifier{..} MutateReadPlan{mrMutatePlan, mrMedia} ctxApiRequest@ApiRequest{iPreferences=Preferences{..}, ..} resultSet serverTimingParams = case resultSet of
   RSStandard{..} -> do
     let
       pkCols = case mrMutatePlan of { Insert{insPkCols} -> insPkCols; _ -> mempty;}
-      response = gucResponse rsGucStatus rsGucHeaders
       prefHeader = prefAppliedHeader $
         Preferences (if null pkCols && isNothing (qsOnConflict iQueryParams) then Nothing else preferResolution)
         preferRepresentation Nothing preferCount preferTransaction preferMissing
@@ -136,33 +135,41 @@ createResponse QualifiedIdentifier{..} MutateReadPlan{mrMutatePlan, mrMedia} ctx
           , prefHeader
           ] ++ serverTimingHeader serverTimingParams
 
-    case preferRepresentation of
-      Just Full -> response HTTP.status201 (headers ++ contentTypeHeaders mrMedia ctxApiRequest) (LBS.fromStrict rsBody)
-      Just None -> response HTTP.status201 headers mempty
-      Just HeadersOnly -> response HTTP.status201 headers mempty
-      Nothing -> response HTTP.status201 headers mempty
+    let status = HTTP.status201
+    let (headers', bod) = case preferRepresentation of
+          Just Full -> (headers ++ contentTypeHeaders mrMedia ctxApiRequest, LBS.fromStrict rsBody)
+          Just None -> (headers, mempty)
+          Just HeadersOnly -> (headers, mempty)
+          Nothing -> (headers, mempty)
+
+    (ovStatus, ovHeaders) <- overrideStatusHeaders rsGucStatus rsGucHeaders status headers'
 
+    Right $ PgrstResponse ovStatus ovHeaders bod
   RSPlan plan ->
-    Wai.responseLBS HTTP.status200 (contentTypeHeaders mrMedia ctxApiRequest) $ LBS.fromStrict plan
+    Right $ PgrstResponse HTTP.status200 (contentTypeHeaders mrMedia ctxApiRequest) $ LBS.fromStrict plan
 
-updateResponse :: MutateReadPlan -> ApiRequest -> ResultSet -> Maybe ServerTimingParams -> Wai.Response
+updateResponse :: MutateReadPlan -> ApiRequest -> ResultSet -> Maybe ServerTimingParams -> Either Error.Error PgrstResponse
 updateResponse MutateReadPlan{mrMedia} ctxApiRequest@ApiRequest{iPreferences=Preferences{..}} resultSet serverTimingParams = case resultSet of
   RSStandard{..} -> do
     let
-      response = gucResponse rsGucStatus rsGucHeaders
       contentRangeHeader =
         Just . RangeQuery.contentRangeH 0 (rsQueryTotal - 1) $
           if shouldCount preferCount then Just rsQueryTotal else Nothing
       prefHeader = prefAppliedHeader $ Preferences Nothing preferRepresentation Nothing preferCount preferTransaction preferMissing
       headers = catMaybes [contentRangeHeader, prefHeader] ++ serverTimingHeader serverTimingParams
 
-    case preferRepresentation of
-        Just Full -> response HTTP.status200 (headers ++ contentTypeHeaders mrMedia ctxApiRequest) (LBS.fromStrict rsBody)
-        Just None -> response HTTP.status204 headers mempty
-        _ -> response HTTP.status204 headers mempty
+    let 
+      (status, headers', body) = case preferRepresentation of
+                          Just Full -> (HTTP.status200, headers ++ contentTypeHeaders mrMedia ctxApiRequest, LBS.fromStrict rsBody)
+                          Just None -> (HTTP.status204, headers, mempty)
+                          _ -> (HTTP.status204, headers, mempty)
+
+    (ovStatus, ovHeaders) <- overrideStatusHeaders rsGucStatus rsGucHeaders status headers'
+
+    Right $ PgrstResponse ovStatus ovHeaders body
 
   RSPlan plan ->
-    Wai.responseLBS HTTP.status200 (contentTypeHeaders mrMedia ctxApiRequest) $ LBS.fromStrict plan
+    Right $ PgrstResponse HTTP.status200 (contentTypeHeaders mrMedia ctxApiRequest) $ LBS.fromStrict plan
 
 singleUpsertResponse :: MutateReadPlan -> ApiRequest -> ResultSet -> Maybe ServerTimingParams -> Wai.Response
 singleUpsertResponse MutateReadPlan{mrMedia} ctxApiRequest@ApiRequest{iPreferences=Preferences{..}} resultSet serverTimingParams = case resultSet of