github.com/blend/go-sdk@v1.20220411.3/examples/db/statement-timeout/main.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package main
     9  
    10  import (
    11  	"context"
    12  	"fmt"
    13  	"log"
    14  	"time"
    15  
    16  	"github.com/blend/go-sdk/db"
    17  	"github.com/blend/go-sdk/ex"
    18  )
    19  
    20  const (
    21  	longQueryTemplate = "SELECT id, pg_sleep(%f) FROM might_sleep WHERE id = 1337;"
    22  	separator         = "=================================================="
    23  )
    24  
    25  func createConn(ctx context.Context) (*db.Connection, error) {
    26  	pool, err := db.New(db.OptConfigFromEnv())
    27  	if err != nil {
    28  		return nil, err
    29  	}
    30  
    31  	err = pool.Open()
    32  	if err != nil {
    33  		return nil, err
    34  	}
    35  
    36  	err = pool.Connection.PingContext(ctx)
    37  	if err != nil {
    38  		return nil, err
    39  	}
    40  
    41  	log.Printf("DSN=%q\n", pool.Config.CreateDSN())
    42  	return pool, nil
    43  }
    44  
    45  func intentionallyLongQuery(ctx context.Context, pool *db.Connection, cfg *config) error {
    46  	type resultRow struct {
    47  		ID      int    `db:"id"`
    48  		PGSleep string `db:"pg_sleep"`
    49  	}
    50  
    51  	s := float64(cfg.PGSleep) / float64(time.Second)
    52  
    53  	statement := fmt.Sprintf(longQueryTemplate, s)
    54  	q := pool.QueryContext(ctx, statement)
    55  
    56  	r := resultRow{}
    57  	log.Println("Starting query")
    58  	found, err := q.Out(&r)
    59  	if err != nil {
    60  		return err
    61  	}
    62  	if !found {
    63  		return ex.New("`SELECT id, pg_sleep(%f) ...;` query returned no results")
    64  	}
    65  
    66  	return nil
    67  }
    68  
    69  func main() {
    70  	log.SetFlags(0)
    71  	log.SetOutput(newLogWriter())
    72  	cfg := getConfig()
    73  
    74  	// 1. Set the `DB_STATEMENT_TIMEOUT` environment variable.
    75  	log.Println(separator)
    76  	cfg.Print()
    77  	err := cfg.SetEnvironment()
    78  	if err != nil {
    79  		log.Fatal(err)
    80  	}
    81  
    82  	deadline := time.Now().Add(cfg.ContextTimeout)
    83  	ctx, cancel := context.WithDeadline(context.Background(), deadline)
    84  	defer cancel()
    85  
    86  	// 2. Parse config / open / ping
    87  	// 3. Make sure `statement_timeout` is in the connection string (it gets printed)
    88  	log.Println(separator)
    89  	pool, err := createConn(ctx)
    90  	if err != nil {
    91  		log.Fatal(err)
    92  	}
    93  	defer cleanUp(pool)
    94  
    95  	// 4. Demonstrate that the observed statement timeout on an open connection is
    96  	//    `StatementTimeout`.
    97  	log.Println(separator)
    98  	timeout, err := ensureStatementTimeout(ctx, pool, cfg)
    99  	if err != nil {
   100  		log.Fatal(err)
   101  	}
   102  	log.Printf("statement_timeout=%s\n", timeout)
   103  
   104  	// 5. Create a table schema and insert data to seed the database.
   105  	err = seedDatabase(ctx, pool)
   106  	if err != nil {
   107  		log.Fatal(err)
   108  	}
   109  
   110  	// 6. Run query that intentionally runs for a long time.
   111  	log.Println(separator)
   112  	err = intentionallyLongQuery(ctx, pool, cfg)
   113  	if err == nil {
   114  		log.Fatal(ex.New("Expected statement contention to occur"))
   115  	}
   116  
   117  	// 7. Display the error / errors in as verbose a way as possible.
   118  	log.Println("***")
   119  	err = displayError(err)
   120  	if err != nil {
   121  		log.Fatal(err)
   122  	}
   123  }