Use single and unique entity ID to deal with entity info wrapper (#4704)

* Use single and unique entity ID to deal with entity info wrapper This drops our reliance on a per-entity-type ID and instead just uses a single entity ID in the entity info wrapper. The intention is to simplify the code and make adding new entities simpler. Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com> Fix entity ID serialized in wrong place and data race issue Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com> Use local error variable for EEA test Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com> Add more guardrails Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com> fix more unit tests Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com> remove unused function Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com> fix tests Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com> * Use entity ID instead of repo ID in message meta in gh handler test Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com> --------- Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com>
mindersec · Oct 10, 2024 · 2622be5 · 2622be5
1 parent eb15217
commit 2622be5
Show file tree

Hide file tree

Showing 19 changed files with 302 additions and 463 deletions.
diff --git a/database/mock/store.go b/database/mock/store.go
diff --git a/internal/controlplane/handlers_githubwebhooks.go b/internal/controlplane/handlers_githubwebhooks.go
@@ -799,7 +799,7 @@ func (s *Server) processRelevantRepositoryEvent(
 		WithProjectID(repoEntity.Entity.ProjectID).
 		WithProviderID(repoEntity.Entity.ProviderID).
 		WithRepository(pbRepo).
-		WithRepositoryID(repoEntity.Entity.ID)
+		WithID(repoEntity.Entity.ID)
 
 	return &processingResult{
 		topic:   events.TopicQueueEntityEvaluate,

diff --git a/internal/controlplane/handlers_githubwebhooks_test.go b/internal/controlplane/handlers_githubwebhooks_test.go
@@ -1404,7 +1404,7 @@ func (s *UnitTestSuite) TestHandleGitHubWebHook() {
 				require.Equal(t, "https://api.github.com/", received.Metadata["source"])
 				require.Equal(t, providerID.String(), received.Metadata["provider_id"])
 				require.Equal(t, projectID.String(), received.Metadata[entities.ProjectIDEventKey])
-				require.Equal(t, repositoryID.String(), received.Metadata["repository_id"])
+				require.Equal(t, repositoryID.String(), received.Metadata[entities.EntityIDEventKey])
 			},
 		},
 		{
@@ -1455,7 +1455,7 @@ func (s *UnitTestSuite) TestHandleGitHubWebHook() {
 				require.Equal(t, "https://api.github.com/", received.Metadata["source"])
 				require.Equal(t, providerID.String(), received.Metadata["provider_id"])
 				require.Equal(t, projectID.String(), received.Metadata[entities.ProjectIDEventKey])
-				require.Equal(t, repositoryID.String(), received.Metadata["repository_id"])
+				require.Equal(t, repositoryID.String(), received.Metadata[entities.EntityIDEventKey])
 			},
 		},
 		{
@@ -1602,7 +1602,7 @@ func (s *UnitTestSuite) TestHandleGitHubWebHook() {
 				require.Equal(t, "https://api.github.com/", received.Metadata["source"])
 				require.Equal(t, providerID.String(), received.Metadata["provider_id"])
 				require.Equal(t, projectID.String(), received.Metadata[entities.ProjectIDEventKey])
-				require.Equal(t, repositoryID.String(), received.Metadata["repository_id"])
+				require.Equal(t, repositoryID.String(), received.Metadata[entities.EntityIDEventKey])
 			},
 		},
 		{
@@ -1653,7 +1653,7 @@ func (s *UnitTestSuite) TestHandleGitHubWebHook() {
 				require.Equal(t, "https://api.github.com/", received.Metadata["source"])
 				require.Equal(t, providerID.String(), received.Metadata["provider_id"])
 				require.Equal(t, projectID.String(), received.Metadata[entities.ProjectIDEventKey])
-				require.Equal(t, repositoryID.String(), received.Metadata["repository_id"])
+				require.Equal(t, repositoryID.String(), received.Metadata[entities.EntityIDEventKey])
 			},
 		},
 		{
@@ -1704,7 +1704,7 @@ func (s *UnitTestSuite) TestHandleGitHubWebHook() {
 				require.Equal(t, "https://api.github.com/", received.Metadata["source"])
 				require.Equal(t, providerID.String(), received.Metadata["provider_id"])
 				require.Equal(t, projectID.String(), received.Metadata[entities.ProjectIDEventKey])
-				require.Equal(t, repositoryID.String(), received.Metadata["repository_id"])
+				require.Equal(t, repositoryID.String(), received.Metadata[entities.EntityIDEventKey])
 			},
 		},
 		{
@@ -1755,7 +1755,7 @@ func (s *UnitTestSuite) TestHandleGitHubWebHook() {
 				require.Equal(t, "https://api.github.com/", received.Metadata["source"])
 				require.Equal(t, providerID.String(), received.Metadata["provider_id"])
 				require.Equal(t, projectID.String(), received.Metadata[entities.ProjectIDEventKey])
-				require.Equal(t, repositoryID.String(), received.Metadata["repository_id"])
+				require.Equal(t, repositoryID.String(), received.Metadata[entities.EntityIDEventKey])
 			},
 		},
 		{
@@ -1854,7 +1854,7 @@ func (s *UnitTestSuite) TestHandleGitHubWebHook() {
 				require.Equal(t, "https://api.github.com/", received.Metadata["source"])
 				require.Equal(t, providerID.String(), received.Metadata["provider_id"])
 				require.Equal(t, projectID.String(), received.Metadata[entities.ProjectIDEventKey])
-				require.Equal(t, repositoryID.String(), received.Metadata["repository_id"])
+				require.Equal(t, repositoryID.String(), received.Metadata[entities.EntityIDEventKey])
 			},
 		},
 		{
@@ -1944,7 +1944,7 @@ func (s *UnitTestSuite) TestHandleGitHubWebHook() {
 				require.Equal(t, "https://api.github.com/", received.Metadata["source"])
 				require.Equal(t, providerID.String(), received.Metadata["provider_id"])
 				require.Equal(t, projectID.String(), received.Metadata[entities.ProjectIDEventKey])
-				require.Equal(t, repositoryID.String(), received.Metadata["repository_id"])
+				require.Equal(t, repositoryID.String(), received.Metadata[entities.EntityIDEventKey])
 			},
 		},
 

diff --git a/internal/db/store.go b/internal/db/store.go
@@ -33,7 +33,7 @@ type GetTypedEntitiesOptions struct {
 type ExtendQuerier interface {
 	Querier
 	GetRuleEvaluationByProfileIdAndRuleType(ctx context.Context, profileID uuid.UUID,
-		ruleName sql.NullString, entityID uuid.NullUUID, ruleTypeName sql.NullString) (*ListRuleEvaluationsByProfileIdRow, error)
+		ruleName sql.NullString, entityID uuid.UUID, ruleTypeName sql.NullString) (*ListRuleEvaluationsByProfileIdRow, error)
 	UpsertPropertyValueV1(ctx context.Context, params UpsertPropertyValueV1Params) (Property, error)
 	GetPropertyValueV1(ctx context.Context, entityID uuid.UUID, key string) (PropertyValueV1, error)
 	GetAllPropertyValuesV1(ctx context.Context, entityID uuid.UUID) ([]PropertyValueV1, error)
@@ -119,12 +119,15 @@ func (q *Queries) GetRuleEvaluationByProfileIdAndRuleType(
 	ctx context.Context,
 	profileID uuid.UUID,
 	ruleName sql.NullString,
-	entityID uuid.NullUUID,
+	entityID uuid.UUID,
 	ruleTypeName sql.NullString,
 ) (*ListRuleEvaluationsByProfileIdRow, error) {
 	params := ListRuleEvaluationsByProfileIdParams{
-		ProfileID:    profileID,
-		EntityID:     entityID,
+		ProfileID: profileID,
+		EntityID: uuid.NullUUID{
+			UUID:  entityID,
+			Valid: true,
+		},
 		RuleName:     ruleName,
 		RuleTypeName: ruleTypeName,
 	}

diff --git a/internal/eea/eea.go b/internal/eea/eea.go
@@ -299,7 +299,7 @@ func (e *EEA) buildRepositoryInfoWrapper(
 
 	return entities.NewEntityInfoWrapper().
 		WithRepository(r).
-		WithRepositoryID(repoID).
+		WithID(repoID).
 		WithProjectID(projID).
 		WithProviderID(ent.Entity.ProviderID), nil
 }
@@ -331,7 +331,7 @@ func (e *EEA) buildArtifactInfoWrapper(
 	eiw := entities.NewEntityInfoWrapper().
 		WithProjectID(projID).
 		WithArtifact(a).
-		WithArtifactID(artID).
+		WithID(artID).
 		WithProviderID(ent.Entity.ProviderID)
 	return eiw, nil
 }
@@ -350,8 +350,6 @@ func (e *EEA) buildPullRequestInfoWrapper(
 		return nil, fmt.Errorf("entity %s does not belong to project %s", prID, projID)
 	}
 
-	repoID := ent.Entity.OriginatedFrom
-
 	rawPR, err := e.entityFetcher.EntityWithPropertiesAsProto(ctx, ent, e.provMan)
 	if err != nil {
 		return nil, fmt.Errorf("error converting entity to protobuf: %w", err)
@@ -363,9 +361,8 @@ func (e *EEA) buildPullRequestInfoWrapper(
 	}
 
 	return entities.NewEntityInfoWrapper().
-		WithRepositoryID(repoID).
 		WithProjectID(projID).
 		WithPullRequest(pr).
-		WithPullRequestID(prID).
+		WithID(prID).
 		WithProviderID(ent.Entity.ProviderID), nil
 }
diff --git a/internal/eea/eea_test.go b/internal/eea/eea_test.go
@@ -109,9 +109,11 @@ func TestAggregator(t *testing.T) {
 
 	inf := entities.NewEntityInfoWrapper().
 		WithRepository(&minderv1.Repository{}).
-		WithRepositoryID(repoID).
+		WithID(repoID).
 		WithProjectID(projectID).
 		WithProviderID(providerID)
+	msg, err := inf.BuildMessage()
+	require.NoError(t, err, "expected no error when building message")
 
 	<-evt.Running()
 
@@ -121,9 +123,7 @@ func TestAggregator(t *testing.T) {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			msg, err := inf.BuildMessage()
-			require.NoError(t, err, "expected no error when building message")
-			err = evt.Publish(rateLimitedMessageTopic, msg.Copy())
+			err := evt.Publish(rateLimitedMessageTopic, msg.Copy())
 			require.NoError(t, err, "expected no error when publishing message")
 		}()
 	}