Use central entities table for EEA logic (#4229)

This replaces the logic from the EEA from the convoluted per-entity column to use a single reference from the central entitites table. Closes: #4168 Signed-off-by: Juan Antonio Osorio <[email protected]>
mindersec · Aug 21, 2024 · 29c9a13 · 29c9a13
1 parent 6354e9e
commit 29c9a13
Show file tree

Hide file tree

Showing 11 changed files with 94 additions and 107 deletions.
diff --git a/database/migrations/000099_eea_entity_instance_idx.down.sql b/database/migrations/000099_eea_entity_instance_idx.down.sql
@@ -0,0 +1,22 @@
+-- Copyright 2024 Stacklok, Inc
+--
+-- Licensed under the Apache License, Version 2.0 (the "License");
+-- you may not use this file except in compliance with the License.
+-- You may obtain a copy of the License at
+--
+--      http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+
+BEGIN;
+
+-- Drop the unique index on entity_instance_id from the entity_execution_lock and flush_cache tables
+
+DROP INDEX entity_execution_lock_entity_instance_idx ON entity_execution_lock;
+DROP INDEX flush_cache_entity_instance_idx ON flush_cache;
+
+COMMIT;
diff --git a/database/migrations/000099_eea_entity_instance_idx.up.sql b/database/migrations/000099_eea_entity_instance_idx.up.sql
@@ -0,0 +1,22 @@
+-- Copyright 2024 Stacklok, Inc
+--
+-- Licensed under the Apache License, Version 2.0 (the "License");
+-- you may not use this file except in compliance with the License.
+-- You may obtain a copy of the License at
+--
+--      http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+
+BEGIN;
+
+-- Add entity_instance_id as a unique index to the entity_execution_lock and flush_cache tables
+
+CREATE UNIQUE INDEX entity_execution_lock_entity_instance_idx ON entity_execution_lock (entity_instance_id);
+CREATE UNIQUE INDEX flush_cache_entity_instance_idx ON flush_cache (entity_instance_id);
+
+COMMIT;
diff --git a/database/mock/store.go b/database/mock/store.go
diff --git a/database/query/entity_execution_lock.sql b/database/query/entity_execution_lock.sql
@@ -23,11 +23,10 @@ INSERT INTO entity_execution_lock(
     sqlc.narg(pull_request_id)::UUID,
     sqlc.arg(project_id)::UUID,
     sqlc.arg(entity_instance_id)::UUID
-) ON CONFLICT(entity, COALESCE(repository_id, '00000000-0000-0000-0000-000000000000'::UUID), COALESCE(artifact_id, '00000000-0000-0000-0000-000000000000'::UUID), COALESCE(pull_request_id, '00000000-0000-0000-0000-000000000000'::UUID))
+) ON CONFLICT(entity_instance_id)
 DO UPDATE SET
     locked_by = gen_random_uuid(),
-    last_lock_time = NOW(),
-    entity_instance_id = sqlc.arg(entity_instance_id)::UUID
+    last_lock_time = NOW()
 WHERE entity_execution_lock.last_lock_time < (NOW() - (@interval::TEXT || ' seconds')::interval)
 RETURNING *;
 
@@ -37,19 +36,11 @@ RETURNING *;
 
 -- name: ReleaseLock :exec
 DELETE FROM entity_execution_lock
-WHERE entity = sqlc.arg(entity)::entities AND
-    COALESCE(repository_id, '00000000-0000-0000-0000-000000000000'::UUID) = COALESCE(sqlc.narg(repository_id)::UUID, '00000000-0000-0000-0000-000000000000'::UUID) AND
-    COALESCE(artifact_id, '00000000-0000-0000-0000-000000000000'::UUID) = COALESCE(sqlc.narg(artifact_id)::UUID, '00000000-0000-0000-0000-000000000000'::UUID) AND
-    COALESCE(pull_request_id, '00000000-0000-0000-0000-000000000000'::UUID) = COALESCE(sqlc.narg(pull_request_id)::UUID, '00000000-0000-0000-0000-000000000000'::UUID) AND
-    locked_by = sqlc.arg(locked_by)::UUID;
+WHERE entity_instance_id = sqlc.arg(entity_instance_id) AND locked_by = sqlc.arg(locked_by)::UUID;
 
 -- name: UpdateLease :exec
 UPDATE entity_execution_lock SET last_lock_time = NOW()
