Add response.headers on GET/POST/PATCH/PUT/DELETE

2020-01-21 12:28:05 -05:00
parent 7dade7f466
commit c7f78fa7fc
9 changed files with 202 additions and 115 deletions
@@ -131,24 +131,26 @@ app dbStructure proc cols conf apiRequest =
                             then limitedQuery cq ((+ 1) <$> maxRows) -- LIMIT maxRows + 1 so we can determine below that maxRows was surpassed
                             else cq
                  stm = createReadStatement q cQuery (contentType == CTSingularJSON) shouldCount
-                        (contentType == CTTextCSV) bField
+                        (contentType == CTTextCSV) bField pgVer
                  explStm = createExplainStatement cq
              row <- H.statement () stm
-              let (tableTotal, queryTotal, _ , body) = row
-              total <- if | plannedCount   -> H.statement () explStm
-                          | estimatedCount -> if tableTotal > (fromIntegral <$> maxRows)
-                                                then do estTotal <- H.statement () explStm
-                                                        pure $ if estTotal > tableTotal then estTotal else tableTotal
-                                                else pure tableTotal
-                          | otherwise      -> pure tableTotal
-              let (status, contentRange) = rangeStatusHeader topLevelRange queryTotal total
-              return $
-                if contentType == CTSingularJSON && queryTotal /= 1
-                  then errorResponseFor . singularityError $ queryTotal
-                  else responseLBS status
-                    [toHeader contentType, contentRange,
-                     contentLocationH tName (iCanonicalQS apiRequest)]
-                    (if headersOnly then mempty else toS body)
+              let (tableTotal, queryTotal, _ , body, gucHeaders) = row
+              case gucHeaders of
+                Left _ -> return . errorResponseFor $ GucHeadersError
+                Right hs -> do
+                  total <- if | plannedCount   -> H.statement () explStm
+                              | estimatedCount -> if tableTotal > (fromIntegral <$> maxRows)
+                                                    then do estTotal <- H.statement () explStm
+                                                            pure $ if estTotal > tableTotal then estTotal else tableTotal
+                                                    else pure tableTotal
+                              | otherwise      -> pure tableTotal
+                  let (status, contentRange) = rangeStatusHeader topLevelRange queryTotal total
+                  return $
+                    if contentType == CTSingularJSON && queryTotal /= 1
+                      then errorResponseFor . singularityError $ queryTotal
+                      else responseLBS status
+                        ([toHeader contentType, contentRange, contentLocationH tName (iCanonicalQS apiRequest)] ++ (gucHToHeader <$> hs))
+                        (if headersOnly then mempty else toS body)

        (ActionCreate, TargetIdent (QualifiedIdentifier tSchema tName), Just pJson) ->
          case mutateSqlParts tSchema tName of
@@ -157,31 +159,31 @@ app dbStructure proc cols conf apiRequest =
              let pkCols = tablePKCols dbStructure tSchema tName
                  stm = createWriteStatement sq mq
                    (contentType == CTSingularJSON) True
-                    (contentType == CTTextCSV) (iPreferRepresentation apiRequest) pkCols
+                    (contentType == CTTextCSV) (iPreferRepresentation apiRequest) pkCols pgVer
              row <- H.statement (toS $ pjRaw pJson) stm
