Storage: GuaranteedUpdate fix & other improvements (#85206)

make GuaranteedUpdate work when ignoring not found errors, increase poll frequency, fix Delete

pkg/services/apiserver/storage/entity/storage.go

@@ -511,9 +511,14 @@ func (s *Storage) GuaranteedUpdate(
 		Subresource: requestInfo.Subresource,
 	}
 
-	err := s.Get(ctx, k.String(), storage.GetOptions{}, destination)
+	getErr := s.Get(ctx, k.String(), storage.GetOptions{}, destination)
-	if err != nil {
+	if getErr != nil {
-		return err
+		if ignoreNotFound && apierrors.IsNotFound(getErr) {
+			// destination is already set to zero value
+			// we'll create the resource
+		} else {
+			return getErr
+		}
 	}
 
 	accessor, err := meta.Accessor(destination)
@@ -544,6 +549,27 @@ func (s *Storage) GuaranteedUpdate(
 		return err
 	}
 
+	// if we have a non-nil getErr, then we've ignored a not found error
+	if getErr != nil {
+		// object does not exist, create it
+		req := &entityStore.CreateEntityRequest{
+			Entity: e,
+		}
+
+		rsp, err := s.store.Create(ctx, req)
+		if err != nil {
+			return err
+		}
+
+		err = entityToResource(rsp.Entity, destination, s.codec)
+		if err != nil {
+			return apierrors.NewInternalError(err)
+		}
+
+		return nil
+	}
+
+	// update the existing object
 	req := &entityStore.UpdateEntityRequest{
 		Entity:          e,
 		PreviousVersion: previousVersion,

pkg/services/store/entity/sqlstash/sql_storage_server.go

@@ -453,6 +453,8 @@ func (s *sqlEntityServer) Create(ctx context.Context, r *entity.CreateEntityRequ
 		// Update resource version
 		current.ResourceVersion = s.snowflake.Generate().Int64()
 
+		current.Action = entity.Entity_CREATED
+
 		values := map[string]any{
 			"guid":             current.Guid,
 			"key":              current.Key,
@@ -482,7 +484,7 @@ func (s *sqlEntityServer) Create(ctx context.Context, r *entity.CreateEntityRequ
 			"origin_key":       current.Origin.Key,
 			"origin_ts":        current.Origin.Time,
 			"message":          current.Message,
-			"action":           entity.Entity_CREATED,
+			"action":           current.Action,
 		}
 
 		// 1. Add row to the `entity_history` values
@@ -658,6 +660,8 @@ func (s *sqlEntityServer) Update(ctx context.Context, r *entity.UpdateEntityRequ
 		// Update resource version
 		current.ResourceVersion = s.snowflake.Generate().Int64()
 
+		current.Action = entity.Entity_UPDATED
+
 		values := map[string]any{
 			// below are only set in history table
 			"guid":       current.Guid,
@@ -689,7 +693,7 @@ func (s *sqlEntityServer) Update(ctx context.Context, r *entity.UpdateEntityRequ
 			"origin_key":       current.Origin.Key,
 			"origin_ts":        current.Origin.Time,
 			"message":          current.Message,
-			"action":           entity.Entity_UPDATED,
+			"action":           current.Action,
 		}
 
 		// 1. Add the `entity_history` values
@@ -709,7 +713,6 @@ func (s *sqlEntityServer) Update(ctx context.Context, r *entity.UpdateEntityRequ
 		delete(values, "name")
 		delete(values, "created_at")
 		delete(values, "created_by")
-		delete(values, "action")
 
 		err = s.dialect.Update(
 			ctx,
@@ -825,6 +828,8 @@ func (s *sqlEntityServer) doDelete(ctx context.Context, tx *session.SessionTx, e
 	// Update resource version
 	ent.ResourceVersion = s.snowflake.Generate().Int64()
 
+	ent.Action = entity.Entity_DELETED
+
 	// Set updated at/by
 	ent.UpdatedAt = time.Now().UnixMilli()
 	modifier, err := appcontext.User(ctx)
@@ -854,6 +859,10 @@ func (s *sqlEntityServer) doDelete(ctx context.Context, tx *session.SessionTx, e
 		return err
 	}
 
+	if ent.Origin == nil {
+		ent.Origin = &entity.EntityOriginInfo{}
+	}
+
 	values := map[string]any{
 		// below are only set in history table
 		"guid":       ent.Guid,
@@ -885,7 +894,7 @@ func (s *sqlEntityServer) doDelete(ctx context.Context, tx *session.SessionTx, e
 		"origin_key":       ent.Origin.Key,
 		"origin_ts":        ent.Origin.Time,
 		"message":          ent.Message,
-		"action":           entity.Entity_DELETED,
+		"action":           ent.Action,
 	}
 
 	// 1. Add the `entity_history` values
@@ -1429,7 +1438,7 @@ func (s *sqlEntityServer) poller(stream chan *entity.Entity) {
 	var err error
 	since := s.snowflake.Generate().Int64()
 
-	t := time.NewTicker(5 * time.Second)
+	t := time.NewTicker(1 * time.Second)
 	defer t.Stop()
 
 	for range t.C {

pkg/services/store/entity/tests/common_test.go

@@ -58,7 +58,7 @@ func createServiceAccountAdminToken(t *testing.T, env *server.TestEnv) (string,
 
 type testContext struct {
 	authToken string
-	client    entity.EntityStoreServer
+	client    entity.EntityStoreClient
 	user      *user.SignedInUser
 	ctx       context.Context
 }
@@ -88,9 +88,11 @@ func createTestContext(t *testing.T) testContext {
 	store, err := sqlstash.ProvideSQLEntityServer(eDB)
 	require.NoError(t, err)
 
+	client := entity.NewEntityStoreClientLocal(store)
+
 	return testContext{
 		authToken: authToken,
-		client:    store,
+		client:    client,
 		user:      serviceAccountUser,
 		ctx:       appcontext.WithUser(context.Background(), serviceAccountUser),
 	}