package main

import (
    "database/sql"
    "fmt"
    "log"
    "net/http"

    _ "github.com/lib/pq"
)

var db *sql.DB

func slowQuery() error {
    _, err := db.Exec("SELECT pg_sleep(10)")
    return err
}

func main() {
    var err error

    db, err = sql.Open("postgres", "postgres://user:pa$$word@localhost/example_db")
    if err != nil {
        log.Fatal(err)
    }

    if err = db.Ping(); err != nil {
        log.Fatal(err)
    }

    mux := http.NewServeMux()
    mux.HandleFunc("/", exampleHandler)

    log.Println("Listening...")
    err = http.ListenAndServe(":5000", mux)
    if err != nil {
        log.Fatal(err)
    }
}

func exampleHandler(w http.ResponseWriter, r *http.Request) {
    err := slowQuery()
    if err != nil {
        serverError(w, err)
        return
    }

    fmt.Fprintln(w, "OK")
}

func serverError(w http.ResponseWriter, err error) {
    log.Printf("ERROR: %s", err.Error())
    http.Error(w, "Sorry, something went wrong", http.StatusInternalServerError)
}

package main

import (
    "context" // New import
    "database/sql"
    "fmt"
    "log"
    "net/http"
    "time" // New import

    _ "github.com/lib/pq"
)

var db *sql.DB

func slowQuery(ctx context.Context) error {
    // Create a new child context with a 5-second timeout, using the
    // provided ctx parameter as the parent.
    ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
    defer cancel()

    // Pass the child context (the one with the timeout) as the first
    // parameter to ExecContext().
    _, err := db.ExecContext(ctx, "SELECT pg_sleep(10)")
    return err
}

...

func exampleHandler(w http.ResponseWriter, r *http.Request) {
    // Pass the request context to slowQuery(), so it can be used as the 
    // parent context.
    err := slowQuery(r.Context())
    if err != nil {
        serverError(w, err)
        return
    }

    fmt.Fprintln(w, "OK")
}

...

package main

import (
    "context"
    "database/sql"
    "errors" // New import
    "fmt"
    "log"
    "net/http"
    "time"

    _ "github.com/lib/pq"
)

var db *sql.DB

func slowQuery(ctx context.Context) error {
    ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
    defer cancel()

    _, err := db.ExecContext(ctx, "SELECT pg_sleep(10)")
    // If we get a "pq: canceling statement..." error wrap it with the 
    // context error before returning.
    if err != nil && err.Error() == "pq: canceling statement due to user request" {
        return fmt.Errorf("%w: %v", ctx.Err(), err)
    }

    return err
}

...

func exampleHandler(w http.ResponseWriter, r *http.Request) {
    err := slowQuery(r.Context())
    if err != nil {
        // Check if the returned error equals or wraps context.Canceled and 
        // record a warning if it does.
        switch {
        case errors.Is(err, context.Canceled):
            serverWarning(err)
        default:
            serverError(w, err)
        }
        return
    }

    fmt.Fprintln(w, "OK")
}

func serverWarning(err error) {
    log.Printf("WARNING: %s", err.Error())
}

...

...

func main() {
    var err error

    db, err = sql.Open("postgres", "postgres://user:pa$$word@localhost/example_db")
    if err != nil {
        log.Fatal(err)
    }

    // Create a context with a 10-second timeout, using the empty 
    // context.Background() as the parent.
    ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
    defer cancel()

    // Use this when testing the connection pool.
    if err = db.PingContext(ctx); err != nil {
        log.Fatal(err)
    }

    mux := http.NewServeMux()
    mux.HandleFunc("/", exampleHandler)

    log.Println("Listening...")
    err = http.ListenAndServe(":5000", mux)
    if err != nil {
        log.Fatal(err)
    }
}

...

func setTimeout(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
        defer cancel()

        // This gives you a copy of the request with a the request context 
        // changed to the new context with the 5-second timeout created 
        // above.
        r = r.WithContext(ctx)
        next.ServeHTTP(w, r)
    })
}