-              let (_, queryTotal, fs, body) = row
-                  headers = catMaybes [
-                      if null fs
-                        then Nothing
-                        else Just $ locationH tName fs
-                    , if iPreferRepresentation apiRequest == Full
-                        then Just $ toHeader contentType
-                        else Nothing
-                    , Just $ contentRangeH 1 0 $
-                        if shouldCount then Just queryTotal else Nothing
-                    , if null pkCols && isNothing (iOnConflict apiRequest)
-                        then Nothing
-                        else (\x -> ("Preference-Applied", show x)) <$> iPreferResolution apiRequest
-                    ]
-              if contentType == CTSingularJSON
-                 && queryTotal /= 1
-                then do
-                  HT.condemn
-                  return . errorResponseFor . singularityError $ queryTotal
-              else
-                return . responseLBS status201 headers $
-                  if iPreferRepresentation apiRequest == Full
-                    then toS body else ""
+              let (_, queryTotal, fs, body, gucHeaders) = row
+              case gucHeaders of
+                Left _ -> return . errorResponseFor $ GucHeadersError
+                Right hdrs -> do
+                  let headers = catMaybes [
+                          if null fs
+                            then Nothing
+                            else Just $ locationH tName fs
+                        , if iPreferRepresentation apiRequest == Full
+                            then Just $ toHeader contentType
+                            else Nothing
+                        , Just $ contentRangeH 1 0 $
+                            if shouldCount then Just queryTotal else Nothing
+                        , if null pkCols && isNothing (iOnConflict apiRequest)
+                            then Nothing
+                            else (\x -> ("Preference-Applied", show x)) <$> iPreferResolution apiRequest
+                        ] ++ (gucHToHeader <$> hdrs)
+                  if contentType == CTSingularJSON && queryTotal /= 1
+                    then do
+                      HT.condemn
+                      return . errorResponseFor . singularityError $ queryTotal
+                    else
+                      return . responseLBS status201 headers $ if iPreferRepresentation apiRequest == Full then toS body else mempty

        (ActionUpdate, TargetIdent (QualifiedIdentifier tSchema tName), Just pJson) ->
          case mutateSqlParts tSchema tName of
@@ -189,26 +191,27 @@ app dbStructure proc cols conf apiRequest =
            Right (sq, mq) -> do
              let stm = createWriteStatement sq mq
                    (contentType == CTSingularJSON) False (contentType == CTTextCSV)
-                    (iPreferRepresentation apiRequest) []
+                    (iPreferRepresentation apiRequest) [] pgVer
              row <- H.statement (toS $ pjRaw pJson) stm
-              let (_, queryTotal, _, body) = row
-                  updateIsNoOp = S.null cols
-                  contentRangeHeader = contentRangeH 0 (queryTotal - 1) $
-                        if shouldCount then Just queryTotal else Nothing
-                  headers | iPreferRepresentation apiRequest == Full = [toHeader contentType, contentRangeHeader]
-                          | otherwise                                = [contentRangeHeader]
-                  status | queryTotal == 0 && not updateIsNoOp      = status404
-                         | iPreferRepresentation apiRequest == Full = status200
-                         | otherwise                                = status204
-              if contentType == CTSingularJSON
-                 && queryTotal /= 1
-                then do
-                  HT.condemn
-                  return . errorResponseFor . singularityError $ queryTotal
-                else
-                  return $ if iPreferRepresentation apiRequest == Full
-                    then responseLBS status headers (toS body)
-                    else responseLBS status headers mempty
+              let (_, queryTotal, _, body, gucHeaders) = row
+              case gucHeaders of
+                Left _ -> return . errorResponseFor $ GucHeadersError
+                Right hdrs -> do
+                  let
+                    updateIsNoOp = S.null cols
+                    status | queryTotal == 0 && not updateIsNoOp      = status404
+                           | iPreferRepresentation apiRequest == Full = status200
+                           | otherwise                                = status204
+                    contentRangeHeader = contentRangeH 0 (queryTotal - 1) $ if shouldCount then Just queryTotal else Nothing
+                    headers = [contentRangeHeader] ++
+                              [if iPreferRepresentation apiRequest == Full then toHeader contentType else mempty] ++
+                              (gucHToHeader <$> hdrs)
+                  if contentType == CTSingularJSON && queryTotal /= 1
+                    then do
+                      HT.condemn
+                      return . errorResponseFor . singularityError $ queryTotal
+                    else
+                      return . responseLBS status headers $ if iPreferRepresentation apiRequest == Full then toS body else mempty

        (ActionSingleUpsert, TargetIdent (QualifiedIdentifier tSchema tName), Just ProcessedJSON{pjRaw, pjType, pjKeys}) ->
          case mutateSqlParts tSchema tName of
@@ -227,19 +230,22 @@ app dbStructure proc cols conf apiRequest =
              else do
                row <- H.statement (toS pjRaw) $
                       createWriteStatement sq mq (contentType == CTSingularJSON) False
