Merge pull request #349 from ydb-platform/metadata

added metadata to `trace.Driver.OnInvoke` and `trace.Driver.OnNewStre…

Author: asmyasnikov

CHANGELOG.md

@@ -1,4 +1,5 @@
-* Changed `database/sql` driver `prepare` behaviour to `nop` with proxing call to conn exec/query with keep-in-cache flag 
+* Changed `database/sql` driver `prepare` behaviour to `nop` with proxing call to conn exec/query with keep-in-cache flag
+* Added metadata to `trace.Driver.OnInvoke` and `trace.Driver.OnNewStream` done events
 
 ## v3.34.0
 * Improved the `xsql` errors mapping to `driver.ErrBadConn` 

internal/conn/conn.go

@@ -10,6 +10,7 @@ import (
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/stats"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/backoff"
@@ -305,10 +306,11 @@ func (c *conn) Invoke(
 			trace.Method(method),
 		)
 		cc *grpc.ClientConn
+		md = metadata.MD{}
 	)
 
 	defer func() {
-		onDone(err, issues, opID, c.GetState())
+		onDone(err, issues, opID, c.GetState(), md)
 	}()
 
 	cc, err = c.take(ctx)
@@ -321,7 +323,7 @@ func (c *conn) Invoke(
 
 	ctx, sentMark := markContext(ctx)
 
-	err = cc.Invoke(ctx, method, req, res, opts...)
+	err = cc.Invoke(ctx, method, req, res, append(opts, grpc.Trailer(&md))...)
 	if err != nil {
 		if wrapping {
 			err = xerrors.FromGRPCError(err,
@@ -386,7 +388,7 @@ func (c *conn) NewStream(
 
 	defer func() {
 		if err != nil {
-			streamRecv(err)(c.GetState(), err)
+			streamRecv(err)(err, c.GetState(), metadata.MD{})
 		}
 	}()
 

internal/conn/grpc_client_stream.go

@@ -21,7 +21,7 @@ type grpcClientStream struct {
 	wrapping bool
 	sentMark *modificationMark
 	onDone   func(ctx context.Context)
-	recv     func(error) func(trace.ConnState, error)
+	recv     func(error) func(error, trace.ConnState, map[string][]string)
 }
 
 func (s *grpcClientStream) CloseSend() (err error) {
@@ -77,7 +77,7 @@ func (s *grpcClientStream) RecvMsg(m interface{}) (err error) {
 	defer func() {
 		onDone := s.recv(xerrors.HideEOF(err))
 		if err != nil {
-			onDone(s.c.GetState(), xerrors.HideEOF(err))
+			onDone(xerrors.HideEOF(err), s.c.GetState(), s.ClientStream.Trailer())
 			s.onDone(s.ClientStream.Context())
 		}
 	}()

log/driver.go

@@ -289,17 +289,19 @@ func Driver(l Logger, details trace.Details) (t trace.Driver) {
 			start := time.Now()
 			return func(info trace.DriverConnInvokeDoneInfo) {
 				if info.Error == nil {
-					l.Tracef(`invoke done {endpoint:%v,method:"%s",latency:"%v"}`,
+					l.Tracef(`invoke done {endpoint:%v,method:"%s",latency:"%v",,metadata:%v}`,
 						endpoint,
 						method,
 						time.Since(start),
+						info.Metadata,
 					)
 				} else {
-					l.Warnf(`invoke failed {endpoint:%v,method:"%s",latency:"%v",error:"%s",version:"%s"}`,
+					l.Warnf(`invoke failed {endpoint:%v,method:"%s",latency:"%v",error:"%s",,metadata:%v,version:"%s"}`,
 						endpoint,
 						method,
 						time.Since(start),
 						info.Error,
+						info.Metadata,
 						meta.Version,
 					)
 				}
@@ -337,17 +339,19 @@ func Driver(l Logger, details trace.Details) (t trace.Driver) {
 				}
 				return func(info trace.DriverConnNewStreamDoneInfo) {
 					if info.Error == nil {
-						l.Tracef(`streaming done {endpoint:%v,method:"%s",latency:"%v"}`,
+						l.Tracef(`streaming done {endpoint:%v,method:"%s",latency:"%v",metadata:%v}`,
 							endpoint,
 							method,
 							time.Since(start),
+							info.Metadata,
 						)
 					} else {
-						l.Warnf(`streaming done {endpoint:%v,method:"%s",latency:"%v",error:"%s",version:"%s"}`,
+						l.Warnf(`streaming done {endpoint:%v,method:"%s",latency:"%v",error:"%s",,metadata:%v,version:"%s"}`,
 							endpoint,
 							method,
 							time.Since(start),
 							info.Error,
+							info.Metadata,
 							meta.Version,
 						)
 					}

trace/driver.go

@@ -236,10 +236,11 @@ type (
 		Method   Method
 	}
 	DriverConnInvokeDoneInfo struct {
-		Error  error
-		Issues []Issue
-		OpID   string
-		State  ConnState
+		Error    error
+		Issues   []Issue
+		OpID     string
+		State    ConnState
+		Metadata map[string][]string
 	}
 	DriverConnNewStreamStartInfo struct {
 		// Context make available context in trace callback function.
@@ -254,8 +255,9 @@ type (
 		Error error
 	}
 	DriverConnNewStreamDoneInfo struct {
-		State ConnState
-		Error error
+		Error    error
+		State    ConnState
+		Metadata map[string][]string
 	}
 	DriverBalancerInitStartInfo struct {
 		// Context make available context in trace callback function.