[#22147] YSQL, QueryDiagnostics: Pgss support for query diagnostics

Summary:
- This diff builds over 47f40bb and adds the pg_stat_statements.csv file within the bundle output.
- All of the pgss data except the normalised query string is extracted from the QueryDesc.
- For the normalised query, query_offset and query_len is stored within the query_diagnostics shared hash_table, which are later utilized to extract normalised query from pgss_query_texts.stat
- PGSS data accumulates during the bundle lifetime, flushing to the filesystem only at bundle completion.
- note that the `query latency histogram` is part of [[ #22324 | #22324 ]]
Jira: DB-11075

Test Plan: ./yb_build.sh --java-test TestYsqlUpgrade#checkPgssData

Reviewers: asaha, hbhanawat, amartsinchyk, telgersma

Reviewed By: telgersma

Subscribers: ybase, yql

Differential Revision: https://phorge.dev.yugabyte.com/D35234

java/yb-pgsql/src/test/java/org/yb/pgsql/TestYbQueryDiagnostics.java

@@ -52,12 +52,13 @@ public void setUp() throws Exception {
         }
     }
 
-    private String getQueryIdOfPreparedStmt(Statement statement) throws Exception {
+    private String getQueryIdFromPgStatStatements(Statement statement, String pattern)
+                                                  throws Exception {
         /* Get query id of the prepared statement */
         ResultSet resultSet = statement.executeQuery("SELECT queryid FROM pg_stat_statements " +
-                                                     "WHERE query LIKE 'PREPARE%'");
+                                                     "WHERE query LIKE '" + pattern + "'");
         if (!resultSet.next())
-            fail("Query id not found for the prepared statement");
+            fail("Query id not found in pg_stat_statements");
 
         return resultSet.getString("queryid");
     }
@@ -129,7 +130,7 @@ public void checkBindVariablesData() throws Exception {
 
         try (Statement statement = connection.createStatement()) {
             /* Run query diagnostics on the prepared stmt */
-            String queryId = getQueryIdOfPreparedStmt(statement);
+            String queryId = getQueryIdFromPgStatStatements(statement, "PREPARE%");
             Path bundleDataPath = runQueryDiagnostics(statement, queryId, params);
 
             /* Generate some data to be dumped */
@@ -180,7 +181,7 @@ public void testYbQueryDiagnosticsStatus() throws Exception {
 
         try (Statement statement = connection.createStatement()) {
             /* Run query diagnostics on the prepared stmt */
-            String queryId = getQueryIdOfPreparedStmt(statement);
+            String queryId = getQueryIdFromPgStatStatements(statement, "PREPARE%");
             QueryDiagnosticsParams successfulRunParams = new QueryDiagnosticsParams(
                 diagnosticsInterval,
                 100 /* explainSampleRate */,
@@ -365,4 +366,61 @@ private void runMultipleBundles(Statement statement, int n) throws Exception {
         }
         fail("Buffer never wrapped around");
     }
+
+    @Test
+    public void checkPgssData() throws Exception {
+        int diagnosticsInterval = 10;
+        QueryDiagnosticsParams queryDiagnosticsParams = new QueryDiagnosticsParams(
+            diagnosticsInterval,
+            100 /* explainSampleRate */,
+            true /* explainAnalyze */,
+            true /* explainDist */,
+            false /* explainDebug */,
+            0 /* bindVarQueryMinDuration */);
+
+        try (Statement statement = connection.createStatement()) {
+            statement.execute("SELECT pg_sleep(0.5)");
+
+            String queryId = getQueryIdFromPgStatStatements(statement, "%pg_sleep%");
+            Path bundleDataPath = runQueryDiagnostics(statement, queryId, queryDiagnosticsParams);
+
+            statement.execute("SELECT pg_sleep(0.1)");
+            statement.execute("select * from pg_class");
+            statement.execute("select pg_sleep(0.2)");
+
+            /* Thread sleeps for diagnosticsInterval + 1 sec to ensure that the bundle has expired*/
+            Thread.sleep((diagnosticsInterval + 1) * 1000);
+
+            Path pgssPath = bundleDataPath.resolve("pg_stat_statements.csv");
+
+            assertTrue("pg_stat_statements file does not exist", Files.exists(pgssPath));
+            assertGreaterThan("pg_stat_statements.csv file is empty",
+                              Files.size(pgssPath) , 0L);
+
+            /* Unit is ms */
+            float epsilon = 10;
+            int expectedTotalTime = 300;
+            int expectedMinTime = 100;
+            int expectedMaxTime = 200;
+            int expectedMeanTime = 150;
+            List<String> pgssData = Files.readAllLines(pgssPath);
+            String[] tokens = pgssData.get(1).split(",");
+
+            assertEquals("pg_stat_statements data size is not as expected",
+                         2, pgssData.size());
+            assertEquals("pg_stat_statements query is incorrect", queryId, tokens[0]);
+            assertTrue("pg_stat_statements contains unnecessary data",
+                       !tokens[1].contains("pg_class"));
+            assertEquals("Number of calls are incorrect", "2", tokens[2]);
+            /* pg_stat_statements outputs data in ms */
+            assertLessThan("total_time is incorrect",
+                           Math.abs(Float.parseFloat(tokens[3]) - expectedTotalTime), epsilon);
+            assertLessThan("min_time is incorrect",
+                           Math.abs(Float.parseFloat(tokens[4]) - expectedMinTime), epsilon);
+            assertLessThan("max_time is incorrect",
+                           Math.abs(Float.parseFloat(tokens[5]) - expectedMaxTime), epsilon);
+            assertLessThan("mean_time is incorrect",
+                           Math.abs(Float.parseFloat(tokens[6]) - expectedMeanTime), epsilon);
+        }
+    }
 }

src/postgres/contrib/pg_stat_statements/pg_stat_statements.c

@@ -92,6 +92,7 @@
 
 
 #include "yb/yql/pggate/webserver/pgsql_webserver_wrapper.h"
+#include "yb_query_diagnostics.h"
 
 PG_MODULE_MAGIC;
 
@@ -426,6 +427,7 @@ static int query_buffer_helper(FILE *file, FILE *qfile, int qlen,
 	pgssReaderContext *context);
 static void enforce_bucket_factor(int * value);
 static bool yb_track_nested_queries(void);
+static void YbGetPgssNormalizedQueryText(Size query_offset, int query_len, char *normalized_query);
 
 /*
  * Module load callback
@@ -579,6 +581,12 @@ _PG_init(void)
 
 	/* Function pointer to check if nested queries should be tracked in ASH */
 	yb_ash_track_nested_queries = yb_track_nested_queries;
+
+	/*
+	 * Initializing the function pointer required by yb_query_diagnostics.c
+	 * to get normalized query text.
+	 */
+	yb_get_normalized_query = &YbGetPgssNormalizedQueryText;
 }
 
 /*
@@ -1662,6 +1670,10 @@ pgss_store(const char *query, uint64 queryId,
 			gc_qtexts();
 	}
 
+	if (YBIsQueryDiagnosticsEnabled())
+		YbSetPgssNormalizedQueryText(queryId, entry->query_offset,
+									 entry->query_len);
+
 	/* Increment the counts, except when jstate is not NULL */
 	if (!jstate)
 	{
@@ -3807,3 +3819,18 @@ yb_track_nested_queries(void)
 {
 	return pgss_track == PGSS_TRACK_ALL;
 }
+
+
+static void
+YbGetPgssNormalizedQueryText(Size query_offset, int query_len, char *normalized_query)
+{
+	char	   *qbuffer = NULL;
+	Size		qbuffer_size = 0;
+
+	qbuffer = qtext_load_file(&qbuffer_size);
+	memcpy(normalized_query, qtext_fetch(query_offset, query_len,
+										 qbuffer, qbuffer_size), query_len);
+	normalized_query[query_len - 1] = '\0'; /* Ensure null-termination */
+
+	free(qbuffer);
+}