Should context cancelation be handled in application code or by row.Scan/row.Next

[MatthiasKunnen]

I'm currently looking into creating an iterator, stream based, approach of consuming rows at sqlc-dev/sqlc#720 (comment).

When the context provided to conn.Query is canceled, row.Scan does not return the cancellation error and row.Next does not return false. This is consistent with github.com/go-sql-driver/mysql but it leaves me wondering why.

Example code:

package main

import (
	"context"
	"log"

	"github.com/jackc/pgx/v5"
)

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	conn, err := pgx.Connect(ctx, `connstring`) // Same behavior if this were background context
	if err != nil {
		log.Fatalf("Unable to connect to database: %v", err)
	}
	defer conn.Close(context.Background())

	rows, err := conn.Query(ctx, `SELECT generate_series(1, 10) AS row_number`)
	if err != nil {
		log.Fatalf("Query failed: %v", err)
	}
	defer rows.Close()
	for rows.Next() {
		// Why do we need to do this and does pgx not handle this?
		select {
		case <-ctx.Done():
			break
		default:
		}

		var id int
		err := rows.Scan(&id)
		if err != nil {
			log.Fatalf("Scan failed: %v", err)
		}
		log.Printf("Row returned: %d\n", id)
		// Simulate cancel by e.g. CTRL+C or some other timeout.
		// Expecting row.Next to return false or row.Scan to error on next call.
		cancel()
	}
	if err := rows.Err(); err != nil {
		log.Fatalf("error: %v", err)
	}
}

I had assumed conn.Query an implementation of QueryContext which states:

QueryContext must honor the context timeout and return when the context is canceled.

[jackc]

This is consistent with github.com/go-sql-driver/mysql but it leaves me wondering why.

Context is used to cancel potentially long running or blocking operations. Neither Rows.Next() nor Rows.Scan() are either. The context will block any further network reads. Presumably, the rows have already been read over the wire. If you were to change your generate_series query to something that returned so many rows that pgx had to read from the network multiple times then you would get a Rows.Next() => false.

[MatthiasKunnen]

Thanks you for your answer.

I can confirm that when generating a million rows, Next returns false roughly 450 to 550 iterations after cancel. I would assume that it depends on the amount of data read.

So I would conclude that pgx handles the context correctly but if processing the extra rows after a cancellation has already ocurred is undesirable, e.g. if it performs an expensive action in response to the rows, the application should perform its own context checking. Is this a fair assessment?

[jackc]

So I would conclude that pgx handles the context correctly but if processing the extra rows after a cancellation has already ocurred is undesirable, e.g. if it performs an expensive action in response to the rows, the application should perform its own context checking. Is this a fair assessment?

That is correct.