[#25286] YSQL, QueryDiagnostics: Convert test flag to preview flag

Summary: This diff removes the `TEST_yb_enable_query_diagnostics` and add a non runtime preview flag `ysql_yb_enable_query_diagnostics`. Jira: DB-14486 Test Plan: ./yb_build.sh --java-test 'org.yb.pgsql.TestYbQueryDiagnostics' Reviewers: asaha Reviewed By: asaha Subscribers: hbhanawat, yql Differential Revision: https://phorge.dev.yugabyte.com/D40637
yugabyte · Dec 17, 2024 · 85e1877 · 85e1877
1 parent 9db0ed1
commit 85e1877
Show file tree

Hide file tree

Showing 12 changed files with 41 additions and 44 deletions.
diff --git a/java/yb-pgsql/src/test/java/org/yb/pgsql/TestYbQueryDiagnostics.java b/java/yb-pgsql/src/test/java/org/yb/pgsql/TestYbQueryDiagnostics.java
@@ -110,7 +110,8 @@ public QueryDiagnosticsStatus(Path path, String status, String description,
     public void setUp() throws Exception {
         /* Set Gflags and restart cluster */
         Map<String, String> flagMap = super.getTServerFlags();
-        flagMap.put("TEST_yb_enable_query_diagnostics", "true");
+        flagMap.put("allowed_preview_flags_csv", "ysql_yb_enable_query_diagnostics");
+        flagMap.put("ysql_yb_enable_query_diagnostics", "true");
 
         /* Required for some of the fields within schema details */
         flagMap.put("ysql_beta_features", "true");
@@ -123,7 +124,8 @@ public void setUp() throws Exception {
     public void setUp(int queryDiagnosticsCircularBufferSize) throws Exception {
         /* Set Gflags and restart cluster */
         Map<String, String> flagMap = super.getTServerFlags();
-        flagMap.put("TEST_yb_enable_query_diagnostics", "true");
+        flagMap.put("allowed_preview_flags_csv", "ysql_yb_enable_query_diagnostics");
+        flagMap.put("ysql_yb_enable_query_diagnostics", "true");
         appendToYsqlPgConf(flagMap,
                             "yb_query_diagnostics_circular_buffer_size=" +
                             queryDiagnosticsCircularBufferSize);

diff --git a/src/postgres/contrib/pg_stat_statements/pg_stat_statements.c b/src/postgres/contrib/pg_stat_statements/pg_stat_statements.c
@@ -1842,7 +1842,7 @@ pgss_store(const char *query, uint64 queryId,
 			gc_qtexts();
 	}
 
-	if (YBIsQueryDiagnosticsEnabled())
+	if (yb_enable_query_diagnostics)
 		YbSetPgssNormalizedQueryText(queryId, entry->query_offset,
 									 entry->query_len);
 
@@ -1956,7 +1956,7 @@ pgss_store(const char *query, uint64 queryId,
 Datum
 pg_stat_statements_reset_1_7(PG_FUNCTION_ARGS)
 {
-	if (YBIsQueryDiagnosticsEnabled())
+	if (yb_enable_query_diagnostics)
 		*yb_pgss_last_reset_time = GetCurrentTimestamp();
 
 	Oid			userid;
@@ -1978,7 +1978,7 @@ pg_stat_statements_reset_1_7(PG_FUNCTION_ARGS)
 Datum
 pg_stat_statements_reset(PG_FUNCTION_ARGS)
 {
-	if (YBIsQueryDiagnosticsEnabled())
+	if (yb_enable_query_diagnostics)
 		*yb_pgss_last_reset_time = GetCurrentTimestamp();
 
 	entry_reset(0, 0, 0);

diff --git a/src/postgres/src/backend/postmaster/postmaster.c b/src/postgres/src/backend/postmaster/postmaster.c
@@ -1073,7 +1073,7 @@ PostmasterMain(int argc, char *argv[])
 		ApplyLauncherRegister();
 
 	/* Register the query diagnostics background worker */
-	if (YBIsEnabledInPostgresEnvVar() && YBIsQueryDiagnosticsEnabled())
+	if (YBIsEnabledInPostgresEnvVar() && yb_enable_query_diagnostics)
 		YbQueryDiagnosticsBgWorkerRegister();
 
 	/* Register ASH collector */

diff --git a/src/postgres/src/backend/storage/ipc/ipci.c b/src/postgres/src/backend/storage/ipc/ipci.c
@@ -153,7 +153,7 @@ CalculateShmemSize(int *num_semaphores)
 	if (YBIsEnabledInPostgresEnvVar() && yb_enable_ash)
 		size = add_size(size, YbAshShmemSize());
 
-	if (YBIsEnabledInPostgresEnvVar() && YBIsQueryDiagnosticsEnabled())
+	if (YBIsEnabledInPostgresEnvVar() && yb_enable_query_diagnostics)
 		size = add_size(size, YbQueryDiagnosticsShmemSize());
 
 	/* include additional requested shmem from preload libraries */
@@ -308,7 +308,7 @@ CreateSharedMemoryAndSemaphores(void)
 	if (YBIsEnabledInPostgresEnvVar() && yb_enable_ash)
 		YbAshShmemInit();
 
-	if (YBIsEnabledInPostgresEnvVar() && YBIsQueryDiagnosticsEnabled())
+	if (YBIsEnabledInPostgresEnvVar() && yb_enable_query_diagnostics)
 		YbQueryDiagnosticsShmemInit();
 
 #ifdef EXEC_BACKEND

diff --git a/src/postgres/src/backend/utils/misc/guc.c b/src/postgres/src/backend/utils/misc/guc.c
@@ -3037,6 +3037,18 @@ static struct config_bool ConfigureNamesBool[] =
 		NULL, NULL, NULL
 	},
 
+	{
+		{"yb_enable_query_diagnostics", PGC_POSTMASTER, STATS_MONITORING,
+			gettext_noop("Enables the collection of query diagnostics data "
+						 "for YSQL queries, facilitating the creation of diagnostic bundles."),
+			NULL,
+			GUC_NOT_IN_SAMPLE
+		},
+		&yb_enable_query_diagnostics,
+		false,
+		NULL, NULL, NULL
+	},
+
 	/* End-of-list marker */
 	{
 		{NULL, 0, 0, NULL, NULL}, NULL, false, NULL, NULL, NULL

diff --git a/src/postgres/src/backend/utils/misc/pg_yb_utils.c b/src/postgres/src/backend/utils/misc/pg_yb_utils.c
@@ -982,7 +982,7 @@ YBInitPostgresBackend(
 			if (yb_enable_ash)
 				YbAshInit();
 
-			if (YBIsEnabledInPostgresEnvVar() && YBIsQueryDiagnosticsEnabled())
+			if (yb_enable_query_diagnostics)
 				YbQueryDiagnosticsInstallHook();
 		}
 

diff --git a/src/postgres/src/backend/utils/misc/yb_query_diagnostics.c b/src/postgres/src/backend/utils/misc/yb_query_diagnostics.c
@@ -89,6 +89,7 @@ typedef struct YbQueryDiagnosticsBundles
 } YbQueryDiagnosticsBundles;
 
 /* GUC variables */
+bool yb_enable_query_diagnostics;
 int yb_query_diagnostics_bg_worker_interval_ms;
 int yb_query_diagnostics_circular_buffer_size;
 
@@ -369,10 +370,10 @@ yb_get_query_diagnostics_status(PG_FUNCTION_ARGS)
 	MemoryContext oldcontext;
 
 	/* Ensure that query diagnostics is enabled */
-	if (!YBIsQueryDiagnosticsEnabled())
+	if (!yb_enable_query_diagnostics)
 		ereport(ERROR,
 				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
-				 errmsg("TEST_yb_enable_query_diagnostics gflag must be true")));
+				 errmsg("ysql_yb_enable_query_diagnostics gflag must be true")));
 
 	/* check to see if caller supports us returning a tuplestore */
 	if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo))
@@ -2087,11 +2088,11 @@ FetchParams(YbQueryDiagnosticsParams *params, FunctionCallInfo fcinfo)
 Datum
 yb_query_diagnostics(PG_FUNCTION_ARGS)
 {
-	if (!YBIsQueryDiagnosticsEnabled())
+	if (!yb_enable_query_diagnostics)
 		ereport(ERROR,
 				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
 				 errmsg("query diagnostics is not enabled"),
-				 errhint("Set TEST_yb_enable_query_diagnostics gflag to true")));
+				 errhint("Set ysql_yb_enable_query_diagnostics gflag to true")));
 
 	YbQueryDiagnosticsMetadata metadata;
 	metadata.start_time = GetCurrentTimestamp();
