fix SLO test database-sql-query (#1836)

Co-authored-by: Konstantin Prokopenko <[email protected]>
Co-authored-by: Aleksey Myasnikov <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -1,3 +1,6 @@
+* Fixed several issues in `database/sql` error handling
+* Improved context management for query result streams.
+
 ## v3.113.0
 * Added experimental `config.WithDisableOptimisticUnban` option to disable fast node unban after pessimization
 * Fixed respect start offset from `topicoptions.WithReaderGetPartitionStartOffset` for commit messages

internal/query/client.go

@@ -337,7 +337,7 @@ func (c *Client) QueryRow(ctx context.Context, q string, opts ...options.Execute
 		onDone(finalErr)
 	}()
 
-	row, err := clientQueryRow(ctx, c.pool(), q, settings, withTrace(c.config.Trace()))
+	row, err := clientQueryRow(ctx, c.pool(), q, settings, withStreamResultTrace(c.config.Trace()))
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}
@@ -348,7 +348,7 @@ func (c *Client) QueryRow(ctx context.Context, q string, opts ...options.Execute
 func clientExec(ctx context.Context, pool sessionPool, q string, opts ...options.Execute) (finalErr error) {
 	settings := options.ExecuteSettings(opts...)
 	err := do(ctx, pool, func(ctx context.Context, s *Session) (err error) {
-		streamResult, err := s.execute(ctx, q, settings, withTrace(s.trace))
+		streamResult, err := s.execute(ctx, q, settings, withStreamResultTrace(s.trace))
 		if err != nil {
 			return xerrors.WithStackTrace(err)
 		}
@@ -397,7 +397,7 @@ func clientQuery(ctx context.Context, pool sessionPool, q string, opts ...option
 ) {
 	settings := options.ExecuteSettings(opts...)
 	err = do(ctx, pool, func(ctx context.Context, s *Session) (err error) {
-		streamResult, err := s.execute(ctx, q, options.ExecuteSettings(opts...), withTrace(s.trace))
+		streamResult, err := s.execute(ctx, q, options.ExecuteSettings(opts...), withStreamResultTrace(s.trace))
 		if err != nil {
 			return xerrors.WithStackTrace(err)
 		}
@@ -487,7 +487,7 @@ func (c *Client) QueryResultSet(
 		onDone(finalErr, rowsCount)
 	}()
 
-	rs, rowsCount, err = clientQueryResultSet(ctx, c.pool(), q, settings, withTrace(c.config.Trace()))
+	rs, rowsCount, err = clientQueryResultSet(ctx, c.pool(), q, settings, withStreamResultTrace(c.config.Trace()))
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}

internal/query/execute_query.go

@@ -131,8 +131,8 @@ func execute(
 	}
 
 	r, err := newResult(ctx, stream, append(opts,
-		withStatsCallback(settings.StatsCallback()),
-		withOnClose(executeCancel),
+		withStreamResultStatsCallback(settings.StatsCallback()),
+		withStreamResultOnClose(executeCancel),
 	)...)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)