fix: preserve DCA relationships after failed processing BED-10064

packages/go/analysis/ad/owns.go

@@ -18,6 +18,7 @@ package ad
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 
@@ -76,17 +77,17 @@ func PostOwnsAndWriteOwner(ctx context.Context, db graph.Database, localGroupDat
 	}
 
 	sink := post.NewFilteredRelationshipSink(ctx, "PostOwnsAndWriteOwner", db, ownsWriteOwnerTracker)
-	defer sink.Done()
+	defer sink.Abort()
 
 	if err := postOwnsEdges(ctx, db, sink, dsHeuristicsCache, anyEnforced, adminGroupIds); err != nil {
-		return sink.Stats(), err
+		return sink.Stats(), errors.Join(err, sink.Abort())
 	}
 
 	if err := postWriteOwnerEdges(ctx, db, sink, dsHeuristicsCache, anyEnforced); err != nil {
-		return sink.Stats(), err
+		return sink.Stats(), errors.Join(err, sink.Abort())
 	}
 
-	return sink.Stats(), nil
+	return sink.Stats(), sink.Done()
 }
 
 func postOwnsEdges(ctx context.Context, db graph.Database, sink *post.FilteredRelationshipSink, dsHeuristicsCache map[string]bool, anyEnforced bool, adminGroupIds cardinality.Duplex[uint64]) error {

packages/go/analysis/azure/post.go

@@ -18,6 +18,7 @@ package azure
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 	"strings"
@@ -225,15 +226,15 @@ func AppRoleAssignments(ctx context.Context, db graph.Database) (*post.AtomicPos
 		return &post.AtomicPostProcessingStats{}, err
 	} else {
 		sink := post.NewFilteredRelationshipSink(ctx, "Azure App Role Assignments Post Processing", db, appRoleAssignmentTracker)
-		defer sink.Done()
+		defer sink.Abort()
 
 		for _, tenant := range tenants {
 			if err := postAppRoleAssignmentsForTenant(ctx, db, sink, tenant); err != nil {
-				return &post.AtomicPostProcessingStats{}, err
+				return &post.AtomicPostProcessingStats{}, errors.Join(err, sink.Abort())
 			}
 		}
 
-		return sink.Stats(), nil
+		return sink.Stats(), sink.Done()
 	}
 }
 
@@ -523,15 +524,15 @@ func CreateAZAddOwnerEdge(ctx context.Context, db graph.Database) (*post.AtomicP
 		return &post.AtomicPostProcessingStats{}, err
 	} else {
 		sink := post.NewFilteredRelationshipSink(ctx, "Azure Add Owner Post Processing", db, addOwnerTracker)
-		defer sink.Done()
+		defer sink.Abort()
 
 		for _, tenant := range tenantNodes {
 			if err := postAzureAddOwner(ctx, db, sink, tenant); err != nil {
-				return &post.AtomicPostProcessingStats{}, err
+				return &post.AtomicPostProcessingStats{}, errors.Join(err, sink.Abort())
 			}
 		}
 
-		return sink.Stats(), nil
+		return sink.Stats(), sink.Done()
 	}
 }
 
@@ -837,27 +838,28 @@ func UserRoleAssignments(ctx context.Context, db graph.Database) (*post.AtomicPo
 		return &post.AtomicPostProcessingStats{}, err
 	} else {
 		sink := post.NewFilteredRelationshipSink(ctx, "Azure User Role Assignments Post Processing", db, userRoleAssignmentTracker)
-		defer sink.Done()
+		defer sink.Abort()
 
 		for _, tenant := range tenantNodes {
 			if roleAssignments, err := FetchTenantRoleAssignments(ctx, db, tenant); err != nil {
-				return &post.AtomicPostProcessingStats{}, err
+				return &post.AtomicPostProcessingStats{}, errors.Join(err, sink.Abort())
 			} else {
 				if err := postAzureResetPassword(ctx, db, sink, tenant, roleAssignments); err != nil {
-					return &post.AtomicPostProcessingStats{}, err
+					return &post.AtomicPostProcessingStats{}, errors.Join(err, sink.Abort())
 				} else {
 					postAzureGlobalAdmins(ctx, sink, roleAssignments, tenant)
 					postAzurePrivilegedRoleAdmins(ctx, sink, roleAssignments, tenant)
 					postAzurePrivilegedAuthAdmins(ctx, sink, roleAssignments, tenant)
 
 					if err := postAzureAddMembers(ctx, sink, roleAssignments); err != nil {
 						slog.Error("Azure AddMember Post-Processing Failure", attr.Error(err))
+						return &post.AtomicPostProcessingStats{}, errors.Join(err, sink.Abort())
 					}
 				}
 			}
 		}
 
-		return sink.Stats(), nil
+		return sink.Stats(), sink.Done()
 	}
 }
 

