Golang sql 查詢逾時設定 query with timeout

Go 設定sql查詢超過一定時間後自動取消的方式如下。

範例環境：

• Go 1.17
• PostgreSQL 14

範例

例如下面GetAllEmployees()在查詢PostgreSQL的employee時利用pg_sleep(5)延遲查詢5秒造成逾時主動停止查詢並返回錯誤。

利用Context.WithTimeout()建立timeout 3秒的Context參數。defer返回的cancel CancelFunc確保離開函式時釋放Context持有的資源。

呼叫DB.QueryContext()傳入timeout Context參數執行資料庫查詢，若查詢逾時會自動取消查詢並返回錯誤pq: canceling statement due to user request。

main.go

package main

import (
    "context"
    "database/sql"
    "fmt"
    "time"

    _ "github.com/lib/pq"
)

type Employee struct {
    ID        int64
    Name      string
    Age       int
}

const (
    HOST     = "localhost"
    PORT     = "5432"
    DATABASE = "postgres"
    USER     = "admin"
    PASSWORD = "12345"
    SSL      = "disable"
)

func connect() *sql.DB {
    driver := "postgres"
    dsn := fmt.Sprintf(
        "host=%s port=%s user=%s password=%s dbname=%s sslmode=%s",
        HOST, PORT, USER, PASSWORD, DATABASE, SSL)

    db, err := sql.Open(driver, dsn)
    if err != nil {
        panic("open database error")
    }
    return db
}

func main() {
    ctx := context.Background()
    db := connect()
    emps, err := GetAllEmployees(ctx, db)
    if err != nil {
        panic(err)
    }
    fmt.Println(emps)
}

func GetAllEmployees(ctx context.Context, db *sql.DB) ([]Employee, error) {
    queryCtx, cancel := context.WithTimeout(ctx, 3*time.Second) // create timeout Context
    defer cancel()                                              // make sure release resources held by Context when exiting function

    rows, err := db.QueryContext(queryCtx, "SELECT *, pg_sleep(5) FROM employee") // delay 5 seconcds query execution
    if err != nil {
        fmt.Println("query timeout")
        return nil, err
    }
    defer rows.Close()

    var emps []Employee
    for rows.Next() {
        var e Employee
        err = rows.Scan(&e.ID, &e.Name, &e.Age)
        if err != nil {
            return nil, err
        }
        emps = append(emps, e)
    }
    return emps, nil
}

執行印出以下。

query timeout
panic: pq: canceling statement due to user request

goroutine 1 [running]:
main.main()
        /../main.go:46 +0xa5
exit status 2

github。

• Canceling in-progress operations
• Golang sql PostgreSQL查詢範例