Merge "Pass request context to frame observer" into kiwicom/master

kiwicom · Oct 30, 2023 · d7209ab · d7209ab
2 parents fb119cd + ab841b4
commit d7209ab
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/conn.go b/conn.go
@@ -1228,6 +1228,9 @@ func (c *Conn) exec(ctx context.Context, req frameBuilder, tracer Tracer) (*fram
 		if v := resp.framer.header.version.version(); v != c.version {
 			return nil, NewErrProtocol("unexpected protocol version in response: got %d expected %d", v, c.version)
 		}
+		if resp.framer != nil {
+			resp.framer.observer.reqCtx = ctx
+		}
 
 		return resp.framer, nil
 	case <-timeoutCh:

diff --git a/frame.go b/frame.go
@@ -410,6 +410,7 @@ type framer struct {
 type frameParseObserver struct {
 	head          ObservedFrameHeader
 	frameObserver FrameObserver
+	reqCtx        context.Context
 }
 
 func (fpo *frameParseObserver) observeFrame(ff *framer, f frame) {
@@ -425,7 +426,11 @@ func (fpo *frameParseObserver) observeFrame(ff *framer, f frame) {
 		of.RowCount = rows.numRows
 		of.RowsSize = rows.rowsContentSize
 	}
-	fpo.frameObserver.ObserveFrame(context.TODO(), of)
+	ctx := fpo.reqCtx
+	if ctx == nil {
+		ctx = context.TODO()
+	}
+	fpo.frameObserver.ObserveFrame(ctx, of)
 }
 
 func newFramer(compressor Compressor, version byte) *framer {