feat: support propagation of OpenTelemetry context from clients with SQLcommenter (#569)

Author: pdecat

fdw.go

@@ -323,7 +323,7 @@ func goFdwBeginForeignScan(node *C.ForeignScanState, eflags C.int) {
 
 	// Extract trace context from session variables for scan operation
 	var traceContext string
-	if traceContextPtr := C.getTraceContextFromSession(); traceContextPtr != nil {
+	if traceContextPtr := C.getTraceContext(); traceContextPtr != nil {
 		traceContext = C.GoString(traceContextPtr)
 		log.Printf("[TRACE] Extracted trace context from session for scan: %s", traceContext)
 	} else {

fdw/fdw.c

@@ -6,6 +6,7 @@
 #include "access/xact.h"
 #include "utils/guc.h"
 #include "utils/builtins.h"
+#include "tcop/tcopprot.h"
 
 extern PGDLLEXPORT void _PG_init(void);
 
@@ -16,6 +17,11 @@ static char *convertUUID(char *uuid);
 static void pgfdw_xact_callback(XactEvent event, void *arg);
 static void exitHook(int code, Datum arg);
 
+/* Trace context extraction functions */
+static char *extractTraceContextFromSession(void);
+static char *extractTraceContextFromQueryComments(void);
+static char *extractTraceContext(void);
+
 void *serializePlanState(FdwPlanState *state);
 FdwExecState *initializeExecState(void *internalstate);
 // Required by postgres, doing basic checks to ensure compatibility,
@@ -103,29 +109,154 @@ static char *extractTraceContextFromSession(void)
     const char *traceparent = GetConfigOption("steampipe.traceparent", true, false);
     const char *tracestate = GetConfigOption("steampipe.tracestate", true, false);
     char *result = NULL;
-    
+
     // Format the result string for Go layer consumption
     if (traceparent != NULL) {
         if (tracestate != NULL) {
             result = psprintf("traceparent=%s;tracestate=%s", traceparent, tracestate);
         } else {
             result = psprintf("traceparent=%s", traceparent);
         }
-        
+
         elog(DEBUG1, "extracted trace context: %s", result);
     } else {
         elog(DEBUG2, "no trace context found in session variables");
     }
-    
+
     return result;
 }
 
 /*
- * Public wrapper for extractTraceContextFromSession - callable from Go
+ * Extract OpenTelemetry trace context from SQL query comments (SQLcommenter format)
+ * Parses comments like: /*traceparent='00-...',tracestate='rojo=...'*\/
+ * Returns a formatted string containing traceparent and tracestate, or NULL if not found
  */
-char *getTraceContextFromSession(void)
+static char *extractTraceContextFromQueryComments(void)
 {
-    return extractTraceContextFromSession();
+    const char *query_string = debug_query_string;
+    char *result = NULL;
+    char *traceparent = NULL;
+    char *tracestate = NULL;
+
+    if (query_string == NULL) {
+        elog(DEBUG2, "no query string available for SQLcommenter parsing");
+        return NULL;
+    }
+
+    elog(DEBUG2, "parsing SQLcommenter from query: %.100s...", query_string);
+
+    // Look for SQL comments in the format /*...*/
+    const char *comment_start = strstr(query_string, "/*");
+    while (comment_start != NULL) {
+        const char *comment_end = strstr(comment_start, "*/");
+        if (comment_end == NULL) {
+            break; // Malformed comment, skip
+        }
+
+        // Extract the comment content
+        size_t comment_len = comment_end - comment_start - 2; // Exclude /* and */
+        char *comment_content = palloc(comment_len + 1);
+        strncpy(comment_content, comment_start + 2, comment_len);
+        comment_content[comment_len] = '\0';
+
+        elog(DEBUG2, "found SQL comment: %s", comment_content);
+
+        // Parse key-value pairs in the comment
+        char *token = strtok(comment_content, ",");
+        while (token != NULL) {
+            // Trim whitespace
+            while (*token == ' ' || *token == '\t') token++;
+
+            // Look for traceparent or tracestate
+            if (strncmp(token, "traceparent=", 12) == 0) {
+                char *value = token + 12;
+                // Remove quotes if present
+                if (*value == '\'' || *value == '"') {
+                    char quote_char = *value;
+                    value++;
+                    char *end_quote = strrchr(value, quote_char);
+                    if (end_quote) *end_quote = '\0';
+                }
+                if (traceparent) pfree(traceparent);
+                traceparent = pstrdup(value);
+                elog(DEBUG2, "extracted traceparent from SQLcommenter: %s", traceparent);
+            } else if (strncmp(token, "tracestate=", 11) == 0) {
+                char *value = token + 11;
+                // Remove quotes if present
+                if (*value == '\'' || *value == '"') {
+                    char quote_char = *value;
+                    value++;
+                    char *end_quote = strrchr(value, quote_char);
+                    if (end_quote) *end_quote = '\0';
+                }
+                if (tracestate) pfree(tracestate);
+                tracestate = pstrdup(value);
+                elog(DEBUG2, "extracted tracestate from SQLcommenter: %s", tracestate);
+            }
+
+            token = strtok(NULL, ",");
+        }
+
+        pfree(comment_content);
+
+        // Look for next comment
+        comment_start = strstr(comment_end + 2, "/*");
+    }
+
+    // Format the result string for Go layer consumption
+    if (traceparent != NULL) {
+        if (tracestate != NULL) {
+            result = psprintf("traceparent=%s;tracestate=%s", traceparent, tracestate);
+        } else {
+            result = psprintf("traceparent=%s", traceparent);
+        }
+
+        elog(DEBUG1, "extracted trace context from SQLcommenter: %s", result);
+    } else {
+        elog(DEBUG2, "no trace context found in SQL comments");
+    }
+
+    // Clean up
+    if (traceparent) pfree(traceparent);
+    if (tracestate) pfree(tracestate);
+
+    return result;
+}
+
+/*
+ * Extract trace context with fallback strategy:
+ * 1. Try PostgreSQL session variables first
+ * 2. Fall back to SQLcommenter in query comments
+ * 3. Return NULL if neither found
+ */
+static char *extractTraceContext(void)
+{
+    char *result = NULL;
+
+    // First try session variables (primary method)
+    result = extractTraceContextFromSession();
+    if (result != NULL) {
+        elog(DEBUG1, "using trace context from session variables");
+        return result;
+    }
+
+    // Fall back to SQLcommenter (secondary method)
+    result = extractTraceContextFromQueryComments();
+    if (result != NULL) {
+        elog(DEBUG1, "using trace context from SQLcommenter");
+        return result;
+    }
+
+    elog(DEBUG2, "no trace context found in session variables or SQLcommenter");
+    return NULL;
+}
+
+/*
+ * Public wrapper for extractTraceContext - callable from Go
+ */
+char *getTraceContext(void)
+{
+    return extractTraceContext();
 }
 
 static void fdwGetForeignRelSize(PlannerInfo *root, RelOptInfo *baserel, Oid foreigntableid)
@@ -147,9 +278,9 @@ static void fdwGetForeignRelSize(PlannerInfo *root, RelOptInfo *baserel, Oid for
   // Save plan state information
   baserel->fdw_private = planstate;
   planstate->foreigntableid = foreigntableid;
-  
-  // Extract trace context from session variables
-  char *traceContext = extractTraceContextFromSession();
+
+  // Extract trace context with fallback strategy (session variables -> SQLcommenter)
+  char *traceContext = extractTraceContext();
   if (traceContext != NULL) {
       planstate->trace_context_string = pstrdup(traceContext);
       pfree(traceContext);

fdw/fdw_helpers.h

@@ -108,4 +108,4 @@ static inline char *nameStr(Name n) { return NameStr(*n); }
 char *tagTypeToString(NodeTag type);
 
 // trace context
-char *getTraceContextFromSession(void);
+char *getTraceContext(void);

hub/hub_base.go

@@ -418,8 +418,10 @@ func (h *hubBase) traceContextForScan(table string, columns []string, limit int6
 	return &telemetry.TraceCtx{Ctx: ctx, Span: span}
 }
 
-// parseTraceContext parses trace context string from session variables
-// Format: "traceparent=00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01;tracestate=rojo=00f067aa0ba902b7"
+// parseTraceContext parses trace context string from session variables or SQLcommenter
+// Supports both formats:
+// - Session variables: "traceparent=00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01;tracestate=rojo=00f067aa0ba902b7"
+// - SQLcommenter: "traceparent='00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',tracestate='rojo=00f067aa0ba902b7'"
 func (h *hubBase) parseTraceContext(traceContextString string) context.Context {
 	log.Printf("[DEBUG] parseTraceContext called with: %s", traceContextString)
 
@@ -430,14 +432,30 @@ func (h *hubBase) parseTraceContext(traceContextString string) context.Context {
 
 	carrier := propagation.MapCarrier{}
 
-	// Parse the trace context string format: "traceparent=..;tracestate=.."
-	parts := strings.Split(traceContextString, ";")
-	log.Printf("[DEBUG] Split trace context into %d parts: %v", len(parts), parts)
+	// Detect format and parse accordingly
+	var parts []string
+	if strings.Contains(traceContextString, ",") {
+		// SQLcommenter format: "traceparent='...',tracestate='...'"
+		parts = strings.Split(traceContextString, ",")
+		log.Printf("[DEBUG] Detected SQLcommenter format, split into %d parts: %v", len(parts), parts)
+	} else {
+		// Session variable format: "traceparent=..;tracestate=.."
+		parts = strings.Split(traceContextString, ";")
+		log.Printf("[DEBUG] Detected session variable format, split into %d parts: %v", len(parts), parts)
+	}
 
 	for _, part := range parts {
 		if kv := strings.SplitN(part, "=", 2); len(kv) == 2 {
 			key := strings.TrimSpace(kv[0])
 			value := strings.TrimSpace(kv[1])
+
+			// Remove quotes from SQLcommenter format
+			if (strings.HasPrefix(value, "'") && strings.HasSuffix(value, "'")) ||
+				(strings.HasPrefix(value, "\"") && strings.HasSuffix(value, "\"")) {
+				value = value[1 : len(value)-1]
+				log.Printf("[DEBUG] Removed quotes from value: %s", value)
+			}
+
 			carrier[key] = value
 			log.Printf("[DEBUG] Added to carrier: %s = %s", key, value)
 		} else {