Why transaction timeout in pgx doesn't work

I used to use the pq library, now I want to use pgx. I have a task to abort a transaction if it has been running in the db for too long. pq is able to abort a transaction as soon as the context has expired, but pgx waits for the query to be completed and only returns an error on the next query or commit.

Is this how it should be, or am I doing something wrong?

How can I terminate the transaction in this case?

import (
    "context"
    "database/sql"
    "fmt"
    _ "github.com/jackc/pgx/v5/stdlib"
    _ "github.com/lib/pq"
    "log"
    "time"
)

func main() {
    fmt.Println("================= PQ ==============")
    pqWork()
    fmt.Println("================= ==============")

    fmt.Println("================= PGX ==============")
    pqxWork()
    fmt.Println("================= ==============")
}

func pqWork() {
    connStr := "host=localhost port=5432 user=... password=... dbname=... "
    db, err := sql.Open("postgres", connStr)
    if err != nil {
        panic(err)
    }
    defer db.Close()

    exec(db)
}

func pqxWork() {
    connStr := "host=localhost port=5432 user=... password=... dbname=... "
    db, err := sql.Open("pgx", connStr)
    if err != nil {
        panic(err)
    }
    defer db.Close()

    exec(db)
}

func exec(db *sql.DB) {
    ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
    defer cancel()
    begin := time.Now()
    tx, err := db.BeginTx(ctx, nil)
    if err != nil {
        panic(err)
    }

    _, err = tx.Exec("select pg_sleep(10)")
    if err != nil {
        fmt.Println("err. time:", time.Since(begin).Milliseconds(), "ms")
        tx.Rollback()
        return
    }
    fmt.Println("success. time:", time.Since(begin).Milliseconds(), "ms")
    err = tx.Commit()
    if err != nil {
        panic(err)
    }
}

Result

================= PQ ==============
err. time:  524 ms
================= ==============
================= PGX ==============
success. time:  10032 ms
panic: sql: transaction has already been committed or rolled back

From `Exec` description:

Exec executes a query that doesn't return rows. For example: an INSERT and UPDATE.

Exec uses [context.Background] internally; to specify the context, use [Tx.ExecContext].

So you just need to replace Exec call with ExecContext like this and that does the trick:

_, err = tx.ExecContext(ctx, "select pg_sleep(10)")