Create/Delete pull requests in central entities table

database/mock/fixtures/store.go

@@ -124,6 +124,9 @@ func WithSuccessfulUpsertPullRequest(
 		mockStore.EXPECT().
 			UpsertPullRequest(gomock.Any(), gomock.Any()).
 			Return(pullRequest, nil)
+		mockStore.EXPECT().
+			CreateOrEnsureEntityByID(gomock.Any(), gomock.Any()).
+			Return(db.EntityInstance{}, nil)
 	}
 }
 
@@ -132,6 +135,9 @@ func WithSuccessfulDeletePullRequest() func(*mockdb.MockStore) {
 		mockStore.EXPECT().
 			DeletePullRequest(gomock.Any(), gomock.Any()).
 			Return(nil)
+		mockStore.EXPECT().
+			DeleteEntityByName(gomock.Any(), gomock.Any()).
+			Return(nil)
 	}
 }
 
@@ -154,3 +160,20 @@ func WithSuccessfulUpsertArtifact(
 			Return(artifact, nil)
 	}
 }
+
+func WithTransaction() func(*mockdb.MockStore) {
+	return func(mockStore *mockdb.MockStore) {
+		mockStore.EXPECT().
+			BeginTransaction().
+			Return(nil, nil)
+		mockStore.EXPECT().
+			GetQuerierWithTransaction(gomock.Any()).
+			Return(mockStore)
+		mockStore.EXPECT().
+			Commit(gomock.Any()).
+			Return(nil)
+		mockStore.EXPECT().
+			Rollback(gomock.Any()).
+			Return(nil)
+	}
+}

database/query/entities.sql

@@ -23,12 +23,33 @@ INSERT INTO entity_instances (
 ) VALUES ($1, $2, $3, sqlc.arg(project_id), sqlc.arg(provider_id), sqlc.narg(originated_from))
 RETURNING *;
 
+
+-- CreateOrEnsureEntityByID adds an entry to the entity_instances table if it does not exist, or returns the existing entry.
+
+-- name: CreateOrEnsureEntityByID :one
+INSERT INTO entity_instances (
+    id,
+    entity_type,
+    name,
+    project_id,
+    provider_id,
+    originated_from
+) VALUES ($1, $2, $3, sqlc.arg(project_id), sqlc.arg(provider_id), sqlc.narg(originated_from))
+ON CONFLICT (id) DO NOTHING
+RETURNING *;
+
 -- DeleteEntity removes an entity from the entity_instances table for a project.
 
 -- name: DeleteEntity :exec
 DELETE FROM entity_instances
 WHERE id = $1 AND project_id = $2;
 
+-- DeleteEntityByName removes an entity from the entity_instances table for a project.
+
+-- name: DeleteEntityByName :exec
+DELETE FROM entity_instances
+WHERE name = sqlc.arg(name) AND project_id = $1;
+
 -- GetEntityByID retrieves an entity by its ID for a project or hierarchy of projects.
 
 -- name: GetEntityByID :one

internal/controlplane/handlers_githubwebhooks.go

@@ -1438,22 +1438,60 @@ func (s *Server) reconcilePrWithDb(
 	// published, see here
 	// https://pkg.go.dev/github.com/google/go-github/[email protected]/github#PullRequestEvent
 	case webhookActionEventOpened, webhookActionEventReopened:
-		dbPr, err := s.store.UpsertPullRequest(ctx, db.UpsertPullRequestParams{
-			RepositoryID: dbrepo.ID,
-			PrNumber:     prEvalInfo.Number,
+		var err error
+		retPr, err = db.WithTransaction(s.store, func(t db.ExtendQuerier) (*db.PullRequest, error) {
+			dbPr, err := t.UpsertPullRequest(ctx, db.UpsertPullRequestParams{
+				RepositoryID: dbrepo.ID,
+				PrNumber:     prEvalInfo.Number,
+			})
+			if err != nil {
+				return nil, err
+			}
+
+			_, err = t.CreateOrEnsureEntityByID(ctx, db.CreateOrEnsureEntityByIDParams{
+				ID:         dbPr.ID,
+				EntityType: db.EntitiesPullRequest,
+				Name:       pullRequestName(dbrepo.RepoOwner, dbrepo.RepoName, prEvalInfo.Number),
+				ProjectID:  dbrepo.ProjectID,
+				ProviderID: dbrepo.ProviderID,
+				OriginatedFrom: uuid.NullUUID{
+					UUID:  dbrepo.ID,
+					Valid: true,
+				},
+			})
+			if err != nil {
+				return nil, err
+			}
+
+			return &dbPr, nil
 		})
 		if err != nil {
 			return nil, fmt.Errorf(
 				"cannot upsert PR %d in repo %s/%s",
 				prEvalInfo.Number, dbrepo.RepoOwner, dbrepo.RepoName)
 		}
-		retPr = &dbPr
 		retErr = nil
 	case webhookActionEventClosed:
-		err := s.store.DeletePullRequest(ctx, db.DeletePullRequestParams{
-			RepositoryID: dbrepo.ID,
-			PrNumber:     prEvalInfo.Number,
+		_, err := db.WithTransaction(s.store, func(t db.ExtendQuerier) (*db.PullRequest, error) {
+			err := t.DeletePullRequest(ctx, db.DeletePullRequestParams{
+				RepositoryID: dbrepo.ID,
+				PrNumber:     prEvalInfo.Number,
+			})
+			if err != nil {
+				return nil, err
+			}
+
+			err = t.DeleteEntityByName(ctx, db.DeleteEntityByNameParams{
+				Name:      pullRequestName(dbrepo.RepoOwner, dbrepo.RepoName, prEvalInfo.Number),
+				ProjectID: dbrepo.ProjectID,
+			})
+			if err != nil {
+				return nil, err
+			}
+
+			return nil, nil
 		})
+
 		if err != nil && !errors.Is(err, sql.ErrNoRows) {
 			return nil, fmt.Errorf("cannot delete PR record %d in repo %s/%s",
 				prEvalInfo.Number, dbrepo.RepoOwner, dbrepo.RepoName)
@@ -1485,3 +1523,7 @@ func updatePullRequestInfoFromProvider(
 	prEvalInfo.RepoName = dbrepo.RepoName
 	return nil
 }
+
+func pullRequestName(owner, name string, number int64) string {
+	return fmt.Sprintf("%s/%s/%d", owner, name, number)
+}

internal/controlplane/handlers_githubwebhooks_test.go

@@ -3041,6 +3041,7 @@ func (s *UnitTestSuite) TestHandleGitHubWebHook() {
 				df.WithSuccessfulUpsertPullRequest(
 					db.PullRequest{},
 				),
+				df.WithTransaction(),
 			),
 			topic:      events.TopicQueueEntityEvaluate,
 			statusCode: http.StatusOK,
@@ -3092,6 +3093,7 @@ func (s *UnitTestSuite) TestHandleGitHubWebHook() {
 					},
 				),
 				df.WithSuccessfulDeletePullRequest(),
+				df.WithTransaction(),
 			),
 			topic:      events.TopicQueueEntityEvaluate,
 			statusCode: http.StatusOK,