@@ -2115,11 +2116,11 @@ yb_query_diagnostics(PG_FUNCTION_ARGS)
 Datum
 yb_cancel_query_diagnostics(PG_FUNCTION_ARGS)
 {
-	if (!YBIsQueryDiagnosticsEnabled())
+	if (!yb_enable_query_diagnostics)
 		ereport(ERROR,
 				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
 				 errmsg("query diagnostics is not enabled"),
-				 errhint("Set TEST_yb_enable_query_diagnostics gflag to true")));
+				 errhint("Set ysql_yb_enable_query_diagnostics gflag to true")));
 
 	int64			query_id = PG_GETARG_INT64(0);
 	YbQueryDiagnosticsEntry *entry;

diff --git a/src/postgres/src/common/pg_yb_common.c b/src/postgres/src/common/pg_yb_common.c
@@ -72,12 +72,12 @@ bool
 YBShouldAllowRunningAsAnyUser()
 {
 	if (YBIsEnabledInPostgresEnvVar())
-    {
+	{
 		return true;
 	}
 	static int cached_value = -1;
 	if (cached_value == -1)
-    {
+	{
 		cached_value = YBCIsEnvVarTrue("YB_PG_ALLOW_RUNNING_AS_ANY_USER");
 	}
 	return cached_value;
@@ -88,7 +88,7 @@ bool YBIsInitDbModeEnvVarSet()
 
 	static int cached_value = -1;
 	if (cached_value == -1)
-    {
+	{
 		cached_value = YBCIsEnvVarTrue("YB_PG_INITDB_MODE");
 	}
 	return cached_value;
@@ -260,20 +260,3 @@ Oid YBGetDatabaseOidFromEnv(const char *database_name)
 	}
 	return InvalidOid;
 }
-
-/*
- * Note: This function is used for the test flag only.
- * Once the associated feature is fully developed and stable, this function will be removed.
- * The flag is defined this way and not in ybc_pggate.cc because it is used in the ipic.c file,
- * which is initialized before the pggate api.
- */
-bool
-YBIsQueryDiagnosticsEnabled()
-{
-	static int cached_value = -1;
-	if (cached_value == -1)
-	{
-		cached_value = YBCIsEnvVarTrue("FLAGS_TEST_yb_enable_query_diagnostics");
-	}
-	return cached_value;
-}
diff --git a/src/postgres/src/include/common/pg_yb_common.h b/src/postgres/src/include/common/pg_yb_common.h
@@ -163,9 +163,4 @@ extern bool YBIsMajorUpgradeInitDb();
  */
 Oid YBGetDatabaseOidFromEnv(const char *database_name);
 
-/**
- * Returns whether the query diagnostics feature is enabled.
- */
-extern bool YBIsQueryDiagnosticsEnabled();
-
 #endif /* PG_YB_COMMON_H */
diff --git a/src/postgres/src/include/yb_query_diagnostics.h b/src/postgres/src/include/yb_query_diagnostics.h
@@ -54,6 +54,7 @@
 #define YB_QD_MAX_EXPLAIN_PLAN_LEN 16384
 
 /* GUC variables */
+extern bool yb_enable_query_diagnostics;
 extern int yb_query_diagnostics_bg_worker_interval_ms;
 extern int yb_query_diagnostics_circular_buffer_size;
 

diff --git a/src/yb/yql/pgwrapper/pg_ash-test.cc b/src/yb/yql/pgwrapper/pg_ash-test.cc
@@ -125,7 +125,9 @@ class PgBgWorkersTest : public PgAshSingleNode {
     options->extra_tserver_flags.push_back("--enable_pg_cron=true");
     options->extra_tserver_flags.push_back(
         "--ysql_pg_conf_csv=cron.yb_job_list_refresh_interval=1");
-    options->extra_tserver_flags.push_back("--TEST_yb_enable_query_diagnostics=true");
+    options->extra_tserver_flags.push_back(
+    "--allowed_preview_flags_csv=ysql_yb_enable_query_diagnostics");
+    options->extra_tserver_flags.push_back("--ysql_yb_enable_query_diagnostics=true");
     options->extra_tserver_flags.push_back(Format("--TEST_yb_ash_wait_code_to_sleep_at=$0",
         to_underlying(ash::WaitStateCode::kCatalogRead)));
     options->extra_tserver_flags.push_back(Format("--TEST_yb_ash_sleep_at_wait_state_ms=$0",

diff --git a/src/yb/yql/pgwrapper/pg_wrapper.cc b/src/yb/yql/pgwrapper/pg_wrapper.cc
@@ -298,9 +298,6 @@ DEFINE_RUNTIME_PG_FLAG(int32, yb_toast_catcache_threshold, -1,
 DEFINE_RUNTIME_PG_FLAG(string, yb_read_after_commit_visibility, "strict",
   "Determines the behavior of read-after-commit-visibility guarantee.");
 
-DEFINE_test_flag(bool, yb_enable_query_diagnostics, false,
-              "True to enable Query Diagnostics");
-
 DEFINE_RUNTIME_PG_FLAG(bool, yb_enable_fkey_catcache, true,
     "Enable preloading of foreign key information into the relation cache.");
 
@@ -318,6 +315,10 @@ DEFINE_NON_RUNTIME_string(ysql_cron_database_name, "yugabyte",
 DEFINE_NON_RUNTIME_bool(ysql_trust_local_yugabyte_connections, true,
             "Trust YSQL connections via the local socket from the yugabyte user.");
 
+DEFINE_NON_RUNTIME_PG_PREVIEW_FLAG(bool, yb_enable_query_diagnostics, false,
+    "Enables the collection of query diagnostics data for YSQL queries, "
+    "facilitating the creation of diagnostic bundles.");
+
 DECLARE_bool(enable_pg_cron);
 
 using gflags::CommandLineFlagInfo;