package main

// Summary: Using txn := db.Begin() followed by txn.Exec(foo),
// txn.{Commit,Rollback} as shown below does not appear to work (see log output at
// bottom).  However passing everything through db.Exec(foo) does,
// including sending db.Exec(`BEGIN`), db.Exec('ROLLBACK'), etc. (see
// code in PR at https://github.com/cockroachdb/docs/pull/5049 for the
// db.Exec version).

import (
	"fmt"
	"log"
	"math"
	"math/rand"
	"time"

	// Import GORM-related packages.
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/postgres"

	// Necessary in order to check for transaction retry error codes.
	// See implementation below in `transferFunds`.
	"github.com/lib/pq"
)

// Account is our model, which corresponds to the "accounts" database
// table.
type Account struct {
	ID      int `gorm:"primary_key"`
	Balance int
}

func main() {
	// Connect to the "bank" database as the "maxroach" user.
	const addr = "postgresql://root@localhost:26257/bank?sslmode=disable"
	db, err := gorm.Open("postgres", addr)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	// Set to `true` and GORM will print out all DB queries.
	db.LogMode(true)

	// Automatically create the "accounts" table based on the Account
	// model.
	db.AutoMigrate(&Account{})

	// Insert two rows into the "accounts" table.
	db.Create(&Account{ID: 1, Balance: 1000})
	db.Create(&Account{ID: 2, Balance: 250})

	// The sequence of steps in this section is:
	// 1. Print account balances.
	// 2. Set up some Accounts and transfer funds between them.
	// 3. Print account balances again to verify the transfer occurred.

	// Print balances before transfer.
	printBalances(db)

	var amount = 100
	var fromAccount Account
	var toAccount Account

	db.First(&fromAccount, 1)
	db.First(&toAccount, 2)

	// Transfer funds between accounts.  To handle any possible
	// transaction retry errors, we add a retry loop with exponential
	// backoff to the transfer logic (see implementation below).
	if err := transferFunds(db, fromAccount, toAccount, amount); err != nil {
		// If the error is returned, it's either:
		//   1. Not a transaction retry error, i.e., some other kind of
		//   database error that you should handle here.
		//   2. A transaction retry error that has occurred more than N
		//   times (defined by the `maxRetries` variable inside
		//   `transferFunds`), in which case you will need to figure out
		//   why your database access is resulting in so much contention
		//   (see 'Understanding and avoiding transaction contention':
		//   https://www.cockroachlabs.com/docs/stable/performance-best-practices-overview.html#understanding-and-avoiding-transaction-contention)
		fmt.Println(err)
	}

	// Print balances after transfer to ensure that it worked.
	printBalances(db)

	// Delete accounts so we can start fresh when we want to run this
	// program again.
	deleteAccounts(db)
}

func transferFunds(db *gorm.DB, fromAccount Account, toAccount Account, amount int) error {
	if fromAccount.Balance < amount {
		return fmt.Errorf("account %d balance %d is lower than transfer amount %d", fromAccount.ID, fromAccount.Balance, amount)
	}

	var maxRetries = 3
	for retries := 0; retries <= maxRetries; retries++ {
		if retries == maxRetries {
			return fmt.Errorf("hit max of %d retries, aborting", retries)
		}

		// db.Exec("BEGIN")
		txn := db.Begin()
		txn.Exec(`SELECT now()`) // disable server-side auto-retries
		if err := txn.Exec(
			`SELECT crdb_internal.force_retry('1s'::INTERVAL)`,

			// `UPSERT INTO accounts (id, balance) VALUES
			//    (?, ((SELECT balance FROM accounts WHERE id = ?) - ?)),
			//    (?, ((SELECT balance FROM accounts WHERE id = ?) + ?))`,
			// fromAccount.ID, fromAccount.ID, amount,
			// toAccount.ID, toAccount.ID, amount
		).Error; err != nil {

			// We need to cast GORM's db.Error to *pq.Error so we can
			// detect the Postgres transaction retry error code and
			// handle retries appropriately.
			pqErr := err.(*pq.Error)

			if pqErr.Code == "40001" {
				// Since this is a transaction retry error, we
				// ROLLBACK the transaction and sleep a little before
				// trying again.  Each time through the loop we sleep
				// for a little longer than the last time
				// (A.K.A. exponential backoff).
				txn.Rollback()
				// db.Exec("ROLLBACK")
				var sleepMs = math.Pow(2, float64(retries)) * 100 * (rand.Float64() + 0.5)
				time.Sleep(time.Millisecond * time.Duration(sleepMs))
			} else {
				return err
			}
		} else {
			// Happy case.  All went well, so we commit and break out
			// of the retry loop.
			txn.Commit()
			// db.Exec("COMMIT")
			break
		}
	}
	return nil
}