-                                            (contentType == CTTextCSV) (iPreferRepresentation apiRequest) []
-                let (_, queryTotal, _, body) = row
-                -- Makes sure the querystring pk matches the payload pk
-                -- e.g. PUT /items?id=eq.1 { "id" : 1, .. } is accepted, PUT /items?id=eq.14 { "id" : 2, .. } is rejected
-                -- If this condition is not satisfied then nothing is inserted, check the WHERE for INSERT in QueryBuilder.hs to see how it's done
-                if queryTotal /= 1
-                  then do
-                    HT.condemn
-                    return . errorResponseFor $ PutMatchingPkError
-                  else
-                    return $ if iPreferRepresentation apiRequest == Full
-                      then responseLBS status200 [toHeader contentType] (toS body)
-                      else responseLBS status204 [] ""
+                                            (contentType == CTTextCSV) (iPreferRepresentation apiRequest) [] pgVer
+                let (_, queryTotal, _, body, gucHeaders) = row
+                case gucHeaders of
+                  Left _ -> return . errorResponseFor $ GucHeadersError
+                  Right hdrs -> do
+                    let headers = toHeader contentType : (gucHToHeader <$> hdrs)
+                        status = if iPreferRepresentation apiRequest == Full then status200 else status204
+                    -- Makes sure the querystring pk matches the payload pk
+                    -- e.g. PUT /items?id=eq.1 { "id" : 1, .. } is accepted, PUT /items?id=eq.14 { "id" : 2, .. } is rejected
+                    -- If this condition is not satisfied then nothing is inserted, check the WHERE for INSERT in QueryBuilder.hs to see how it's done
+                    if queryTotal /= 1
+                      then do
+                        HT.condemn
+                        return . errorResponseFor $ PutMatchingPkError
+                      else
+                        return . responseLBS status headers $ if iPreferRepresentation apiRequest == Full then toS body else mempty

        (ActionDelete, TargetIdent (QualifiedIdentifier tSchema tName), Nothing) ->
          case mutateSqlParts tSchema tName of
@@ -248,20 +254,25 @@ app dbStructure proc cols conf apiRequest =
              let stm = createWriteStatement sq mq
                    (contentType == CTSingularJSON) False
                    (contentType == CTTextCSV)
-                    (iPreferRepresentation apiRequest) []
+                    (iPreferRepresentation apiRequest) [] pgVer
              row <- H.statement mempty stm
-              let (_, queryTotal, _, body) = row
-                  contentRangeHeader = contentRangeH 1 0 $
-                        if shouldCount then Just queryTotal else Nothing
-              if contentType == CTSingularJSON
-                 && queryTotal /= 1
-                then do
-                  HT.condemn
-                  return . errorResponseFor . singularityError $ queryTotal
-                else
-                  return $ if iPreferRepresentation apiRequest == Full
-                    then responseLBS status200 [toHeader contentType, contentRangeHeader] (toS body)
-                    else responseLBS status204 [contentRangeHeader] ""
+              let (_, queryTotal, _, body, gucHeaders) = row
+              case gucHeaders of
+                Left _ -> return . errorResponseFor $ GucHeadersError
+                Right hdrs -> do
+                  let
+                    status = if iPreferRepresentation apiRequest == Full then status200 else status204
+                    contentRangeHeader = contentRangeH 1 0 $ if shouldCount then Just queryTotal else Nothing
+                    headers = [contentRangeHeader] ++
+                              [if iPreferRepresentation apiRequest == Full then toHeader contentType else mempty] ++
+                              (gucHToHeader <$> hdrs)
+                  if contentType == CTSingularJSON
+                     && queryTotal /= 1
+                    then do
+                      HT.condemn
+                      return . errorResponseFor . singularityError $ queryTotal
+                    else
+                      return . responseLBS status headers $ if iPreferRepresentation apiRequest == Full then toS body else mempty

        (ActionInfo, TargetIdent (QualifiedIdentifier tSchema tTable), Nothing) ->
          let mTable = find (\t -> tableName t == tTable && tableSchema t == tSchema) (dbTables dbStructure) in
@@ -282,7 +293,7 @@ app dbStructure proc cols conf apiRequest =
                pq = requestToCallProcQuery qi (specifiedProcArgs cols proc) returnsScalar preferParams
                stm = callProcStatement returnsScalar pq q cq shouldCount (contentType == CTSingularJSON)
                        (contentType == CTTextCSV) (contentType `elem` rawContentTypes) (preferParams == Just MultipleObjects)
