prestodb/presto-go-client

Support of context canceling

Opened this issue · 2 comments

Is there any plan for support of canceling with context?
This feature is implemented on other database/sql drivers like on this one for PG.

A common scenario for the usage of this feature is when a user closes the request but the query remains on running state and it will cause to occupying system resources on heavy queries.

@meysampg I have a similar use case, and this is what I'm currently doing that gets the job done, hope it helps!

func doQuery(ctx context.Context, db *sql.DB, query string) error {
    // prepend a unique identifier to the query string
    query = fmt.Sprintf("/* my-program-%s */ %s", generateID(), query)
    rows, err := db.QueryContext(ctx, query)
    if err != nil {
        // ...
    }
    defer rows.Close()

    // Grab the ID of the query from the runtime table that stores all active running queries
    // This is why we used the unique identifier, so that way two queries with the same SQL won't
    // collide on this table.
    var queryID string
    if err := db.QueryRow("SELECT query_id FROM system.runtime.queries WHERE query = ?', query).Scan(&queryID); err != nil {
    // ...
    }
    for rows.Next() {
        // do what you will
    }
    // Handle the error returned from context cancellation and deadline expiration separately
    // from normal errors so we can run the cancel logic.
    // Can even handle them separately if you'd like.
    if err := rows.Err(); err == context.Canceled || err == context.DeadlineExceeded {
        killQuery := fmt.Sprintf("CALL system.runtime.kill_query('%s')", queryID)
        if err := db.Exec(killQuery); err != nil {
            // something went wrong cancelling
            return err
        }
     } else if err != nil {
         // all other errors
        return err
     }

     return nil
}

@ryantking Brilliant! Thanks!