Fix to skip SQL normalization in pgss_store()

geldata · Oct 1, 2024 · f6e7de1 · f6e7de1
1 parent cbd6c4a
commit f6e7de1
Showing 1 changed file with 24 additions and 6 deletions.
diff --git a/edb_stat_statements/edb_stat_statements.c b/edb_stat_statements/edb_stat_statements.c
@@ -72,15 +72,33 @@ _PG_init(void) {
  */
 static void
 edbss_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate) {
-    const char *orig_sourcetext = pstate->p_sourcetext;
+    if (jstate) {
+        // Keep old data
+        const char *orig_sourcetext = pstate->p_sourcetext;
+        int orig_stmt_location = query->stmt_location;
+        int orig_clocations_count = jstate->clocations_count;
+
+        // Parse EdgeDB query info JSON and overwrite
+        if (edbss_overwrite_stmt_info(pstate, query)) {
+            query->stmt_location = -1;  // CleanQuerytext will fix this and stmt_len
+            jstate->clocations_count = 0;  // Skip the SQL generate_normalized_query()
+        } else {
+            // Skip all non-EdgeDB queries
+            jstate = NULL;
+        }
 
-    if (!edbss_overwrite_stmt_info(pstate, query))
-        jstate = NULL;
+        // This would call pgss_store()
+        if (prev_post_parse_analyze_hook)
+            prev_post_parse_analyze_hook(pstate, query, jstate);
 
-    if (prev_post_parse_analyze_hook)
-        prev_post_parse_analyze_hook(pstate, query, jstate);
+        // Restore values after pgss_store()
+        pstate->p_sourcetext = orig_sourcetext;
+        query->stmt_location = orig_stmt_location;
+        if (jstate)
+            jstate->clocations_count = orig_clocations_count;
 
-    pstate->p_sourcetext = orig_sourcetext;
+    } else if (prev_post_parse_analyze_hook)
+        prev_post_parse_analyze_hook(pstate, query, jstate);
 }
 
 /*