From 7e960214287ad86acf84c31f7b5952d4bf5fd1b9 Mon Sep 17 00:00:00 2001
From: Alex Suraci <suraci.alex@gmail.com>
Date: Thu, 15 Nov 2018 13:17:42 -0500
Subject: [PATCH] retry on serialization errors

---
 storage/sql/sql.go | 39 +++++++++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/storage/sql/sql.go b/storage/sql/sql.go
index dc6be4a1..ddcfae9e 100644
--- a/storage/sql/sql.go
+++ b/storage/sql/sql.go
@@ -6,10 +6,10 @@ import (
 	"regexp"
 	"time"
 
+	"github.com/lib/pq"
 	"github.com/sirupsen/logrus"
 
 	// import third party drivers
-	_ "github.com/lib/pq"
 	_ "github.com/mattn/go-sqlite3"
 )
 
@@ -51,19 +51,34 @@ var (
 		// NOTE(ericchiang): For some reason using `SET SESSION CHARACTERISTICS AS TRANSACTION` at a
 		// session level didn't work for some edge cases. Might be something worth exploring.
 		executeTx: func(db *sql.DB, fn func(sqlTx *sql.Tx) error) error {
-			tx, err := db.Begin()
-			if err != nil {
-				return err
-			}
-			defer tx.Rollback()
+			for {
+				tx, err := db.Begin()
+				if err != nil {
+					return err
+				}
 
-			if _, err := tx.Exec(`SET TRANSACTION ISOLATION LEVEL SERIALIZABLE;`); err != nil {
-				return err
+				defer tx.Rollback()
+
+				if _, err := tx.Exec(`SET TRANSACTION ISOLATION LEVEL SERIALIZABLE;`); err != nil {
+					return err
+				}
+
+				if err := fn(tx); err != nil {
+					return err
+				}
+
+				err = tx.Commit()
+				if err != nil {
+					if pqErr, ok := err.(*pq.Error); ok && pqErr.Code == "40001" {
+						// serialization error; retry
+						continue
+					}
+
+					return err
+				}
+
+				return nil
 			}
-			if err := fn(tx); err != nil {
-				return err
-			}
-			return tx.Commit()
 		},
 
 		supportsTimezones: true,