-WHERE entity = $1 AND
-COALESCE(repository_id, '00000000-0000-0000-0000-000000000000'::UUID) = COALESCE(sqlc.narg(repository_id), '00000000-0000-0000-0000-000000000000'::UUID) AND
-COALESCE(artifact_id, '00000000-0000-0000-0000-000000000000'::UUID) = COALESCE(sqlc.narg(artifact_id)::UUID, '00000000-0000-0000-0000-000000000000'::UUID) AND
-COALESCE(pull_request_id, '00000000-0000-0000-0000-000000000000'::UUID) = COALESCE(sqlc.narg(pull_request_id)::UUID, '00000000-0000-0000-0000-000000000000'::UUID) AND
-locked_by = sqlc.arg(locked_by)::UUID;
+WHERE entity_instance_id = $1 AND locked_by = sqlc.arg(locked_by)::UUID;
 
 -- name: EnqueueFlush :one
 INSERT INTO flush_cache(
@@ -66,16 +57,13 @@ INSERT INTO flush_cache(
     sqlc.narg(pull_request_id)::UUID,
     sqlc.arg(project_id)::UUID,
     sqlc.arg(entity_instance_id)::UUID
-) ON CONFLICT(entity, COALESCE(repository_id, '00000000-0000-0000-0000-000000000000'::UUID), COALESCE(artifact_id, '00000000-0000-0000-0000-000000000000'::UUID), COALESCE(pull_request_id, '00000000-0000-0000-0000-000000000000'::UUID))
+) ON CONFLICT(entity_instance_id)
 DO NOTHING
 RETURNING *;
 
 -- name: FlushCache :one
 DELETE FROM flush_cache
-WHERE entity = $1 AND
-    COALESCE(repository_id, '00000000-0000-0000-0000-000000000000'::UUID) = COALESCE(sqlc.narg(repository_id)::UUID, '00000000-0000-0000-0000-000000000000'::UUID) AND
-    COALESCE(artifact_id, '00000000-0000-0000-0000-000000000000'::UUID) = COALESCE(sqlc.narg(artifact_id)::UUID, '00000000-0000-0000-0000-000000000000'::UUID) AND
-    COALESCE(pull_request_id, '00000000-0000-0000-0000-000000000000'::UUID) = COALESCE(sqlc.narg(pull_request_id)::UUID, '00000000-0000-0000-0000-000000000000'::UUID)
+WHERE entity_instance_id= $1
 RETURNING *;
 
 -- name: ListFlushCache :many

diff --git a/internal/db/entity_execution_lock.sql.go b/internal/db/entity_execution_lock.sql.go
diff --git a/internal/db/querier.go b/internal/db/querier.go
diff --git a/internal/eea/eea.go b/internal/eea/eea.go
@@ -199,7 +199,10 @@ func (e *EEA) FlushMessageHandler(msg *message.Message) error {
 		return fmt.Errorf("error unmarshalling payload: %w", err)
 	}
 
-	repoID, artifactID, pullRequestID := inf.GetEntityDBIDs()
+	eID, err := inf.GetID()
+	if err != nil {
+		return fmt.Errorf("error getting entity ID: %w", err)
+	}
 
 	logger := zerolog.Ctx(ctx).With().
 		Str("component", "EEA").
