internal/database: support bulk upsert

Add DB.BulkUpsert, which adds an ON CONFLICT clause to the INSERT
that replaces existing column values.

Change-Id: I59f36be0bcb0c0854f42da489e265f2a1396c439
Reviewed-on: https://team-review.git.corp.google.com/c/golang/discovery/+/766360
Reviewed-by: Julie Qiu <julieqiu@google.com>

1 files changed, 51 insertions, 9 deletions

diff --git a/internal/database/database_test.go b/internal/database/database_test.go
index 2ff6f9da..cb6acbd9 100644
--- a/internal/database/database_test.go
+++ b/internal/database/database_test.go
@@ -83,14 +83,6 @@ func TestBulkInsert(t *testing.T) {
 		},
 		{
 
-			name:           "test-conflict-no-action-true",
-			columns:        []string{"colA"},
-			values:         []interface{}{"valueA", "valueA"},
-			conflictAction: OnConflictDoNothing,
-			wantCount:      1,
-		},
-		{
-
 			name:         "insert-returning",
 			columns:      []string{"colA", "colB"},
 			values:       []interface{}{"valueA1", "valueB1", "valueA2", "valueB2"},
@@ -99,13 +91,21 @@ func TestBulkInsert(t *testing.T) {
 		},
 		{
 
-			name:    "test-conflict-no-action-false",
+			name:    "test-conflict",
 			columns: []string{"colA"},
 			values:  []interface{}{"valueA", "valueA"},
 			wantErr: true,
 		},
 		{
 
+			name:           "test-conflict-do-nothing",
+			columns:        []string{"colA"},
+			values:         []interface{}{"valueA", "valueA"},
+			conflictAction: OnConflictDoNothing,
+			wantCount:      1,
+		},
+		{
+
 			// This should execute the statement
 			// INSERT INTO series (path) VALUES ('''); TRUNCATE series CASCADE;)');
 			// which will insert a row with path value:
@@ -219,6 +219,48 @@ func TestLargeBulkInsert(t *testing.T) {
 	}
 }
 
+func TestBulkUpsert(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), testTimeout*3)
+	defer cancel()
+	if _, err := testDB.Exec(ctx, `CREATE TEMPORARY TABLE test_replace (C1 int PRIMARY KEY, C2 int);`); err != nil {
+		t.Fatal(err)
+	}
+	for _, values := range [][]interface{}{
+		{2, 4, 4, 8},                 // First, insert some rows.
+		{1, -1, 2, -2, 3, -3, 4, -4}, // Then replace those rows while inserting others.
+	} {
+		err := testDB.Transact(ctx, sql.LevelDefault, func(tx *DB) error {
+			return tx.BulkUpsert(ctx, "test_replace", []string{"C1", "C2"}, values, []string{"C1"})
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		var got []interface{}
+		err = testDB.RunQuery(ctx, `SELECT C1, C2 FROM test_replace ORDER BY C1`, func(rows *sql.Rows) error {
+			var a, b int
+			if err := rows.Scan(&a, &b); err != nil {
+				return err
+			}
+			got = append(got, a, b)
+			return nil
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !cmp.Equal(got, values) {
+			t.Errorf("%v: got %v, want %v", values, got, values)
+		}
+	}
+}
+
+func TestBuildUpsertConflictAction(t *testing.T) {
+	got := buildUpsertConflictAction([]string{"a", "b"}, []string{"c", "d"})
+	want := "ON CONFLICT (c, d) DO UPDATE SET a=excluded.a, b=excluded.b"
+	if got != want {
+		t.Errorf("got %q, want %q", got, want)
+	}
+}
+
 func TestDBAfterTransactFails(t *testing.T) {
 	ctx := context.Background()
 	var tx *DB