func printBalances(db *gorm.DB) {
	var accounts []Account
	db.Find(&accounts)
	fmt.Printf("Balance at '%s':\n", time.Now())
	for _, account := range accounts {
		fmt.Printf("%d %d\n", account.ID, account.Balance)
	}
}

func deleteAccounts(db *gorm.DB) error {
	// Used to tear down the accounts table so we can re-run this
	// program.
	err := db.Exec("DELETE from accounts where ID > 0").Error
	if err != nil {
		return err
	}
	return nil
}

// -*- mode: compilation; default-directory: "~/work/code/deemphasize-savepoints/golang/" -*-
// Compilation started at Wed Jul 24 13:30:55

// make -k 
// go run gorm-sample.go

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:43) 
// [2019-07-24 13:30:55]  [0.77ms]  INSERT INTO "accounts" ("id","balance") VALUES ('1','1000') RETURNING "accounts"."id"  
// [1 rows affected or returned ] 

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:44) 
// [2019-07-24 13:30:55]  [0.69ms]  INSERT INTO "accounts" ("id","balance") VALUES ('2','250') RETURNING "accounts"."id"  
// [1 rows affected or returned ] 

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:140) 
// [2019-07-24 13:30:55]  [0.76ms]  SELECT * FROM "accounts"    
// [2 rows affected or returned ] 
// Balance at '2019-07-24 13:30:55.617543 -0400 EDT m=+0.019774217':
// 1 1000
// 2 250

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:58) 
// [2019-07-24 13:30:55]  [0.48ms]  SELECT * FROM "accounts"  WHERE ("accounts"."id" = 1) ORDER BY "accounts"."id" ASC LIMIT 1  
// [1 rows affected or returned ] 

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:59) 
// [2019-07-24 13:30:55]  [0.50ms]  SELECT * FROM "accounts"  WHERE ("accounts"."id" = 2) ORDER BY "accounts"."id" ASC LIMIT 1  
// [1 rows affected or returned ] 

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:98) 
// [2019-07-24 13:30:55]  [0.61ms]  SELECT now()  
// [1 rows affected or returned ] 

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:99) 
// [2019-07-24 13:30:55]  pq: restart transaction: crdb_internal.force_retry(): TransactionRetryWithProtoRefreshError: forced by crdb_internal.force_retry() 

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:99) 
// [2019-07-24 13:30:55]  [0.90ms]  SELECT crdb_internal.force_retry('1s'::INTERVAL)  
// [0 rows affected or returned ] 

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:98) 
// [2019-07-24 13:30:55]  [0.79ms]  SELECT now()  
// [1 rows affected or returned ] 

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:99) 
// [2019-07-24 13:30:55]  pq: restart transaction: crdb_internal.force_retry(): TransactionRetryWithProtoRefreshError: forced by crdb_internal.force_retry() 

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:99) 
// [2019-07-24 13:30:55]  [1.07ms]  SELECT crdb_internal.force_retry('1s'::INTERVAL)  
// [0 rows affected or returned ] 

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:98) 
// [2019-07-24 13:30:56]  [0.74ms]  SELECT now()  
// [1 rows affected or returned ] 

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:99) 
// [2019-07-24 13:30:56]  pq: restart transaction: crdb_internal.force_retry(): TransactionRetryWithProtoRefreshError: forced by crdb_internal.force_retry() 

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:99) 
// [2019-07-24 13:30:56]  [1.56ms]  SELECT crdb_internal.force_retry('1s'::INTERVAL)  
// [0 rows affected or returned ] 
// hit max of 3 retries, aborting

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:140) 
// [2019-07-24 13:30:56]  [1.44ms]  SELECT * FROM "accounts"    
// [2 rows affected or returned ] 
// Balance at '2019-07-24 13:30:56.533603 -0400 EDT m=+0.935833949':
// 1 1000
// 2 250

// (/Users/rloveland/work/code/deemphasize-savepoints/golang/gorm-sample.go:150) 
// [2019-07-24 13:30:56]  [3.81ms]  DELETE from accounts where ID > 0  
// [2 rows affected or returned ] 

// Compilation finished at Wed Jul 24 13:30:56