@@ -211,12 +214,7 @@ func (e *EEA) FlushMessageHandler(msg *message.Message) error {
 
 	logger.Debug().Msg("flushing event")
 
-	_, err = e.querier.FlushCache(ctx, db.FlushCacheParams{
-		Entity:        entities.EntityTypeToDB(inf.Type),
-		RepositoryID:  repoID,
-		ArtifactID:    artifactID,
-		PullRequestID: pullRequestID,
-	})
+	_, err = e.querier.FlushCache(ctx, eID)
 	// Nothing to do here. If we can't flush the cache, it means
 	// that the event has already been executed.
 	if err != nil && errors.Is(err, sql.ErrNoRows) {

diff --git a/internal/engine/eval_status.go b/internal/engine/eval_status.go
@@ -39,11 +39,16 @@ func (e *executor) createEvalStatusParams(
 	rule *models.RuleInstance,
 ) (*engif.EvalStatusParams, error) {
 	repoID, artID, prID := inf.GetEntityDBIDs()
+	eID, err := inf.GetID()
+	if err != nil {
+		return nil, fmt.Errorf("Error getting ID from entity info wrapper")
+	}
 
 	params := &engif.EvalStatusParams{
 		Rule:          rule,
 		Profile:       profile,
 		EntityType:    entities.EntityTypeToDB(inf.Type),
+		EntityID:      eID,
 		RepoID:        repoID,
 		ArtifactID:    artID,
 		PullRequestID: prID,

diff --git a/internal/engine/executor.go b/internal/engine/executor.go
@@ -246,11 +246,8 @@ func (e *executor) updateLockLease(
 		zerolog.Ctx(ctx).With().Str("execution_id", executionID.String()).Logger())
 
 	if err := e.querier.UpdateLease(ctx, db.UpdateLeaseParams{
-		Entity:        params.EntityType,
-		RepositoryID:  params.RepoID,
-		ArtifactID:    params.ArtifactID,
-		PullRequestID: params.PullRequestID,
-		LockedBy:      executionID,
+		LockedBy:         executionID,
+		EntityInstanceID: params.EntityID,
 	}); err != nil {
 		logger.Err(err).Msg("error updating lock lease")
 		return
@@ -264,10 +261,18 @@ func (e *executor) releaseLockAndFlush(
 	inf *entities.EntityInfoWrapper,
 ) {
 	repoID, artID, prID := inf.GetEntityDBIDs()
+	eID, err := inf.GetID()
+	if err != nil {
+		zerolog.Ctx(ctx).Error().Err(err).Msg("error getting entity id")
+		return
+	}
 
 	logger := zerolog.Ctx(ctx).Info().
 		Str("entity_type", inf.Type.ToString()).
-		Str("execution_id", inf.ExecutionID.String())
+		Str("execution_id", inf.ExecutionID.String()).
+		Str("entity_id", eID.String())
+
+	// TODO: change these to entity_id
 	if repoID.Valid {
 		logger = logger.Str("repo_id", repoID.UUID.String())
 	}
@@ -280,11 +285,8 @@ func (e *executor) releaseLockAndFlush(
 	}
 
 	if err := e.querier.ReleaseLock(ctx, db.ReleaseLockParams{
-		Entity:        entities.EntityTypeToDB(inf.Type),
-		RepositoryID:  repoID,
-		ArtifactID:    artID,
-		PullRequestID: prID,
-		LockedBy:      *inf.ExecutionID,
+		EntityInstanceID: eID,
+		LockedBy:         *inf.ExecutionID,
 	}); err != nil {
 		logger.Err(err).Msg("error updating lock lease")
 	}

diff --git a/internal/engine/executor_test.go b/internal/engine/executor_test.go
@@ -256,24 +256,15 @@ default allow = true`,
 	// Mock update lease for lock
 	mockStore.EXPECT().
 		UpdateLease(gomock.Any(), db.UpdateLeaseParams{
-			Entity: db.EntitiesRepository,
-			RepositoryID: uuid.NullUUID{
-				UUID:  repositoryID,
-				Valid: true,
-			},
-			ArtifactID:    uuid.NullUUID{},
-			PullRequestID: uuid.NullUUID{},
-			LockedBy:      executionID,
+			EntityInstanceID: repositoryID,
+			LockedBy:         executionID,
 		}).Return(nil)
 
 	// Mock release lock
 	mockStore.EXPECT().
 		ReleaseLock(gomock.Any(), db.ReleaseLockParams{
-			Entity:        db.EntitiesRepository,
-			RepositoryID:  uuid.NullUUID{UUID: repositoryID, Valid: true},
-			ArtifactID:    uuid.NullUUID{},
-			PullRequestID: uuid.NullUUID{},
-			LockedBy:      executionID,
+			EntityInstanceID: repositoryID,
+			LockedBy:         executionID,
 		}).Return(nil)
 
 	// -- end expectations

diff --git a/internal/engine/interfaces/interface.go b/internal/engine/interfaces/interface.go
@@ -117,6 +117,7 @@ type EvalStatusParams struct {
 	TaskRunID        uuid.UUID
 	BuildID          uuid.UUID
 	EntityType       db.Entities
+	EntityID         uuid.UUID
 	EvalStatusFromDb *db.ListRuleEvaluationsByProfileIdRow
 	evalErr          error
 	actionsOnOff     map[ActionType]models.ActionOpt