From 1bbeae4290d97997aebed45f2a8b36b32ef7dadd Mon Sep 17 00:00:00 2001 From: Anthonin Bonnefoy Date: Fri, 26 Jul 2024 10:18:02 +0200 Subject: Create module for query_id test There are limited possibilities to check whether a query_id has been set correctly. - Checking pg_stat_activity is not possible in the regress tests as you need a second session to check the reported query_id. - pg_stat_statements can be used indirectly but you're limited to how pgss uses query_id. For example, it doesn't rely on queryId in ExecutorRun. To fix this, this patch introduces a new test module to check queryId. This module install a check to all parse, plan and executor hooks and assert that queryId is set. On top of that, when at the root level, it checks that queryId reported by pgstat is the same as the processed query. --- src/test/modules/Makefile | 1 + src/test/modules/meson.build | 1 + src/test/modules/test_query_id/Makefile | 20 ++ .../test_query_id/expected/test_query_id.out | 51 +++++ src/test/modules/test_query_id/meson.build | 28 +++ .../test_query_id/sql/test_query_id.sql | 46 +++++ .../modules/test_query_id/test_query_id.c | 185 ++++++++++++++++++ 7 files changed, 332 insertions(+) create mode 100644 src/test/modules/test_query_id/Makefile create mode 100644 src/test/modules/test_query_id/expected/test_query_id.out create mode 100644 src/test/modules/test_query_id/meson.build create mode 100644 src/test/modules/test_query_id/sql/test_query_id.sql create mode 100644 src/test/modules/test_query_id/test_query_id.c diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile index 256799f520a..a2391c685a5 100644 --- a/src/test/modules/Makefile +++ b/src/test/modules/Makefile @@ -29,6 +29,7 @@ SUBDIRS = \ test_parser \ test_pg_dump \ test_predtest \ + test_query_id \ test_radixtree \ test_rbtree \ test_regex \ diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build index d8fe059d236..cece5ff85fb 100644 --- a/src/test/modules/meson.build +++ b/src/test/modules/meson.build @@ -28,6 +28,7 @@ subdir('test_oat_hooks') subdir('test_parser') subdir('test_pg_dump') subdir('test_predtest') +subdir('test_query_id') subdir('test_radixtree') subdir('test_rbtree') subdir('test_regex') diff --git a/src/test/modules/test_query_id/Makefile b/src/test/modules/test_query_id/Makefile new file mode 100644 index 00000000000..ffd5dfb0094 --- /dev/null +++ b/src/test/modules/test_query_id/Makefile @@ -0,0 +1,20 @@ +# src/test/modules/test_query_id/Makefile + +MODULE_big = test_query_id +OBJS = \ + $(WIN32RES) \ + test_query_id.o +PGFILEDESC = "test_query_id - test module for query_id" + +REGRESS = test_query_id + +ifdef USE_PGXS +PG_CONFIG = pg_config +PGXS := $(shell $(PG_CONFIG) --pgxs) +include $(PGXS) +else +subdir = src/test/modules/test_query_id +top_builddir = ../../../.. +include $(top_builddir)/src/Makefile.global +include $(top_srcdir)/contrib/contrib-global.mk +endif diff --git a/src/test/modules/test_query_id/expected/test_query_id.out b/src/test/modules/test_query_id/expected/test_query_id.out new file mode 100644 index 00000000000..949ed50e04e --- /dev/null +++ b/src/test/modules/test_query_id/expected/test_query_id.out @@ -0,0 +1,51 @@ +-- Load test_query_id's hooks +LOAD 'test_query_id'; +-- Test query with simple protocol +SELECT 1; + ?column? +---------- + 1 +(1 row) + +-- Test a query with extended protocol +SELECT 1 \bind \g + ?column? +---------- + 1 +(1 row) + +-- Test a ctas (fail) +-- CREATE TABLE m AS SELECT i AS k, (i || ' v')::text v FROM generate_series(1, 16, 3) i; +-- Test CREATE table +CREATE TABLE n(a int); +INSERT INTO n SELECT * FROM generate_series(1, 100); +-- Test alter table +ALTER TABLE n ADD UNIQUE (a); +-- Test cursor use (fail) +-- BEGIN; +-- declare foocur CURSOR FOR SELECT * from n; +-- fetch forward 10 from foocur; +-- fetch forward 1 from foocur; +-- CLOSE foocur; +-- COMMIT; +-- Test procedure declaration +CREATE PROCEDURE proc_with_utility_stmt() +LANGUAGE SQL +AS $$ + SHOW cpu_operator_cost; + SELECT 1; +$$; +-- Test procedure call +CALL proc_with_utility_stmt(); +CREATE FUNCTION simple_function() returns int +LANGUAGE SQL +AS $$ + SELECT 1; +$$; +-- Test inline function +select * from (select * from simple_function()); + simple_function +----------------- + 1 +(1 row) + diff --git a/src/test/modules/test_query_id/meson.build b/src/test/modules/test_query_id/meson.build new file mode 100644 index 00000000000..f4401e2591c --- /dev/null +++ b/src/test/modules/test_query_id/meson.build @@ -0,0 +1,28 @@ +# Copyright (c) 2024, PostgreSQL Global Development Group + +test_query_id_sources = files( + 'test_query_id.c', +) + +if host_system == 'windows' + test_query_id_sources += rc_lib_gen.process(win32ver_rc, extra_args: [ + '--NAME', 'test_query_id', + '--FILEDESC', 'test_query_id - test module for query_ids',]) +endif + +test_query_id = shared_module('test_query_id', + test_query_id_sources, + kwargs: pg_test_mod_args, +) +test_install_libs += test_query_id + +tests += { + 'name': 'test_query_id', + 'sd': meson.current_source_dir(), + 'bd': meson.current_build_dir(), + 'regress': { + 'sql': [ + 'test_query_id', + ], + }, +} diff --git a/src/test/modules/test_query_id/sql/test_query_id.sql b/src/test/modules/test_query_id/sql/test_query_id.sql new file mode 100644 index 00000000000..7176b5412a8 --- /dev/null +++ b/src/test/modules/test_query_id/sql/test_query_id.sql @@ -0,0 +1,46 @@ +-- Load test_query_id's hooks +LOAD 'test_query_id'; + +-- Test query with simple protocol +SELECT 1; + +-- Test a query with extended protocol +SELECT 1 \bind \g + +-- Test a ctas (fail) +-- CREATE TABLE m AS SELECT i AS k, (i || ' v')::text v FROM generate_series(1, 16, 3) i; + +-- Test CREATE table +CREATE TABLE n(a int); +INSERT INTO n SELECT * FROM generate_series(1, 100); + +-- Test alter table +ALTER TABLE n ADD UNIQUE (a); + +-- Test cursor use (fail) +-- BEGIN; +-- declare foocur CURSOR FOR SELECT * from n; +-- fetch forward 10 from foocur; +-- fetch forward 1 from foocur; +-- CLOSE foocur; +-- COMMIT; + +-- Test procedure declaration +CREATE PROCEDURE proc_with_utility_stmt() +LANGUAGE SQL +AS $$ + SHOW cpu_operator_cost; + SELECT 1; +$$; + +-- Test procedure call +CALL proc_with_utility_stmt(); + +CREATE FUNCTION simple_function() returns int +LANGUAGE SQL +AS $$ + SELECT 1; +$$; + +-- Test inline function +select * from (select * from simple_function()); diff --git a/src/test/modules/test_query_id/test_query_id.c b/src/test/modules/test_query_id/test_query_id.c new file mode 100644 index 00000000000..074dbecbb83 --- /dev/null +++ b/src/test/modules/test_query_id/test_query_id.c @@ -0,0 +1,185 @@ +/*-------------------------------------------------------------------------- + * + * test_query_id.c + * Test if query_id is set and matches reported value in pgstat. + * + * Copyright (c) 2024, PostgreSQL Global Development Group + * + * IDENTIFICATION + * src/test/modules/test_query_id/test_query_id.c + * + * ------------------------------------------------------------------------- + */ + +#include "postgres.h" + +#include "executor/executor.h" +#include "optimizer/planner.h" +#include "parser/analyze.h" +#include "storage/lwlock.h" +#include "tcop/utility.h" +#include "utils/backend_status.h" + +PG_MODULE_MAGIC; + +static void test_query_id_post_parse_analyze(ParseState *pstate, Query *query, + JumbleState *jstate); +static PlannedStmt *test_query_id_planner(Query *parse, + const char *query_string, + int cursorOptions, + ParamListInfo boundParams); +static void test_query_id_ExecutorStart(QueryDesc *queryDesc, int eflags); +static void test_query_id_ExecutorRun(QueryDesc *queryDesc, + ScanDirection direction, + uint64 count, bool execute_once); +static void test_query_id_ExecutorFinish(QueryDesc *queryDesc); +static void test_query_id_ExecutorEnd(QueryDesc *queryDesc); +static void test_query_id_ProcessUtility(PlannedStmt *pstmt, const char *queryString, + bool readOnlyTree, + ProcessUtilityContext context, ParamListInfo params, + QueryEnvironment *queryEnv, + DestReceiver *dest, QueryCompletion *qc); + +/* Current nesting depth of planner/ExecutorRun/ProcessUtility calls */ +static int nesting_level = 0; + +static post_parse_analyze_hook_type prev_post_parse_analyze_hook = NULL; +static planner_hook_type prev_planner_hook = NULL; +static ExecutorStart_hook_type prev_ExecutorStart = NULL; +static ExecutorRun_hook_type prev_ExecutorRun = NULL; +static ExecutorFinish_hook_type prev_ExecutorFinish = NULL; +static ExecutorEnd_hook_type prev_ExecutorEnd = NULL; +static ProcessUtility_hook_type prev_ProcessUtility = NULL; + +static void +test_query_id_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate) +{ + if (prev_post_parse_analyze_hook) + prev_post_parse_analyze_hook(pstate, query, jstate); + Assert(query->queryId); + /* queryId will be reported to pgstat just after this hook */ +} + +static PlannedStmt * +test_query_id_planner(Query *parse, + const char *query_string, + int cursorOptions, + ParamListInfo boundParams) +{ + PlannedStmt *result; + + nesting_level++; + if (prev_planner_hook) + result = prev_planner_hook(parse, query_string, cursorOptions, + boundParams); + else + result = standard_planner(parse, query_string, cursorOptions, + boundParams); + nesting_level--; + Assert(parse->queryId); + if (nesting_level == 0) + Assert(parse->queryId == pgstat_get_my_query_id()); + return result; +} + +static void +test_query_id_ExecutorStart(QueryDesc *queryDesc, int eflags) +{ + if (prev_ExecutorStart) + prev_ExecutorStart(queryDesc, eflags); + else + standard_ExecutorStart(queryDesc, eflags); + Assert(queryDesc->plannedstmt->queryId); + if (nesting_level == 0) + Assert(queryDesc->plannedstmt->queryId == pgstat_get_my_query_id()); +} + +static void +test_query_id_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 count, + bool execute_once) +{ + nesting_level++; + if (prev_ExecutorRun) + prev_ExecutorRun(queryDesc, direction, count, execute_once); + else + standard_ExecutorRun(queryDesc, direction, count, execute_once); + nesting_level--; + Assert(queryDesc->plannedstmt->queryId); + if (nesting_level == 0) + Assert(queryDesc->plannedstmt->queryId == pgstat_get_my_query_id()); +} + +static void +test_query_id_ExecutorFinish(QueryDesc *queryDesc) +{ + nesting_level++; + if (prev_ExecutorFinish) + prev_ExecutorFinish(queryDesc); + else + standard_ExecutorFinish(queryDesc); + nesting_level--; + Assert(queryDesc->plannedstmt->queryId); + if (nesting_level == 0) + Assert(queryDesc->plannedstmt->queryId == pgstat_get_my_query_id()); +} + +static void +test_query_id_ExecutorEnd(QueryDesc *queryDesc) +{ + if (prev_ExecutorEnd) + prev_ExecutorEnd(queryDesc); + else + standard_ExecutorEnd(queryDesc); + Assert(queryDesc->plannedstmt->queryId); + if (nesting_level == 0) + Assert(queryDesc->plannedstmt->queryId == pgstat_get_my_query_id()); +} + +/* + * ProcessUtility hook + */ +static void +test_query_id_ProcessUtility(PlannedStmt *pstmt, const char *queryString, + bool readOnlyTree, + ProcessUtilityContext context, + ParamListInfo params, QueryEnvironment *queryEnv, + DestReceiver *dest, QueryCompletion *qc) +{ + nesting_level++; + if (prev_ProcessUtility) + prev_ProcessUtility(pstmt, queryString, readOnlyTree, + context, params, queryEnv, + dest, qc); + else + standard_ProcessUtility(pstmt, queryString, readOnlyTree, + context, params, queryEnv, + dest, qc); + nesting_level--; + Assert(pstmt->queryId); + if (nesting_level == 0) + Assert(pstmt->queryId == pgstat_get_my_query_id()); +} + +void +_PG_init(void) +{ + EnableQueryId(); + + /* + * Install hooks. + */ + prev_post_parse_analyze_hook = post_parse_analyze_hook; + post_parse_analyze_hook = test_query_id_post_parse_analyze; + prev_planner_hook = planner_hook; + planner_hook = test_query_id_planner; + prev_ExecutorStart = ExecutorStart_hook; + ExecutorStart_hook = test_query_id_ExecutorStart; + prev_ExecutorRun = ExecutorRun_hook; + ExecutorRun_hook = test_query_id_ExecutorRun; + prev_ExecutorFinish = ExecutorFinish_hook; + ExecutorFinish_hook = test_query_id_ExecutorFinish; + prev_ExecutorEnd = ExecutorEnd_hook; + ExecutorEnd_hook = test_query_id_ExecutorEnd; + prev_ProcessUtility = ProcessUtility_hook; + ProcessUtility_hook = test_query_id_ProcessUtility; +} -- 2.39.3 (Apple Git-146)