From 247b6f55d01a372cad2b43a57d192774bbc19ccf Mon Sep 17 00:00:00 2001
From: Chris Roberts <croberts@hashicorp.com>
Date: Wed, 28 Sep 2022 08:36:29 -0700
Subject: [PATCH] Add upsert helpers

---
 internal/server/singleprocess/state/state.go | 36 +++++++++++++++-----
 1 file changed, 28 insertions(+), 8 deletions(-)

diff --git a/internal/server/singleprocess/state/state.go b/internal/server/singleprocess/state/state.go
index cb4ba827e..fa44fba03 100644
--- a/internal/server/singleprocess/state/state.go
+++ b/internal/server/singleprocess/state/state.go
@@ -167,14 +167,6 @@ func stateStoreSchema() *memdb.DBSchema {
 	return db
 }
 
-// Provides db for searching
-// NOTE: In most cases this should be used instead of accessing `db`
-// directly when searching for values to ensure all associations are
-// fully loaded in the results.
-func (s *State) search() *gorm.DB {
-	return s.db.Preload(clause.Associations)
-}
-
 // Convert error to a GRPC status error when dealing with lookups
 func lookupErrorToStatus(
 	typeName string, // thing trying to be found (basis, project, etc)
@@ -230,3 +222,31 @@ func errorToStatus(
 
 	return status.Error(codes.Internal, err.Error())
 }
+
+// Provides db for searching
+// NOTE: In most cases this should be used instead of accessing `db`
+// directly when searching for values to ensure all associations are
+// fully loaded in the results.
+func (s *State) search() *gorm.DB {
+	return s.db.Preload(clause.Associations)
+}
+
+func (s *State) upsert(arg interface{}) error {
+	return s.db.Clauses(
+		clause.OnConflict{
+			UpdateAll: true,
+		},
+	).Save(arg).Error
+}
+
+func (s *State) upsertFull(arg interface{}) error {
+	return s.db.Session(
+		&gorm.Session{
+			FullSaveAssociations: true,
+		},
+	).Clauses(
+		clause.OnConflict{
+			UpdateAll: true,
+		},
+	).Save(arg).Error
+}