fix: warehouse skip processing destinations

achettyiitr · achettyiitr · commit 376e332b9221 · 2025-04-17T18:26:59.000+05:30
diff --git a/warehouse/internal/repo/upload.go b/warehouse/internal/repo/upload.go
@@ -78,6 +78,7 @@ type scanFn func(dest ...any) error
 type ProcessOptions struct {
 	SkipIdentifiers                   []string
 	SkipWorkspaces                    []string
+	SkipDestinations                  []string
 	AllowMultipleSourcesForJobsPickup bool
 }
 
@@ -283,12 +284,12 @@ func (u *Uploads) GetToProcess(ctx context.Context, destType string, limit int,
 	partitionIdentifierSQL := `destination_id, namespace`
 
 	if len(opts.SkipIdentifiers) > 0 {
-		skipIdentifiersSQL = `AND ((destination_id || '_' || namespace)) != ALL($5)`
+		skipIdentifiersSQL = `AND ((destination_id || '_' || namespace)) != ALL($6)`
 	}
 
 	if opts.AllowMultipleSourcesForJobsPickup {
 		if len(opts.SkipIdentifiers) > 0 {
-			skipIdentifiersSQL = `AND ((source_id || '_' || destination_id || '_' || namespace)) != ALL($5)`
+			skipIdentifiersSQL = `AND ((source_id || '_' || destination_id || '_' || namespace)) != ALL($6)`
 		}
 		partitionIdentifierSQL = fmt.Sprintf(`%s, %s`, "source_id", partitionIdentifierSQL)
 	}
@@ -308,7 +309,8 @@ func (u *Uploads) GetToProcess(ctx context.Context, destType string, limit int,
 					t.status != $2 AND
 					t.status != $3 %s AND
 					COALESCE(metadata->>'nextRetryTime', NOW()::text)::timestamptz <= NOW() AND
-          			workspace_id <> ALL ($4)
+          			workspace_id <> ALL ($4) AND
+					destination_id <> ALL ($5)
 			) grouped_uploads
 			WHERE
 				grouped_uploads.row_number = 1
@@ -330,12 +332,16 @@ func (u *Uploads) GetToProcess(ctx context.Context, destType string, limit int,
 	if opts.SkipWorkspaces == nil {
 		opts.SkipWorkspaces = []string{}
 	}
+	if opts.SkipDestinations == nil {
+		opts.SkipDestinations = []string{}
+	}
 
 	args := []any{
 		destType,
 		model.ExportedData,
 		model.Aborted,
 		pq.Array(opts.SkipWorkspaces),
+		pq.Array(opts.SkipDestinations),
 	}
 
 	if len(opts.SkipIdentifiers) > 0 {
@@ -388,22 +394,27 @@ func (u *Uploads) UploadJobsStats(ctx context.Context, destType string, opts Pro
 			status != $3 AND
 			status != $4  AND
 			COALESCE((metadata->>'nextRetryTime')::TIMESTAMPTZ, $1::TIMESTAMPTZ) <= $1::TIMESTAMPTZ AND
-			workspace_id <> ALL ($5)`
+			workspace_id <> ALL ($5) AND
+			destination_id <> ALL ($6)`
 
 	if opts.SkipWorkspaces == nil {
 		opts.SkipWorkspaces = []string{}
 	}
+	if opts.SkipDestinations == nil {
+		opts.SkipDestinations = []string{}
+	}
 
 	args := []any{
 		u.now(),
 		destType,
 		model.ExportedData,
 		model.Aborted,
 		pq.Array(opts.SkipWorkspaces),
+		pq.Array(opts.SkipDestinations),
 	}
 
 	if len(opts.SkipIdentifiers) > 0 {
-		query += `AND ((destination_id || '_' || namespace)) != ALL($6)`
+		query += `AND ((destination_id || '_' || namespace)) != ALL($7)`
 		args = append(args, pq.Array(opts.SkipIdentifiers))
 	}
 
diff --git a/warehouse/internal/repo/upload_test.go b/warehouse/internal/repo/upload_test.go
@@ -449,6 +449,40 @@ func TestUploads_GetToProcess(t *testing.T) {
 		require.NoError(t, err)
 		require.Len(t, toProcess, 0)
 	})
+	t.Run("skip destinations", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			db         = setupDB(t)
+			repoUpload = repo.NewUploads(db)
+			priority   = 100
+
+			uploads []model.Upload
+		)
+
+		uploads = append(uploads,
+			prepareUpload(db, sourceID, model.Waiting, priority,
+				time.Date(2021, 1, 1, 0, 0, 0, 0, time.UTC),
+				time.Date(2021, 1, 1, 1, 0, 0, 0, time.UTC),
+			),
+			prepareUpload(db, sourceID, model.ExportingDataFailed, priority,
+				time.Date(2021, 1, 1, 0, 0, 0, 0, time.UTC),
+				time.Date(2021, 1, 1, 1, 0, 0, 0, time.UTC),
+			),
+		)
+		require.Len(t, uploads, 2)
+
+		toProcess, err := repoUpload.GetToProcess(ctx, destType, 10, repo.ProcessOptions{})
+		require.NoError(t, err)
+		require.Len(t, toProcess, 1)
+		require.Equal(t, uploads[0].ID, toProcess[0].ID)
+
+		toProcess, err = repoUpload.GetToProcess(ctx, destType, 10, repo.ProcessOptions{
+			SkipDestinations: []string{destID},
+		})
+		require.NoError(t, err)
+		require.Len(t, toProcess, 0)
+	})
 
 	t.Run("ordering by priority", func(t *testing.T) {
 		t.Parallel()
@@ -726,6 +760,22 @@ func TestUploads_Processing(t *testing.T) {
 		require.Equal(t, []model.Upload{uploads[4]}, s)
 	})
 
+	t.Run("skip destinations", func(t *testing.T) {
+		t.Parallel()
+
+		s, err := repoUpload.GetToProcess(ctx, destType, 10, repo.ProcessOptions{
+			SkipDestinations: []string{"destination_id", "destination_id_4"},
+		})
+		require.NoError(t, err)
+		require.Empty(t, s)
+
+		s, err = repoUpload.GetToProcess(ctx, destType, 10, repo.ProcessOptions{
+			SkipDestinations: []string{"destination_id"},
+		})
+		require.NoError(t, err)
+		require.Equal(t, []model.Upload{uploads[4]}, s)
+	})
+
 	t.Run("multiple sources", func(t *testing.T) {
 		t.Parallel()
 
diff --git a/warehouse/router/router.go b/warehouse/router/router.go
@@ -106,6 +106,7 @@ type Router struct {
 		warehouseSyncFreqIgnore           config.ValueLoader[bool]
 		cronTrackerRetries                config.ValueLoader[int64]
 		uploadBufferTimeInMin             config.ValueLoader[time.Duration]
+		skipDestinationIDs                config.ValueLoader[[]string]
 	}
 
 	stats struct {
@@ -410,6 +411,7 @@ func (r *Router) uploadsToProcess(ctx context.Context, availableWorkers int, ski
 	uploads, err := r.uploadRepo.GetToProcess(ctx, r.destType, availableWorkers, repo.ProcessOptions{
 		SkipIdentifiers:                   skipIdentifiers,
 		SkipWorkspaces:                    r.tenantManager.DegradedWorkspaces(),
+		SkipDestinations:                  r.config.skipDestinationIDs.Load(),
 		AllowMultipleSourcesForJobsPickup: r.config.allowMultipleSourcesForJobsPickup,
 	})
 	if err != nil {
@@ -468,8 +470,9 @@ func (r *Router) uploadsToProcess(ctx context.Context, availableWorkers int, ski
 	}
 
 	jobsStats, err := r.uploadRepo.UploadJobsStats(ctx, r.destType, repo.ProcessOptions{
-		SkipIdentifiers: skipIdentifiers,
-		SkipWorkspaces:  r.tenantManager.DegradedWorkspaces(),
+		SkipIdentifiers:  skipIdentifiers,
+		SkipWorkspaces:   r.tenantManager.DegradedWorkspaces(),
+		SkipDestinations: r.config.skipDestinationIDs.Load(),
 	})
 	if err != nil {
 		return nil, fmt.Errorf("processing stats: %w", err)
@@ -713,6 +716,7 @@ func (r *Router) loadReloadableConfig(whName string) {
 	r.config.warehouseSyncFreqIgnore = r.conf.GetReloadableBoolVar(false, "Warehouse.warehouseSyncFreqIgnore")
 	r.config.cronTrackerRetries = r.conf.GetReloadableInt64Var(5, 1, "Warehouse.cronTrackerRetries")
 	r.config.uploadBufferTimeInMin = r.conf.GetReloadableDurationVar(180, time.Minute, "Warehouse.uploadBufferTimeInMin")
+	r.config.skipDestinationIDs = r.conf.GetReloadableStringSliceVar(nil, fmt.Sprintf(`Warehouse.%v.skipDestinationIDs`, whName))
 }
 
 func (r *Router) loadStats() {
diff --git a/warehouse/router/router_test.go b/warehouse/router/router_test.go
@@ -583,6 +583,7 @@ func TestRouter(t *testing.T) {
 		r.statsFactory = stats.NOP
 		r.conf = config.New()
 		r.config.allowMultipleSourcesForJobsPickup = true
+		r.config.skipDestinationIDs = config.SingleValueLoader([]string{})
 		r.config.stagingFilesBatchSize = config.SingleValueLoader(100)
 		r.config.warehouseSyncFreqIgnore = config.SingleValueLoader(true)
 		r.destType = destinationType
@@ -718,6 +719,7 @@ func TestRouter(t *testing.T) {
 		r.statsFactory = stats.NOP
 		r.conf = config.New()
 		r.config.allowMultipleSourcesForJobsPickup = true
+		r.config.skipDestinationIDs = config.SingleValueLoader([]string{})
 		r.config.stagingFilesBatchSize = config.SingleValueLoader(100)
 		r.config.warehouseSyncFreqIgnore = config.SingleValueLoader(true)
 		r.config.noOfWorkers = config.SingleValueLoader(10)
@@ -869,6 +871,7 @@ func TestRouter(t *testing.T) {
 			r.statsFactory = stats.NOP
 			r.conf = config.New()
 			r.config.allowMultipleSourcesForJobsPickup = true
+			r.config.skipDestinationIDs = config.SingleValueLoader([]string{})
 			r.config.stagingFilesBatchSize = config.SingleValueLoader(100)
 			r.config.warehouseSyncFreqIgnore = config.SingleValueLoader(true)
 			r.config.noOfWorkers = config.SingleValueLoader(0)
