Only retry an INSERT operation on unique constraint violation

The retry logic is used to generate an unique ID, but if there is another error the DB code can fall into an infinite loop. Signed-off-by: IWAMOTO Toshihiro <[email protected]>
kubeflow · Oct 29, 2018 · 4f84dac · 4f84dac
1 parent 0bc5182
commit 4f84dac
Showing 1 changed file with 22 additions and 23 deletions.
diff --git a/pkg/db/interface.go b/pkg/db/interface.go
@@ -109,6 +109,14 @@ func generateRandid() string {
 	return string(rs1Letters[rand.Intn(len(rs1Letters))]) + fmt.Sprintf("%016x", id)[1:]
 }
 
+func isDBDuplicateError(err error) bool {
+	errmsg := strings.ToLower(err.Error())
+	if strings.Contains(errmsg, "unique") || strings.Contains(errmsg, "duplicate") {
+		return true
+	}
+	return false
+}
+
 func (d *dbConn) GetStudyConfig(id string) (*api.StudyConfig, error) {
 	row := d.db.QueryRow("SELECT * FROM studies WHERE id = ?", id)
 
@@ -240,13 +248,10 @@ func (d *dbConn) CreateStudy(in *api.StudyConfig) (string, error) {
 		)
 		if err == nil {
 			break
-		} else {
-			errmsg := strings.ToLower(err.Error())
-			if strings.Contains(errmsg, "unique") || strings.Contains(errmsg, "duplicate") {
-				i--
-				if i > 0 {
-					continue
-				}
+		} else if isDBDuplicateError(err) {
+			i--
+			if i > 0 {
+				continue
 			}
 		}
 		return "", err
@@ -387,13 +392,10 @@ func (d *dbConn) CreateTrial(trial *api.Trial) error {
 		if err == nil {
 			trial.TrialId = trialID
 			break
-		} else {
-			errmsg := strings.ToLower(err.Error())
-			if strings.Contains(errmsg, "unique") || strings.Contains(errmsg, "duplicate") {
-				i--
-				if i > 0 {
-					continue
-				}
+		} else if isDBDuplicateError(err) {
+			i--
+			if i > 0 {
+				continue
 			}
 		}
 		return err
@@ -728,13 +730,10 @@ func (d *dbConn) CreateWorker(worker *api.Worker) (string, error) {
 		if err == nil {
 			worker.WorkerId = workerID
 			break
-		} else {
-			errmsg := strings.ToLower(err.Error())
-			if strings.Contains(errmsg, "unique") || strings.Contains(errmsg, "duplicate") {
-				i--
-				if i > 0 {
-					continue
-				}
+		} else if isDBDuplicateError(err) {
+			i--
+			if i > 0 {
+				continue
 			}
 		}
 		return "", err
@@ -915,7 +914,7 @@ func (d *dbConn) SetSuggestionParam(algorithm string, studyID string, params []*
 		paramID = generateRandid()
 		_, err = d.db.Exec("INSERT INTO suggestion_param VALUES (?, ?, ?, ?)",
 			paramID, algorithm, studyID, strings.Join(ps, ",\n"))
-		if err == nil {
+		if err == nil || !isDBDuplicateError(err) {
 			break
 		}
 	}
@@ -1023,7 +1022,7 @@ func (d *dbConn) SetEarlyStopParam(algorithm string, studyID string, params []*a
 		paramID := generateRandid()
 		_, err = d.db.Exec("INSERT INTO earlystopping_param VALUES (?,?, ?, ?)",
 			paramID, algorithm, studyID, strings.Join(ps, ",\n"))
-		if err == nil {
+		if err == nil || !isDBDuplicateError(err) {
 			break
 		}
 	}