diff --git a/conn.go b/conn.go index cd9773d01..b04469f00 100644 --- 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 index bcf748b76..b876d4c13 100644 --- 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 {