-                        bField (pgVersion dbStructure)
+                        bField pgVer
              row <- H.statement (toS $ pjRaw pJson) stm
              let (tableTotal, queryTotal, body, gucHeaders) = row
                  (status, contentRange) = rangeStatusHeader topLevelRange queryTotal tableTotal
@@ -294,7 +305,7 @@ app dbStructure proc cols conf apiRequest =
                      HT.condemn
                      return . errorResponseFor . singularityError $ queryTotal
                    else
-                      return $ responseLBS status ([toHeader contentType, contentRange] ++ toHeaders hs)
+                      return $ responseLBS status ([toHeader contentType, contentRange] ++ (gucHToHeader <$> hs))
                      (if invMethod == InvHead then mempty else toS body)

        (ActionInspect headersOnly, TargetDefaultSpec tSchema, Nothing) -> do
@@ -325,6 +336,7 @@ app dbStructure proc cols conf apiRequest =
        shouldCount = exactCount || estimatedCount
        topLevelRange = iTopLevelRange apiRequest
        returnsScalar = maybe False procReturnsScalar proc
+        pgVer = pgVersion dbStructure

        readSqlParts s t =
          let
@@ -197,3 +197,9 @@ returningF qi returnings =
  if null returnings
    then "RETURNING 1" -- For mutation cases where there's no ?select, we return 1 to know how many rows were modified
    else "RETURNING " <> intercalate ", " (pgFmtColumn qi <$> returnings)
+
+responseHeadersF :: PgVersion -> SqlFragment
+responseHeadersF pgVer =
+  if pgVer >= pgVersion96
+    then "coalesce(nullif(current_setting('response.headers', true), ''), '[]')" :: Text -- nullif is used because of https://gist.github.com/steve-chavez/8d7033ea5655096903f3b52f8ed09a15
+    else "'[]'" :: Text
@@ -38,12 +38,12 @@ import           Text.InterpolatedString.Perl6   (qc)
    is represented as a list of strings containing variable bindings like
    @"k1=eq.42"@, or the empty list if there is no location header.
 -}
-type ResultsWithCount = (Maybe Int64, Int64, [BS.ByteString], BS.ByteString)
+type ResultsWithCount = (Maybe Int64, Int64, [BS.ByteString], BS.ByteString, Either Text [GucHeader])

 createWriteStatement :: SqlQuery -> SqlQuery -> Bool -> Bool -> Bool ->
-                        PreferRepresentation -> [Text] ->
+                        PreferRepresentation -> [Text] -> PgVersion ->
                        H.Statement ByteString ResultsWithCount
-createWriteStatement selectQuery mutateQuery wantSingle isInsert asCsv rep pKeys =
+createWriteStatement selectQuery mutateQuery wantSingle isInsert asCsv rep pKeys pgVer =
  unicodeStatement sql (param HE.unknown) decodeStandard True
 where
  sql = [qc|
@@ -53,7 +53,8 @@ createWriteStatement selectQuery mutateQuery wantSingle isInsert asCsv rep pKeys
        '' AS total_result_set,
        pg_catalog.count(_postgrest_t) AS page_total,
        {locF} AS header,
-        {bodyF} AS body
+        {bodyF} AS body,
+        {responseHeadersF pgVer} AS response_headers
      FROM ({selectQuery}) _postgrest_t |]

  locF =
@@ -73,11 +74,11 @@ createWriteStatement selectQuery mutateQuery wantSingle isInsert asCsv rep pKeys

  decodeStandard :: HD.Result ResultsWithCount
  decodeStandard =
-   fromMaybe (Nothing, 0, [], "") <$> HD.rowMaybe standardRow
+   fromMaybe (Nothing, 0, [], mempty, Right []) <$> HD.rowMaybe standardRow

-createReadStatement :: SqlQuery -> SqlQuery -> Bool -> Bool -> Bool -> Maybe FieldName ->
+createReadStatement :: SqlQuery -> SqlQuery -> Bool -> Bool -> Bool -> Maybe FieldName -> PgVersion ->
                       H.Statement () ResultsWithCount
