From 15e04903b1805830b1773a833d8d5ed00dbcca01 Mon Sep 17 00:00:00 2001 From: steve-chavez Date: Wed, 10 Sep 2025 19:13:52 -0500 Subject: [PATCH] fix: not logging explain query for estimated count Fixes https://github.com/PostgREST/postgrest/issues/4319 --- CHANGELOG.md | 1 + src/PostgREST/ApiRequest/Preferences.hs | 5 ++ src/PostgREST/Logger.hs | 2 +- src/PostgREST/Query.hs | 67 +++++++++++-------------- test/io/test_io.py | 43 +++++++++++++--- 5 files changed, 71 insertions(+), 47 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d2afeb223..16ba25da7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ This project adheres to [Semantic Versioning](http://semver.org/). ### Fixed - Fix not logging OpenAPI queries when `log-query=main-query` is enabled by @steve-chavez in #4226 +- Fix not logging explain query when `log-query=main-query` is enabled by @steve-chavez in #4319 ### Added diff --git a/src/PostgREST/ApiRequest/Preferences.hs b/src/PostgREST/ApiRequest/Preferences.hs index 086cb2fdc..55369efbb 100644 --- a/src/PostgREST/ApiRequest/Preferences.hs +++ b/src/PostgREST/ApiRequest/Preferences.hs @@ -19,6 +19,7 @@ module PostgREST.ApiRequest.Preferences , PreferMaxAffected(..) , fromHeaders , shouldCount + , shouldExplainCount , prefAppliedHeader ) where @@ -238,6 +239,10 @@ shouldCount :: Maybe PreferCount -> Bool shouldCount prefCount = prefCount == Just ExactCount || prefCount == Just EstimatedCount +shouldExplainCount :: Maybe PreferCount -> Bool +shouldExplainCount prefCount = + prefCount == Just PlannedCount || prefCount == Just EstimatedCount + -- | Whether to commit or roll back transactions. data PreferTransaction = Commit -- ^ Commit transaction - the default. diff --git a/src/PostgREST/Logger.hs b/src/PostgREST/Logger.hs index de71b48ca..c091b8026 100644 --- a/src/PostgREST/Logger.hs +++ b/src/PostgREST/Logger.hs @@ -123,7 +123,7 @@ logWithZTime loggerState txt = do logMainQ :: LoggerState -> MainQuery -> IO () logMainQ loggerState MainQuery{mqOpenAPI=(x, y, z),..} = - let snipts = renderSnippet <$> [mqMain, x, y, z] + let snipts = renderSnippet <$> [mqMain, x, y, z, fromMaybe mempty mqExplain] -- Does not log SQL when it's empty (happens on OPTIONS requests and when the openapi queries are not generated) logQ q = when (q /= mempty) $ logWithZTime loggerState $ showOnSingleLine '\n' $ T.decodeUtf8 q in mapM_ logQ snipts diff --git a/src/PostgREST/Query.hs b/src/PostgREST/Query.hs index 6132f2afd..a5655dfb5 100644 --- a/src/PostgREST/Query.hs +++ b/src/PostgREST/Query.hs @@ -46,7 +46,8 @@ import PostgREST.ApiRequest.Preferences (PreferCount (..), PreferHandling (..), PreferMaxAffected (..), PreferTransaction (..), - Preferences (..)) + Preferences (..), + shouldExplainCount) import PostgREST.ApiRequest.Types (Mutation (..)) import PostgREST.Auth.Types (AuthResult (..)) import PostgREST.Config (AppConfig (..), @@ -86,10 +87,10 @@ data QueryResult data MainQuery = MainQuery { mqTxVars :: SQL.Snippet -- ^ the transaction variables that always run on each query , mqPreReq :: Maybe SQL.Snippet -- ^ the pre-request function that runs if enabled - , mqCount :: SQL.Snippet -- ^ this count query is actually a fragment of the main query, but the same count query also runs after the main one in the case of `count=estimated`, so it's cached here. -- TODO only one of the following queries actually runs on each request, once OpenAPI is removed from core it will be easier to refactor this , mqMain :: SQL.Snippet , mqOpenAPI :: (SQL.Snippet, SQL.Snippet, SQL.Snippet) + , mqExplain :: Maybe SQL.Snippet -- ^ the explain query that gets generated for the "Prefer: count=estimated" case } -- | Standard result set format used for the mqMain query @@ -142,35 +143,51 @@ planIsoLvl AppConfig{configRoleIsoLvl} role actPlan = case actPlan of mainQuery :: ActionPlan -> AppConfig -> ApiRequest -> AuthResult -> Maybe QualifiedIdentifier -> MainQuery -mainQuery (NoDb _) _ _ _ _ = MainQuery mempty Nothing mempty mempty (mempty, mempty, mempty) +mainQuery (NoDb _) _ _ _ _ = MainQuery mempty Nothing mempty (mempty, mempty, mempty) mempty mainQuery (Db plan) conf@AppConfig{..} apiReq@ApiRequest{iPreferences=Preferences{..}} authRes preReq = let genQ = MainQuery (PreQuery.txVarQuery plan conf authRes apiReq) (PreQuery.preReqQuery <$> preReq) in case plan of DbCrud _ WrappedReadPlan{..} -> let countQuery = QueryBuilder.readPlanToCountQuery wrReadPlan in - genQ countQuery (Statements.mainRead wrReadPlan countQuery preferCount configDbMaxRows pMedia wrHandler) (mempty, mempty, mempty) + genQ (Statements.mainRead wrReadPlan countQuery preferCount configDbMaxRows pMedia wrHandler) (mempty, mempty, mempty) + (if shouldExplainCount preferCount then Just (Statements.postExplain countQuery) else Nothing) DbCrud _ MutateReadPlan{..} -> - genQ mempty (Statements.mainWrite mrReadPlan mrMutatePlan pMedia mrHandler preferRepresentation preferResolution) (mempty, mempty, mempty) + genQ (Statements.mainWrite mrReadPlan mrMutatePlan pMedia mrHandler preferRepresentation preferResolution) (mempty, mempty, mempty) mempty DbCrud _ CallReadPlan{..} -> - genQ mempty (Statements.mainCall crProc crCallPlan crReadPlan preferCount pMedia crHandler) (mempty, mempty, mempty) + genQ (Statements.mainCall crProc crCallPlan crReadPlan preferCount pMedia crHandler) (mempty, mempty, mempty) mempty MayUseDb InspectPlan{ipSchema=tSchema} -> - genQ mempty mempty (SqlFragment.accessibleTables tSchema, SqlFragment.accessibleFuncs tSchema, SqlFragment.schemaDescription tSchema) + genQ mempty (SqlFragment.accessibleTables tSchema, SqlFragment.accessibleFuncs tSchema, SqlFragment.schemaDescription tSchema) mempty -- TODO: Generate the Hasql Statement in a diferent module after the OpenAPI functionality is removed -actionQuery :: MainQuery -> DbActionPlan -> AppConfig -> ApiRequest -> SchemaCache -> DbHandler QueryResult +actionQuery :: MainQuery -> DbActionPlan -> AppConfig -> ApiRequest -> SchemaCache -> ExceptT Error SQL.Transaction QueryResult actionQuery MainQuery{..} (DbCrud True plan) conf@AppConfig{..} apiReq _ = do explRes <- lift $ SQL.statement mempty $ SQL.dynamicallyParameterized mqMain planRow configDbPreparedStatements optionalRollback conf apiReq pure $ DbPlanResult (pMedia plan) explRes -actionQuery MainQuery{..} (DbCrud _ plan@WrappedReadPlan{..}) conf@AppConfig{..} apiReq _ = do - resultSet <- lift $ SQL.statement mempty $ dynStmt (HD.singleRow $ standardRow True) +actionQuery MainQuery{..} (DbCrud _ plan@WrappedReadPlan{..}) conf@AppConfig{..} apiReq@ApiRequest{iPreferences=Preferences{..}} _ = do + resultSet@RSStandard{rsTableTotal=tableTotal} <- lift $ SQL.statement mempty $ dynStmt (HD.singleRow $ standardRow True) failNotSingular pMedia resultSet optionalRollback conf apiReq - DbCrudResult plan <$> resultSetWTotal conf apiReq resultSet mqCount + explainTotal <- lift . fmap join $ traverse (\snip -> + SQL.statement mempty $ SQL.dynamicallyParameterized snip decodeExplain configDbPreparedStatements) + mqExplain + + pure $ DbCrudResult plan + resultSet{rsTableTotal=case preferCount of + Just PlannedCount -> explainTotal + Just EstimatedCount -> if tableTotal > (fromIntegral <$> configDbMaxRows) + then max <$> tableTotal <*> explainTotal + else tableTotal + _ -> tableTotal} where dynStmt decod = SQL.dynamicallyParameterized mqMain decod configDbPreparedStatements + decodeExplain :: HD.Result (Maybe Int64) + decodeExplain = + let row = HD.singleRow $ column HD.bytea in + (^? L.nth 0 . L.key "Plan" . L.key "Plan Rows" . L._Integral) <$> row + actionQuery MainQuery{..} (DbCrud _ plan@MutateReadPlan{..}) conf@AppConfig{..} apiReq@ApiRequest{iPreferences=Preferences{..}} _ = do resultSet <- lift $ SQL.statement mempty $ dynStmt decodeRow failMutation resultSet @@ -246,34 +263,6 @@ failPut RSStandard{rsQueryTotal=queryTotal} = lift SQL.condemn throwError $ Error.ApiRequestError Error.PutMatchingPkError -resultSetWTotal :: AppConfig -> ApiRequest -> ResultSet -> SQL.Snippet -> DbHandler ResultSet -resultSetWTotal AppConfig{..} ApiRequest{iPreferences=Preferences{..}} rs@RSStandard{rsTableTotal=tableTotal} countQuery = - case preferCount of - Just PlannedCount -> do - total <- explain - return rs{rsTableTotal=total} - Just EstimatedCount -> - if tableTotal > (fromIntegral <$> configDbMaxRows) then do - total <- max tableTotal <$> explain - return rs{rsTableTotal=total} - else - return rs - Just ExactCount -> - return rs - Nothing -> - return rs - where - explain = - lift . SQL.statement mempty $ - SQL.dynamicallyParameterized (Statements.postExplain 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. diff --git a/test/io/test_io.py b/test/io/test_io.py index b07efc6c2..ca5db1ace 100644 --- a/test/io/test_io.py +++ b/test/io/test_io.py @@ -1029,6 +1029,16 @@ def test_log_query(level, defaultenv): response = postgrest.session.get("/projects") assert response.status_code == 200 + response = postgrest.session.get( + "/projects", headers={"Prefer": "count=estimated"} + ) + assert response.status_code == 206 + + response = postgrest.session.get( + "/projects", headers={"Prefer": "count=planned"} + ) + assert response.status_code == 206 + response = postgrest.session.get("/infinite_recursion") assert response.status_code == 500 @@ -1036,6 +1046,9 @@ def test_log_query(level, defaultenv): root_2xx_regx_ln2 = r".+: WITH base_types AS \(.+\) SELECT pn.nspname AS proc_schema, .+ FROM pg_proc p.+AND p.pronamespace = \$1::regnamespace" root_2xx_regx_ln3 = r".+: SELECT pg_catalog\.obj_description\(\$1::regnamespace, 'pg_namespace'\)" get_2xx_regx = r'.+: WITH pgrst_source AS.+SELECT "public"\."projects"\.\* FROM "public"\."projects".+_postgrest_t' + get_2xx_count_regx = ( + r'.+: EXPLAIN \(FORMAT JSON\) SELECT 1 FROM "public"."projects"' + ) infinite_recursion_5xx_regx = r'.+: WITH pgrst_source AS.+SELECT "public"\."infinite_recursion"\.\* FROM "public"\."infinite_recursion".+_postgrest_t' if level == "crit": @@ -1050,13 +1063,21 @@ def test_log_query(level, defaultenv): assert re.match(infinite_recursion_5xx_regx, output[1]) assert len(output) == 2 elif level == "info": - output = postgrest.read_stdout(nlines=8) - assert re.match(root_2xx_regx_ln1, output[0]) - assert re.match(root_2xx_regx_ln2, output[1]) - assert re.match(root_2xx_regx_ln3, output[2]) - assert re.match(get_2xx_regx, output[4]) - assert re.match(infinite_recursion_5xx_regx, output[7]) - assert len(output) == 8 + output_root = postgrest.read_stdout(nlines=4) + assert re.match(root_2xx_regx_ln1, output_root[0]) + assert re.match(root_2xx_regx_ln2, output_root[1]) + assert re.match(root_2xx_regx_ln3, output_root[2]) + assert len(output_root) == 4 + output_get = postgrest.read_stdout(nlines=8) + assert re.match(get_2xx_regx, output_get[0]) + assert re.match(get_2xx_regx, output_get[2]) + assert re.match(get_2xx_count_regx, output_get[3]) + assert re.match(get_2xx_regx, output_get[5]) + assert re.match(get_2xx_count_regx, output_get[6]) + assert len(output_get) == 8 + output_err = postgrest.read_stdout(nlines=2) + assert re.match(infinite_recursion_5xx_regx, output_err[1]) + assert len(output_err) == 2 elif level == "debug": output_root = postgrest.read_stdout(nlines=8) assert re.match(root_2xx_regx_ln1, output_root[4]) @@ -1066,6 +1087,14 @@ def test_log_query(level, defaultenv): output_get = postgrest.read_stdout(nlines=6) assert re.match(get_2xx_regx, output_get[4]) assert len(output_get) == 6 + output_get_estimated = postgrest.read_stdout(nlines=7) + assert re.match(get_2xx_regx, output_get_estimated[4]) + assert re.match(get_2xx_count_regx, output_get_estimated[5]) + assert len(output_get_estimated) == 7 + output_get_planned = postgrest.read_stdout(nlines=7) + assert re.match(get_2xx_regx, output_get_planned[4]) + assert re.match(get_2xx_count_regx, output_get_planned[5]) + assert len(output_get_planned) == 7 output_err = postgrest.read_stdout(nlines=6) assert re.match(infinite_recursion_5xx_regx, output_err[5]) assert len(output_err) == 6