packages/go/analysis/post/sink.go

@@ -18,6 +18,8 @@ package post
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"log/slog"
 	"runtime"
 	"sync"
@@ -38,28 +40,42 @@ func newPropertiesWithFirstSeen() *graph.Properties {
 
 // FilteredRelationshipSink is an asynchronous graph relationship writer that ensures only new relationships are
 // inserted and removes unused ones. It uses a delta tracker to track changes between graphs and avoids reinserting
-// existing edges. Any edge not visited during processing is treated as obsolete and will be deleted after the
-// operation completes.
+// existing edges. Any edge not visited during successful processing is treated as obsolete and will be deleted
+// by Done. Failed processing must call Abort to preserve existing edges.
 type FilteredRelationshipSink struct {
 	operationName string
 	db            graph.Database
 	edgeTracker   *Tracker
 	jobC          chan EnsureRelationshipJob
 	stats         AtomicPostProcessingStats
 	wg            sync.WaitGroup
+	ctx           context.Context
+	cancel        context.CancelFunc
+	submitLock    sync.RWMutex
+	submitDone    chan struct{}
+	shutdownOnce  sync.Once
+	errorLock     sync.Mutex
+	workerError   error
+	doneError     error
 }
 
 // NewFilteredRelationshipSink creates a new filtered relationship sink initialized with a given database, delta tracker, and operation name.
 func NewFilteredRelationshipSink(ctx context.Context, operationName string, db graph.Database, deltaSubgraph *Tracker) *FilteredRelationshipSink {
-	newSink := &FilteredRelationshipSink{
-		db:            db,
-		edgeTracker:   deltaSubgraph,
-		operationName: operationName,
-		jobC:          make(chan EnsureRelationshipJob),
-		stats:         NewAtomicPostProcessingStats(),
-	}
+	var (
+		workerContext, cancel = context.WithCancel(ctx)
+		newSink               = &FilteredRelationshipSink{
+			db:            db,
+			edgeTracker:   deltaSubgraph,
+			operationName: operationName,
+			jobC:          make(chan EnsureRelationshipJob),
+			stats:         NewAtomicPostProcessingStats(),
+			ctx:           workerContext,
+			cancel:        cancel,
+			submitDone:    make(chan struct{}),
+		}
+	)
 
-	newSink.start(ctx)
+	newSink.start(workerContext)
 	return newSink
 }
 
@@ -82,7 +98,8 @@ func (s *FilteredRelationshipSink) insertWorker(ctx context.Context, commonProps
 				}
 
 				if err := batch.CreateRelationshipByIDs(nextJob.FromID, nextJob.ToID, nextJob.Kind, relProps); err != nil {
-					slog.Error("Create Relationship Error", slog.String("err", err.Error()))
+					s.recordError(fmt.Errorf("%s: creating relationship: %w", s.operationName, err))
+					return err
 				} else {
 					s.stats.AddRelationshipsCreated(nextJob.Kind, 1)
 
@@ -94,12 +111,27 @@ func (s *FilteredRelationshipSink) insertWorker(ctx context.Context, commonProps
 			}
 		}
 
-		return nil
+		return ctx.Err()
 	}); err != nil {
-		slog.Error("FilteredRelationshipSink Error", attr.Error(err))
+		s.recordError(fmt.Errorf("%s: inserting relationships: %w", s.operationName, err))
 	}
 }
 