-createReadStatement selectQuery countQuery isSingle countTotal asCsv binaryField =
+createReadStatement selectQuery countQuery isSingle countTotal asCsv binaryField pgVer =
  unicodeStatement sql HE.noParams decodeStandard False
 where
  sql = [qc|
@@ -88,7 +89,8 @@ createReadStatement selectQuery countQuery isSingle countTotal asCsv binaryField
        {countResultF} AS total_result_set,
        pg_catalog.count(_postgrest_t) AS page_total,
        {noLocationF} AS header,
-        {bodyF} AS body
+        {bodyF} AS body,
+        {responseHeadersF pgVer} AS response_headers
      FROM ( SELECT * FROM {sourceCTEName}) _postgrest_t |]

  (countCTEF, countResultF) = countF countQuery countTotal
@@ -108,8 +110,8 @@ createReadStatement selectQuery countQuery isSingle countTotal asCsv binaryField
    for that common type of query.
 -}
 standardRow :: HD.Row ResultsWithCount
-standardRow = (,,,) <$> nullableColumn HD.int8 <*> column HD.int8
-                    <*> column header <*> column HD.bytea
+standardRow = (,,,,) <$> nullableColumn HD.int8 <*> column HD.int8
+                    <*> column header <*> column HD.bytea <*> column decodeGucHeaders
  where
    header = HD.array $ HD.dimension replicateM $ element HD.bytea

@@ -128,7 +130,7 @@ callProcStatement returnsScalar callProcQuery selectQuery countQuery countTotal
        {countResultF} AS total_result_set,
        pg_catalog.count(_postgrest_t) AS page_total,
        {bodyF} AS body,
-        {responseHeaders} AS response_headers
+        {responseHeadersF pgVer} AS response_headers
      FROM ({selectQuery}) _postgrest_t;|]

    (countCTEF, countResultF) = countF countQuery countTotal
@@ -145,18 +147,12 @@ callProcStatement returnsScalar callProcQuery selectQuery countQuery countTotal
     | multObjects = "json_agg(_postgrest_t.pgrst_scalar)::character varying"
     | otherwise   = "(json_agg(_postgrest_t.pgrst_scalar)->0)::character varying"

-    responseHeaders =
-      if pgVer >= pgVersion96
-        then "coalesce(nullif(current_setting('response.headers', true), ''), '[]')" :: Text -- nullif is used because of https://gist.github.com/steve-chavez/8d7033ea5655096903f3b52f8ed09a15
-        else "'[]'" :: Text
-
    decodeProc :: HD.Result ProcResults
    decodeProc =
-      let row = fromMaybe (Just 0, 0, "[]", "[]") <$> HD.rowMaybe procRow in
-      (\(a, b, c, d) -> (a, b, c, first toS $ JSON.eitherDecode $ toS d)) <$> row
+      fromMaybe (Just 0, 0, mempty, Right []) <$> HD.rowMaybe procRow
      where
        procRow = (,,,) <$> nullableColumn HD.int8 <*> column HD.int8
-                        <*> column HD.bytea <*> column HD.bytea
+                        <*> column HD.bytea <*> column decodeGucHeaders

 createExplainStatement :: SqlQuery -> H.Statement () (Maybe Int64)
 createExplainStatement countQuery =
@@ -178,3 +174,6 @@ createExplainStatement countQuery =

 unicodeStatement :: Text -> HE.Params a -> HD.Result b -> Bool -> H.Statement a b
 unicodeStatement = H.Statement . encodeUtf8
+
+decodeGucHeaders :: HD.Value (Either Text [GucHeader])
+decodeGucHeaders = first toS . JSON.eitherDecode . toS <$> HD.bytea
@@ -408,8 +408,8 @@ instance JSON.FromJSON GucHeader where
    _ -> mzero
  parseJSON _          = mzero

-toHeaders :: [GucHeader] -> [Header]
-toHeaders = map $ \(GucHeader (k, v)) -> (CI.mk $ toS k, toS v)
+gucHToHeader :: GucHeader -> Header
+gucHToHeader (GucHeader (k, v)) = (CI.mk $ toS k, toS v)

 {-|
  This type will hold information about which particular 'Relation' between two tables to choose when there are multiple ones.