refactor: move decoding to Query.hs

Towards solving https://github.com/PostgREST/postgrest/issues/3934. This makes the Statements internal module pure now it only returns Hasql.Snippet.
2025-09-02 16:26:15 -05:00
parent 00d92aa474
commit 3e776213b7
3 changed files with 106 additions and 114 deletions
@@ -1,22 +1,35 @@
 {-# LANGUAGE NamedFieldPuns  #-}
 {-# LANGUAGE RecordWildCards #-}
-- TODO: This module shouldn't depend on SchemaCache
+{-|
+Module      : PostgREST.Query
+Description : PostgREST query executor
+
+This module parametrizes, prepares, executes SQL queries and decodes their results.
+
+TODO: This module shouldn't depend on SchemaCache
+-}
 module PostgREST.Query
  ( Query (..)
  , QueryResult (..)
+  , ResultSet (..)
  , query
  , getSQLQuery
  ) where

+import           Control.Lens                      ((^?))
 import qualified Data.Aeson                        as JSON
 import qualified Data.Aeson.KeyMap                 as KM
-import qualified Data.ByteString                   as BS
+import qualified Data.Aeson.Lens                   as L
+import qualified Data.ByteString                   as BS hiding
+                                                         (break)
+import qualified Data.ByteString.Char8             as BS
 import qualified Data.HashMap.Strict               as HM
 import qualified Data.Set                          as S
 import qualified Hasql.Decoders                    as HD
-import qualified Hasql.DynamicStatements.Snippet   as SQL (Snippet)
+import qualified Hasql.DynamicStatements.Snippet   as SQL hiding (sql)
 import qualified Hasql.DynamicStatements.Statement as SQL
 import qualified Hasql.Session                     as SQL (Session)
+import qualified Hasql.Statement                   as SQL
 import qualified Hasql.Transaction                 as SQL
 import qualified Hasql.Transaction.Sessions        as SQL

@@ -47,7 +60,6 @@ import PostgREST.Plan                    (ActionPlan (..),
                                          InfoPlan (..),
                                          InspectPlan (..))
 import PostgREST.Plan.MutatePlan         (MutatePlan (..))
-import PostgREST.Query.Statements        (ResultSet (..))
 import PostgREST.SchemaCache             (SchemaCache (..))
 import PostgREST.SchemaCache.Identifiers (QualifiedIdentifier (..))
 import PostgREST.SchemaCache.Routine     (Routine (..), RoutineMap)
@@ -73,6 +85,27 @@ data QueryResult
  | MaybeDbResult InspectPlan  (Maybe (TablesMap, RoutineMap, Maybe Text))
  | NoDbResult    InfoPlan

+-- | Standard result set format used for all queries
+data ResultSet
+  = RSStandard
+  { rsTableTotal :: Maybe Int64
+  -- ^ count of all the table rows
+  , rsQueryTotal :: Int64
+  -- ^ count of the query rows
+  , rsLocation   :: [(BS.ByteString, BS.ByteString)]
+  -- ^ The Location header(only used for inserts) is represented as a list of strings containing
+  -- variable bindings like @"k1=eq.42"@, or the empty list if there is no location header.
+  , rsBody       :: BS.ByteString
+  -- ^ the aggregated body of the query
+  , rsGucHeaders :: Maybe BS.ByteString
+  -- ^ the HTTP headers to be added to the response
+  , rsGucStatus  :: Maybe Text
+  -- ^ the HTTP status to be added to the response
+  , rsInserted   :: Maybe Int64
+  -- ^ the number of rows inserted (Only used for upserts)
+  }
+  | RSPlan BS.ByteString -- ^ the plan of the query
+
 query :: AppConfig -> AuthResult -> ApiRequest -> ActionPlan -> SchemaCache -> Query
 query _ _ _ (NoDb x) _ = NoDbQuery $ NoDbResult x
 query config AuthResult{..} apiReq (Db plan) sCache =
@@ -106,7 +139,7 @@ actionQuery (DbCrud plan@WrappedReadPlan{..}) conf@AppConfig{..} apiReq@ApiReque
  (mainActionQuery, mainSQLQuery)
  where
    countQuery = QueryBuilder.readPlanToCountQuery wrReadPlan
-    (result, mainSQLQuery) = Statements.prepareRead
+    result@(SQL.Statement mainSQLQuery _ _ _) = SQL.dynamicallyParameterized (Statements.prepareRead
      (QueryBuilder.readPlanToQuery wrReadPlan)
      (if preferCount == Just EstimatedCount then
         -- LIMIT maxRows + 1 so we can determine below that maxRows was surpassed
@@ -117,18 +150,23 @@ actionQuery (DbCrud plan@WrappedReadPlan{..}) conf@AppConfig{..} apiReq@ApiReque
      (shouldCount preferCount)
      wrMedia
      wrHandler
-      configDbPreparedStatements
+      ) decodeIt configDbPreparedStatements
    mainActionQuery = do
      resultSet <- lift $ SQL.statement mempty result
      failNotSingular wrMedia resultSet
      optionalRollback conf apiReq
      DbCrudResult plan <$> resultSetWTotal conf apiReq resultSet countQuery

+    decodeIt :: HD.Result ResultSet
+    decodeIt = case wrMedia of
+      MTVndPlan{} -> planRow
+      _           -> HD.singleRow $ standardRow True
+
 actionQuery (DbCrud plan@MutateReadPlan{..}) conf@AppConfig{..} apiReq@ApiRequest{iPreferences=Preferences{..}} _ =
  (mainActionQuery, mainSQLQuery)
  where
    (isPut, isInsert, pkCols) = case mrMutatePlan of {Insert{where_,insPkCols} -> ((not . null) where_, True, insPkCols); _ -> (False,False, mempty);}
-    (result, mainSQLQuery) = Statements.prepareWrite
+    result@(SQL.Statement mainSQLQuery _ _ _) = SQL.dynamicallyParameterized (Statements.prepareWrite
      (QueryBuilder.readPlanToQuery mrReadPlan)
      (QueryBuilder.mutatePlanToQuery mrMutatePlan)
      isInsert
@@ -137,8 +175,7 @@ actionQuery (DbCrud plan@MutateReadPlan{..}) conf@AppConfig{..} apiReq@ApiReques
      mrHandler
      preferRepresentation
      preferResolution
-      pkCols
-      configDbPreparedStatements
+      pkCols) decodeIt configDbPreparedStatements
    failMutation resultSet = case mrMutation of
      MutationCreate -> do
        failNotSingular mrMedia resultSet
@@ -156,18 +193,23 @@ actionQuery (DbCrud plan@MutateReadPlan{..}) conf@AppConfig{..} apiReq@ApiReques
      optionalRollback conf apiReq
      pure $ DbCrudResult plan resultSet

+    decodeIt :: HD.Result ResultSet
+    decodeIt = case mrMedia of
+      MTVndPlan{} -> planRow
+      _           -> fromMaybe (RSStandard Nothing 0 mempty mempty Nothing Nothing Nothing) <$> HD.rowMaybe (standardRow False)
+
 actionQuery (DbCall plan@CallReadPlan{..}) conf@AppConfig{..} apiReq@ApiRequest{iPreferences=Preferences{..}} _ =
  (mainActionQuery, mainSQLQuery)
  where
-    (result, mainSQLQuery) = Statements.prepareCall
+    result@(SQL.Statement mainSQLQuery _ _ _) = SQL.dynamicallyParameterized (Statements.prepareCall
      crProc
      (QueryBuilder.callPlanToQuery crCallPlan)
      (QueryBuilder.readPlanToQuery crReadPlan)
      (QueryBuilder.readPlanToCountQuery crReadPlan)
      (shouldCount preferCount)
      crMedia
-      crHandler
-      configDbPreparedStatements
+      crHandler) decodeIt configDbPreparedStatements
+
    mainActionQuery = do
      resultSet <- lift $ SQL.statement mempty result
      optionalRollback conf apiReq
@@ -175,6 +217,11 @@ actionQuery (DbCall plan@CallReadPlan{..}) conf@AppConfig{..} apiReq@ApiRequest{
      failExceedsMaxAffectedPref (preferMaxAffected,preferHandling) resultSet
      pure $ DbCallResult plan resultSet

+    decodeIt :: HD.Result ResultSet
+    decodeIt = case crMedia of
+      MTVndPlan{} -> planRow
+      _           -> fromMaybe (RSStandard (Just 0) 0 mempty mempty Nothing Nothing Nothing) <$> HD.rowMaybe (standardRow True)
+
 actionQuery (MaybeDb plan@InspectPlan{ipSchema=tSchema}) AppConfig{..} _ sCache =
  (mainActionQuery, mempty)
  where
@@ -225,9 +272,16 @@ resultSetWTotal AppConfig{..} ApiRequest{iPreferences=Preferences{..}} rs@RSStan
      return rs
  where
    explain =
-      lift . SQL.statement mempty . Statements.preparePlanRows countQuery $
+      lift . SQL.statement mempty $
+        SQL.dynamicallyParameterized (Statements.preparePlanRows countQuery)
+        decodeIt
        configDbPreparedStatements

+    decodeIt :: HD.Result (Maybe Int64)
+    decodeIt =
+      let row = HD.singleRow $ column HD.bytea in
+      (^? L.nth 0 . L.key "Plan" .  L.key "Plan Rows" . L._Integral) <$> row
+
 -- |
 -- Fail a response if a single JSON object was requested and not exactly one
 -- was found.
@@ -274,3 +328,30 @@ runPreReqQuery conf = lift $ traverse_ (SQL.statement mempty . stmt) (configDbPr
 getSQLQuery :: Query -> ByteString
 getSQLQuery DbQuery{dqSQL} = dqSQL
 getSQLQuery _              = mempty
+
+-- | We use rowList because when doing EXPLAIN (FORMAT TEXT), the result comes as many rows. FORMAT JSON comes as one.
+planRow :: HD.Result ResultSet
+planRow = RSPlan . BS.unlines <$> HD.rowList (column HD.bytea)
+
+column :: HD.Value a -> HD.Row a
+column = HD.column . HD.nonNullable
+
+nullableColumn :: HD.Value a -> HD.Row (Maybe a)
+nullableColumn = HD.column . HD.nullable
+
+arrayColumn :: HD.Value a -> HD.Row [a]
+arrayColumn = column . HD.listArray . HD.nonNullable
+
+standardRow :: Bool -> HD.Row ResultSet
+standardRow noLocation =
+  RSStandard <$> nullableColumn HD.int8 <*> column HD.int8
+             <*> (if noLocation then pure mempty else fmap splitKeyValue <$> arrayColumn HD.bytea)
+             <*> (fromMaybe mempty <$> nullableColumn HD.bytea)
+             <*> nullableColumn HD.bytea
+             <*> nullableColumn HD.text
+             <*> nullableColumn HD.int8
+  where
+    splitKeyValue :: ByteString -> (ByteString, ByteString)
+    splitKeyValue kv =
+      let (k, v) = BS.break (== '=') kv in
+      (k, BS.tail v)
@@ -1,28 +1,15 @@
 {-|
 Module      : PostgREST.Query.Statements
-Description : PostgREST single SQL statements.
-
-This module constructs single SQL statements that can be parametrized and prepared.
-
- It consumes the SqlQuery types generated by the QueryBuilder module.
- It generates the body format and some headers of the final HTTP response.
+Description : PostgREST SQL statements.
 -}
 module PostgREST.Query.Statements
  ( prepareWrite
  , prepareRead
  , prepareCall
  , preparePlanRows
-  , ResultSet (..)
  ) where

-import qualified Data.Aeson.Lens                   as L
-import qualified Data.ByteString.Char8             as BS
-import qualified Hasql.Decoders                    as HD
-import qualified Hasql.DynamicStatements.Snippet   as SQL
-import qualified Hasql.DynamicStatements.Statement as SQL
-import qualified Hasql.Statement                   as SQL
-
-import Control.Lens ((^?))
+import qualified Hasql.DynamicStatements.Snippet as SQL

 import PostgREST.ApiRequest.Preferences
 import PostgREST.MediaType              (MTVndPlanFormat (..),
@@ -33,34 +20,10 @@ import PostgREST.SchemaCache.Routine    (MediaHandler (..), Routine,

 import Protolude

-- | Standard result set format used for all queries
-data ResultSet
-  = RSStandard
-  { rsTableTotal :: Maybe Int64
-  -- ^ count of all the table rows
-  , rsQueryTotal :: Int64
-  -- ^ count of the query rows
-  , rsLocation   :: [(BS.ByteString, BS.ByteString)]
-  -- ^ The Location header(only used for inserts) is represented as a list of strings containing
-  -- variable bindings like @"k1=eq.42"@, or the empty list if there is no location header.
-  , rsBody       :: BS.ByteString
-  -- ^ the aggregated body of the query
-  , rsGucHeaders :: Maybe BS.ByteString
-  -- ^ the HTTP headers to be added to the response
-  , rsGucStatus  :: Maybe Text
-  -- ^ the HTTP status to be added to the response
-  , rsInserted   :: Maybe Int64
-  -- ^ the number of rows inserted (Only used for upserts)
-  }
-  | RSPlan BS.ByteString -- ^ the plan of the query
-
-
 prepareWrite :: SQL.Snippet -> SQL.Snippet -> Bool -> Bool -> MediaType -> MediaHandler ->
-                Maybe PreferRepresentation -> Maybe PreferResolution -> [Text] -> Bool -> (SQL.Statement () ResultSet, ByteString)
-prepareWrite selectQuery mutateQuery isInsert isPut mt handler rep resolution pKeys prepared =
- (result, sql)
+                Maybe PreferRepresentation -> Maybe PreferResolution -> [Text] -> SQL.Snippet
+prepareWrite selectQuery mutateQuery isInsert isPut mt handler rep resolution pKeys = mtSnippet mt snippet
 where
-  result@(SQL.Statement sql _ _ _) = SQL.dynamicallyParameterized (mtSnippet mt snippet) decodeIt prepared
  checkUpsert snip = if isInsert && (isPut || resolution == Just MergeDuplicates) then snip else "''"
  pgrstInsertedF = checkUpsert "nullif(current_setting('pgrst.inserted', true),'')::int"
  snippet =
@@ -89,16 +52,9 @@ prepareWrite selectQuery mutateQuery isInsert isPut mt handler rep resolution pK
    | handler == NoAgg = "SELECT * FROM " <> sourceCTE
    | otherwise        = selectQuery

-  decodeIt :: HD.Result ResultSet
-  decodeIt = case mt of
-    MTVndPlan{} -> planRow
-    _           -> fromMaybe (RSStandard Nothing 0 mempty mempty Nothing Nothing Nothing) <$> HD.rowMaybe (standardRow False)
-
-prepareRead :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType -> MediaHandler -> Bool -> (SQL.Statement () ResultSet, ByteString)
-prepareRead selectQuery countQuery countTotal mt handler prepared =
- (result, sql)
+prepareRead :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType -> MediaHandler -> SQL.Snippet
+prepareRead selectQuery countQuery countTotal mt handler = mtSnippet mt snippet
 where
-  result@(SQL.Statement sql _ _ _) = SQL.dynamicallyParameterized (mtSnippet mt snippet) decodeIt prepared
  snippet =
    "WITH " <> sourceCTE <> " AS ( " <> selectQuery <> " ) " <>
    countCTEF <> " " <>
@@ -113,18 +69,11 @@ prepareRead selectQuery countQuery countTotal mt handler prepared =

  (countCTEF, countResultF) = countF countQuery countTotal

-  decodeIt :: HD.Result ResultSet
-  decodeIt = case mt of
-    MTVndPlan{} -> planRow
-    _           -> HD.singleRow $ standardRow True

 prepareCall :: Routine -> SQL.Snippet -> SQL.Snippet -> SQL.Snippet -> Bool ->
-               MediaType -> MediaHandler -> Bool ->
-               (SQL.Statement () ResultSet, ByteString)
-prepareCall rout callProcQuery selectQuery countQuery countTotal mt handler prepared =
-  (result, sql)
+               MediaType -> MediaHandler -> SQL.Snippet
+prepareCall rout callProcQuery selectQuery countQuery countTotal mt handler = mtSnippet mt snippet
  where
-    result@(SQL.Statement sql _ _ _) = SQL.dynamicallyParameterized (mtSnippet mt snippet) decodeIt prepared
    snippet =
      "WITH " <> sourceCTE <> " AS (" <> callProcQuery <> ") " <>
      countCTEF <>
@@ -141,49 +90,11 @@ prepareCall rout callProcQuery selectQuery countQuery countTotal mt handler prep

    (countCTEF, countResultF) = countF countQuery countTotal

-    decodeIt :: HD.Result ResultSet
-    decodeIt = case mt of
-      MTVndPlan{} -> planRow
-      _           -> fromMaybe (RSStandard (Just 0) 0 mempty mempty Nothing Nothing Nothing) <$> HD.rowMaybe (standardRow True)

-preparePlanRows :: SQL.Snippet -> Bool -> SQL.Statement () (Maybe Int64)
-preparePlanRows countQuery =
-  SQL.dynamicallyParameterized snippet decodeIt
-  where
-    snippet = explainF PlanJSON mempty countQuery
-    decodeIt :: HD.Result (Maybe Int64)
-    decodeIt =
-      let row = HD.singleRow $ column HD.bytea in
-      (^? L.nth 0 . L.key "Plan" .  L.key "Plan Rows" . L._Integral) <$> row
-
-standardRow :: Bool -> HD.Row ResultSet
-standardRow noLocation =
-  RSStandard <$> nullableColumn HD.int8 <*> column HD.int8
-             <*> (if noLocation then pure mempty else fmap splitKeyValue <$> arrayColumn HD.bytea)
-             <*> (fromMaybe mempty <$> nullableColumn HD.bytea)
-             <*> nullableColumn HD.bytea
-             <*> nullableColumn HD.text
-             <*> nullableColumn HD.int8
-  where
-    splitKeyValue :: ByteString -> (ByteString, ByteString)
-    splitKeyValue kv =
-      let (k, v) = BS.break (== '=') kv in
-      (k, BS.tail v)
+preparePlanRows :: SQL.Snippet -> SQL.Snippet
+preparePlanRows = explainF PlanJSON mempty

 mtSnippet :: MediaType -> SQL.Snippet -> SQL.Snippet
 mtSnippet mediaType snippet = case mediaType of
  MTVndPlan _ fmt opts -> explainF fmt opts snippet
  _                    -> snippet
-
-- | We use rowList because when doing EXPLAIN (FORMAT TEXT), the result comes as many rows. FORMAT JSON comes as one.
-planRow :: HD.Result ResultSet
-planRow = RSPlan . BS.unlines <$> HD.rowList (column HD.bytea)
-
-column :: HD.Value a -> HD.Row a
-column = HD.column . HD.nonNullable
-
-nullableColumn :: HD.Value a -> HD.Row (Maybe a)
-nullableColumn = HD.column . HD.nullable
-
-arrayColumn :: HD.Value a -> HD.Row [a]
-arrayColumn = column . HD.listArray . HD.nonNullable
@@ -40,8 +40,8 @@ import PostgREST.Plan                    (CallReadPlan (..),
                                          InfoPlan (..),
                                          InspectPlan (..))
 import PostgREST.Plan.MutatePlan         (MutatePlan (..))
-import PostgREST.Query                   (QueryResult (..))
-import PostgREST.Query.Statements        (ResultSet (..))
+import PostgREST.Query                   (QueryResult (..),
+                                          ResultSet (..))
 import PostgREST.Response.GucHeader      (GucHeader, unwrapGucHeader)
 import PostgREST.SchemaCache             (SchemaCache (..))
 import PostgREST.SchemaCache.Identifiers (QualifiedIdentifier (..),