Skip to content

Commit 1a019f4

Browse files
pdecatpdecat
committed
feat: support propagation of OpenTelemetry context from clients with SQLcommenter
1 parent 6b8a367 commit 1a019f4

File tree

5 files changed

+286
-19
lines changed

5 files changed

+286
-19
lines changed

‎fdw.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -323,7 +323,7 @@ func goFdwBeginForeignScan(node *C.ForeignScanState, eflags C.int) {
323323

324324
// Extract trace context from session variables for scan operation
325325
var traceContext string
326-
if traceContextPtr := C.getTraceContextFromSession(); traceContextPtr != nil {
326+
if traceContextPtr := C.getTraceContext(); traceContextPtr != nil {
327327
traceContext = C.GoString(traceContextPtr)
328328
log.Printf("[TRACE] Extracted trace context from session for scan: %s", traceContext)
329329
} else {

‎fdw/fdw.c

Lines changed: 139 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#include "access/xact.h"
77
#include "utils/guc.h"
88
#include "utils/builtins.h"
9+
#include "tcop/tcopprot.h"
910

1011
extern PGDLLEXPORT void _PG_init(void);
1112

@@ -16,6 +17,11 @@ static char *convertUUID(char *uuid);
1617
static void pgfdw_xact_callback(XactEvent event, void *arg);
1718
static void exitHook(int code, Datum arg);
1819

20+
/* Trace context extraction functions */
21+
static char *extractTraceContextFromSession(void);
22+
static char *extractTraceContextFromQueryComments(void);
23+
static char *extractTraceContext(void);
24+
1925
void *serializePlanState(FdwPlanState *state);
2026
FdwExecState *initializeExecState(void *internalstate);
2127
// Required by postgres, doing basic checks to ensure compatibility,
@@ -103,29 +109,153 @@ static char *extractTraceContextFromSession(void)
103109
const char *traceparent = GetConfigOption("steampipe.traceparent", true, false);
104110
const char *tracestate = GetConfigOption("steampipe.tracestate", true, false);
105111
char *result = NULL;
106-
112+
107113
// Format the result string for Go layer consumption
108114
if (traceparent != NULL) {
109115
if (tracestate != NULL) {
110116
result = psprintf("traceparent=%s;tracestate=%s", traceparent, tracestate);
111117
} else {
112118
result = psprintf("traceparent=%s", traceparent);
113119
}
114-
120+
115121
elog(DEBUG1, "extracted trace context: %s", result);
116122
} else {
117123
elog(DEBUG2, "no trace context found in session variables");
118124
}
119-
125+
120126
return result;
121127
}
122128

123129
/*
124-
* Public wrapper for extractTraceContextFromSession - callable from Go
130+
* Extract OpenTelemetry trace context from SQL query comments (SQLcommenter format)
131+
* Parses comments like: /*traceparent='00-...',tracestate='rojo=...'*\/
132+
* Returns a formatted string containing traceparent and tracestate, or NULL if not found
125133
*/
126-
char *getTraceContextFromSession(void)
134+
static char *extractTraceContextFromQueryComments(void)
127135
{
128-
return extractTraceContextFromSession();
136+
const char *query_string = debug_query_string;
137+
char *result = NULL;
138+
char *traceparent = NULL;
139+
char *tracestate = NULL;
140+
141+
if (query_string == NULL) {
142+
elog(DEBUG2, "no query string available for SQLcommenter parsing");
143+
return NULL;
144+
}
145+
146+
elog(DEBUG2, "parsing SQLcommenter from query: %.100s...", query_string);
147+
148+
// Look for SQL comments in the format /*...*/
149+
const char *comment_start = strstr(query_string, "/*");
150+
while (comment_start != NULL) {
151+
const char *comment_end = strstr(comment_start, "*/");
152+
if (comment_end == NULL) {
153+
break; // Malformed comment, skip
154+
}
155+
156+
// Extract the comment content
157+
size_t comment_len = comment_end - comment_start - 2; // Exclude /* and */
158+
char *comment_content = palloc(comment_len + 1);
159+
strncpy(comment_content, comment_start + 2, comment_len);
160+
comment_content[comment_len] = '\0';
161+
162+
elog(DEBUG2, "found SQL comment: %s", comment_content);
163+
164+
// Parse key-value pairs in the comment
165+
char *token = strtok(comment_content, ",");
166+
while (token != NULL) {
167+
// Trim whitespace
168+
while (*token == ' ' || *token == '\t') token++;
169+
170+
// Look for traceparent or tracestate
171+
if (strncmp(token, "traceparent=", 12) == 0) {
172+
char *value = token + 12;
173+
// Remove quotes if present
174+
if (*value == '\'' || *value == '"') {
175+
char quote_char = *value;
176+
value++;
177+
char *end_quote = strrchr(value, quote_char);
178+
if (end_quote) *end_quote = '\0';
179+
}
180+
if (traceparent) pfree(traceparent);
181+
traceparent = pstrdup(value);
182+
elog(DEBUG2, "extracted traceparent from SQLcommenter: %s", traceparent);
183+
} else if (strncmp(token, "tracestate=", 11) == 0) {
184+
char *value = token + 11;
185+
// Remove quotes if present
186+
if (*value == '\'' || *value == '"') {
187+
value++;
188+
char *end_quote = strrchr(value, *value == '\'' ? '\'' : '"');
189+
if (end_quote) *end_quote = '\0';
190+
}
191+
if (tracestate) pfree(tracestate);
192+
tracestate = pstrdup(value);
193+
elog(DEBUG2, "extracted tracestate from SQLcommenter: %s", tracestate);
194+
}
195+
196+
token = strtok(NULL, ",");
197+
}
198+
199+
pfree(comment_content);
200+
201+
// Look for next comment
202+
comment_start = strstr(comment_end + 2, "/*");
203+
}
204+
205+
// Format the result string for Go layer consumption
206+
if (traceparent != NULL) {
207+
if (tracestate != NULL) {
208+
result = psprintf("traceparent=%s;tracestate=%s", traceparent, tracestate);
209+
} else {
210+
result = psprintf("traceparent=%s", traceparent);
211+
}
212+
213+
elog(DEBUG1, "extracted trace context from SQLcommenter: %s", result);
214+
} else {
215+
elog(DEBUG2, "no trace context found in SQL comments");
216+
}
217+
218+
// Clean up
219+
if (traceparent) pfree(traceparent);
220+
if (tracestate) pfree(tracestate);
221+
222+
return result;
223+
}
224+
225+
/*
226+
* Extract trace context with fallback strategy:
227+
* 1. Try PostgreSQL session variables first
228+
* 2. Fall back to SQLcommenter in query comments
229+
* 3. Return NULL if neither found
230+
*/
231+
static char *extractTraceContext(void)
232+
{
233+
char *result = NULL;
234+
235+
// First try session variables (primary method)
236+
result = extractTraceContextFromSession();
237+
if (result != NULL) {
238+
elog(DEBUG1, "using trace context from session variables");
239+
return result;
240+
}
241+
242+
// Fall back to SQLcommenter (secondary method)
243+
result = extractTraceContextFromQueryComments();
244+
if (result != NULL) {
245+
elog(DEBUG1, "using trace context from SQLcommenter");
246+
return result;
247+
}
248+
249+
elog(DEBUG2, "no trace context found in session variables or SQLcommenter");
250+
return NULL;
251+
}
252+
253+
/*
254+
* Public wrapper for extractTraceContext - callable from Go
255+
*/
256+
char *getTraceContext(void)
257+
{
258+
return extractTraceContext();
129259
}
130260

131261
static void fdwGetForeignRelSize(PlannerInfo *root, RelOptInfo *baserel, Oid foreigntableid)
@@ -147,9 +277,9 @@ static void fdwGetForeignRelSize(PlannerInfo *root, RelOptInfo *baserel, Oid for
147277
// Save plan state information
148278
baserel->fdw_private = planstate;
149279
planstate->foreigntableid = foreigntableid;
150-
151-
// Extract trace context from session variables
152-
char *traceContext = extractTraceContextFromSession();
280+
281+
// Extract trace context with fallback strategy (session variables -> SQLcommenter)
282+
char *traceContext = extractTraceContext();
153283
if (traceContext != NULL) {
154284
planstate->trace_context_string = pstrdup(traceContext);
155285
pfree(traceContext);

‎fdw/fdw_helpers.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,4 +108,4 @@ static inline char *nameStr(Name n) { return NameStr(*n); }
108108
char *tagTypeToString(NodeTag type);
109109

110110
// trace context
111-
char *getTraceContextFromSession(void);
111+
char *getTraceContext(void);

‎hub/hub_base.go

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -418,8 +418,10 @@ func (h *hubBase) traceContextForScan(table string, columns []string, limit int6
418418
return &telemetry.TraceCtx{Ctx: ctx, Span: span}
419419
}
420420

421-
// parseTraceContext parses trace context string from session variables
422-
// Format: "traceparent=00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01;tracestate=rojo=00f067aa0ba902b7"
421+
// parseTraceContext parses trace context string from session variables or SQLcommenter
422+
// Supports both formats:
423+
// - Session variables: "traceparent=00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01;tracestate=rojo=00f067aa0ba902b7"
424+
// - SQLcommenter: "traceparent='00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',tracestate='rojo=00f067aa0ba902b7'"
423425
func (h *hubBase) parseTraceContext(traceContextString string) context.Context {
424426
log.Printf("[DEBUG] parseTraceContext called with: %s", traceContextString)
425427

@@ -430,14 +432,30 @@ func (h *hubBase) parseTraceContext(traceContextString string) context.Context {
430432

431433
carrier := propagation.MapCarrier{}
432434

433-
// Parse the trace context string format: "traceparent=..;tracestate=.."
434-
parts := strings.Split(traceContextString, ";")
435-
log.Printf("[DEBUG] Split trace context into %d parts: %v", len(parts), parts)
435+
// Detect format and parse accordingly
436+
var parts []string
437+
if strings.Contains(traceContextString, ",") {
438+
// SQLcommenter format: "traceparent='...',tracestate='...'"
439+
parts = strings.Split(traceContextString, ",")
440+
log.Printf("[DEBUG] Detected SQLcommenter format, split into %d parts: %v", len(parts), parts)
441+
} else {
442+
// Session variable format: "traceparent=..;tracestate=.."
443+
parts = strings.Split(traceContextString, ";")
444+
log.Printf("[DEBUG] Detected session variable format, split into %d parts: %v", len(parts), parts)
445+
}
436446

437447
for _, part := range parts {
438448
if kv := strings.SplitN(part, "=", 2); len(kv) == 2 {
439449
key := strings.TrimSpace(kv[0])
440450
value := strings.TrimSpace(kv[1])
451+
452+
// Remove quotes from SQLcommenter format
453+
if (strings.HasPrefix(value, "'") && strings.HasSuffix(value, "'")) ||
454+
(strings.HasPrefix(value, "\"") && strings.HasSuffix(value, "\"")) {
455+
value = value[1 : len(value)-1]
456+
log.Printf("[DEBUG] Removed quotes from value: %s", value)
457+
}
458+
441459
carrier[key] = value
442460
log.Printf("[DEBUG] Added to carrier: %s = %s", key, value)
443461
} else {

0 commit comments

Comments
 (0)