+// recordError cancels the pipeline so submissions cannot block after the writer stops.
+func (s *FilteredRelationshipSink) recordError(err error) {
+	if contextError := s.ctx.Err(); contextError != nil && errors.Is(err, contextError) {
+		return
+	}
+
+	s.errorLock.Lock()
+	if s.workerError == nil {
+		s.workerError = err
+		slog.ErrorContext(s.ctx, "FilteredRelationshipSink Error", attr.Error(err))
+	}
+	s.errorLock.Unlock()
+	s.cancel()
+}
+
 // deltaFilterWorker filters out duplicate edges before they reach the insert worker. It checks whether
 // an edge has already been tracked in the delta subgraph; if not, it forwards it to the insert channel.
 func (s *FilteredRelationshipSink) deltaFilterWorker(ctx context.Context, filterC, insertC chan EnsureRelationshipJob) {
@@ -148,8 +180,7 @@ func (s *FilteredRelationshipSink) deleteMissingEdges(ctx context.Context) error
 	return nil
 }
 
-// worker is the main goroutine responsible for managing the entire lifecycle of the sink. It
-// coordinates between filtering, insertion, and deletion phases and handles shutdown gracefully.
+// worker coordinates filtering and insertion. Pruning is reserved for successful completion in Done.
 func (s *FilteredRelationshipSink) worker(ctx context.Context) {
 	defer s.wg.Done()
 
@@ -189,11 +220,6 @@ func (s *FilteredRelationshipSink) worker(ctx context.Context) {
 
 	close(insertC)
 	insertWG.Wait()
-
-	// Remove any lingering edges after the operation completes
-	if err := s.deleteMissingEdges(ctx); err != nil {
-		slog.Error("Failed deleting missing edges after ingestion sink flush", attr.Error(err))
-	}
 }
 
 // Stats returns a pointer to the atomic statistics structure tracking processed relationships.
@@ -209,11 +235,74 @@ func (s *FilteredRelationshipSink) start(ctx context.Context) {
 
 // Submit submits a new job to be processed by the sink.
 func (s *FilteredRelationshipSink) Submit(ctx context.Context, nextJob EnsureRelationshipJob) bool {
-	return channels.Submit(ctx, s.jobC, nextJob)
+	s.submitLock.RLock()
+	defer s.submitLock.RUnlock()
+
+	if ctx.Err() != nil || s.ctx.Err() != nil {
+		return false
+	}
+	select {
+	case <-s.submitDone:
+		return false
+	default:
+	}
+
+	select {
+	case <-ctx.Done():
+		return false
+	case <-s.ctx.Done():
+		return false
+	case <-s.submitDone:
+		return false
+	case s.jobC <- nextJob:
+		return true
+	}
+}
+
+// shutdown stops submissions and waits for workers. The first shutdown call determines whether to prune.
+func (s *FilteredRelationshipSink) shutdown(prune bool) {
+	s.shutdownOnce.Do(func() {
+		defer s.cancel()
+		if !prune {
+			s.cancel()
+		}
+
+		close(s.submitDone)
+		s.submitLock.Lock()
+		close(s.jobC)
+		s.submitLock.Unlock()
+		s.wg.Wait()
+
+		s.errorLock.Lock()
+		s.doneError = s.workerError
+		s.errorLock.Unlock()
+
+		if s.doneError == nil {
+			if err := s.ctx.Err(); err != nil {
+				s.doneError = err
+			} else if prune {
+				s.doneError = s.deleteMissingEdges(s.ctx)
+				if s.doneError != nil {
+					slog.ErrorContext(s.ctx, "Error deleting stale relationships", attr.Error(s.doneError))
+				}
+			}
+		}
+	})
+}
+
+// Done drains submitted jobs and removes obsolete relationships only if processing succeeded.
+// Repeated calls return the same result. Producers must complete successfully before calling Done.
+func (s *FilteredRelationshipSink) Done() error {
+	s.shutdown(true)
+	return s.doneError
 }
 
-// Done signals the end of processing and waits for all workers to complete.
-func (s *FilteredRelationshipSink) Done() {
-	close(s.jobC)
-	s.wg.Wait()
+// Abort cancels processing and waits for workers without deleting obsolete relationships.
+// It returns any writer error, excluding cancellation caused by shutdown. Repeated calls are safe.
+// Calling Done after Abort returns a cancellation or worker error.
+func (s *FilteredRelationshipSink) Abort() error {
+	s.shutdown(false)
+	s.errorLock.Lock()
+	defer s.errorLock.Unlock()
+	return s.workerError
 }