fix

ydb-platform · Oct 17, 2024 · 84bc831 · 84bc831
1 parent 4ed21d0
commit 84bc831
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 15 deletions.
diff --git a/otel/config.go b/otel/config.go
@@ -14,6 +14,8 @@ type (
 	Span     interface {
 		TraceID() string
 
+		Relation(span Span)
+
 		Msg(msg string, attributes ...KeyValue)
 
 		End(attributes ...KeyValue)

diff --git a/otel/noop.go b/otel/noop.go
@@ -16,6 +16,8 @@ type (
 	noopSpan   struct{}
 )
 
+func (noopSpan) Relation(Span) {}
+
 func (noopSpan) TraceID() string {
 	return ""
 }

diff --git a/otel/sql.go b/otel/sql.go
@@ -3,8 +3,6 @@ package otel
 import (
 	"context"
 
-	otelTrace "go.opentelemetry.io/otel/trace"
-
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/kv"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
@@ -206,16 +204,17 @@ func databaseSQL(cfg Config) (t trace.DatabaseSQL) {
 	}
 	t.OnTxExec = func(info trace.DatabaseSQLTxExecStartInfo) func(trace.DatabaseSQLTxExecDoneInfo) {
 		if cfg.Details()&trace.DatabaseSQLTxEvents != 0 {
-			if !isStmtCall(*info.Context) {
-				*info.Context = otelTrace.ContextWithSpan(*info.Context, otelTrace.SpanFromContext(info.TxContext))
-			}
 			start := childSpanWithReplaceCtx(
 				info.Context,
 				info.Call.FunctionID(),
 				kv.String("query", info.Query),
 				kv.String("transaction_id", safeID(info.Tx)),
 			)
 
+			if !isStmtCall(*info.Context) {
+				start.Relation(cfg.SpanFromContext(info.TxContext))
+			}
+
 			return func(info trace.DatabaseSQLTxExecDoneInfo) {
 				finish(
 					start,
@@ -228,16 +227,17 @@ func databaseSQL(cfg Config) (t trace.DatabaseSQL) {
 	}
 	t.OnTxQuery = func(info trace.DatabaseSQLTxQueryStartInfo) func(trace.DatabaseSQLTxQueryDoneInfo) {
 		if cfg.Details()&trace.DatabaseSQLTxEvents != 0 {
-			if !isStmtCall(*info.Context) {
-				*info.Context = otelTrace.ContextWithSpan(*info.Context, otelTrace.SpanFromContext(info.TxContext))
-			}
 			start := childSpanWithReplaceCtx(
 				info.Context,
 				info.Call.FunctionID(),
 				kv.String("query", info.Query),
 				kv.String("transaction_id", safeID(info.Tx)),
 			)
 
+			if !isStmtCall(*info.Context) {
+				start.Relation(cfg.SpanFromContext(info.TxContext))
+			}
+
 			return func(info trace.DatabaseSQLTxQueryDoneInfo) {
 				finish(
 					start,
@@ -250,7 +250,6 @@ func databaseSQL(cfg Config) (t trace.DatabaseSQL) {
 	}
 	t.OnTxPrepare = func(info trace.DatabaseSQLTxPrepareStartInfo) func(trace.DatabaseSQLTxPrepareDoneInfo) {
 		if cfg.Details()&trace.DatabaseSQLTxEvents != 0 {
-			*info.Context = otelTrace.ContextWithSpan(*info.Context, otelTrace.SpanFromContext(info.TxContext))
 			start := childSpanWithReplaceCtx(
 				info.Context,
 				info.Call.FunctionID(),
@@ -270,15 +269,16 @@ func databaseSQL(cfg Config) (t trace.DatabaseSQL) {
 	}
 	t.OnStmtExec = func(info trace.DatabaseSQLStmtExecStartInfo) func(trace.DatabaseSQLStmtExecDoneInfo) {
 		if cfg.Details()&trace.DatabaseSQLStmtEvents != 0 {
-			*info.Context = markStmtCall(
-				otelTrace.ContextWithSpan(*info.Context, otelTrace.SpanFromContext(info.StmtContext)),
-			)
 			start := childSpanWithReplaceCtx(
 				info.Context,
 				info.Call.FunctionID(),
 				kv.String("query", info.Query),
 			)
 
+			start.Relation(cfg.SpanFromContext(info.StmtContext))
+
+			*info.Context = markStmtCall(*info.Context)
+
 			return func(info trace.DatabaseSQLStmtExecDoneInfo) {
 				finish(
 					start,
@@ -291,15 +291,16 @@ func databaseSQL(cfg Config) (t trace.DatabaseSQL) {
 	}
 	t.OnStmtQuery = func(info trace.DatabaseSQLStmtQueryStartInfo) func(trace.DatabaseSQLStmtQueryDoneInfo) {
 		if cfg.Details()&trace.DatabaseSQLStmtEvents != 0 {
-			*info.Context = markStmtCall(
-				otelTrace.ContextWithSpan(*info.Context, otelTrace.SpanFromContext(info.StmtContext)),
-			)
 			start := childSpanWithReplaceCtx(
 				info.Context,
 				info.Call.FunctionID(),
 				kv.String("query", info.Query),
 			)
 
+			start.Relation(cfg.SpanFromContext(info.StmtContext))
+
+			*info.Context = markStmtCall(*info.Context)
+
 			return func(info trace.DatabaseSQLStmtQueryDoneInfo) {
 				finish(
 					start,