feat: Expose error classifier

scheduler/queue/scheduler.go

@@ -33,6 +33,7 @@ type Scheduler struct {
 	metrics           *metrics.Metrics
 	invocationID      string
 	seed              int64
+	errorClassifier   schema.ErrorClassifier
 }
 
 type Option func(*Scheduler)
@@ -61,6 +62,12 @@ func WithInvocationID(invocationID string) Option {
 	}
 }
 
+func WithErrorClassifier(classifier schema.ErrorClassifier) Option {
+	return func(d *Scheduler) {
+		d.errorClassifier = classifier
+	}
+}
+
 func NewShuffleQueueScheduler(logger zerolog.Logger, m *metrics.Metrics, seed int64, opts ...Option) *Scheduler {
 	scheduler := &Scheduler{
 		logger:       logger,
@@ -104,6 +111,7 @@ func (d *Scheduler) Sync(ctx context.Context, tableClients []WorkUnit, resolvedR
 				d.deterministicCQID,
 				d.metrics,
 				msgChan,
+				d.errorClassifier,
 			).work(ctx, activeWorkSignal)
 			return nil
 		})

scheduler/queue/worker.go

@@ -32,7 +32,8 @@ type worker struct {
 	deterministicCQID bool
 	metrics           *metrics.Metrics
 	// message channel for sending SyncError messages
-	msgChan chan<- message.SyncMessage
+	msgChan         chan<- message.SyncMessage
+	errorClassifier schema.ErrorClassifier
 }
 
 func (w *worker) work(ctx context.Context, activeWorkSignal *activeWorkSignal) {
@@ -57,6 +58,7 @@ func newWorker(
 	deterministicCQID bool,
 	m *metrics.Metrics,
 	msgChan chan<- message.SyncMessage,
+	errorClassifier schema.ErrorClassifier,
 ) *worker {
 	return &worker{
 		jobs:              jobs,
@@ -68,6 +70,7 @@ func newWorker(
 		invocationID:      invocationID,
 		metrics:           m,
 		msgChan:           msgChan,
+		errorClassifier:   errorClassifier,
 	}
 }
 
@@ -113,6 +116,11 @@ func (w *worker) resolveTable(ctx context.Context, table *schema.Table, client s
 			close(res)
 		}()
 		if err := table.Resolver(ctx, client, parent, res); err != nil {
+			event := schema.ErrorEvent{Table: table, Client: client, Phase: schema.ErrorPhaseTableResolver}
+			if w.errorClassifier.Suppress(ctx, err, event) {
+				logger.Debug().Err(err).Msg("table resolver finished with suppressed error")
+				return
+			}
 			logger.Error().Err(err).Msg("table resolver finished with error")
 			w.metrics.AddErrors(ctx, 1, selector)
 			// Send SyncError message
@@ -155,7 +163,7 @@ func (w *worker) resolveResource(ctx context.Context, table *schema.Table, clien
 			wg.Add(1)
 			go func() {
 				defer wg.Done()
-				resolvedResources := resolvers.ResolveResourcesChunk(ctx, w.logger, w.metrics, table, client, parent, chunks[i], w.caser)
+				resolvedResources := resolvers.ResolveResourcesChunkWithClassifier(ctx, w.logger, w.metrics, table, client, parent, chunks[i], w.caser, w.errorClassifier)
 				for _, resolvedResource := range resolvedResources {
 					if err := resolvedResource.CalculateCQID(w.deterministicCQID); err != nil {
 						w.logger.Error().Err(err).Str("table", table.Name).Str("client", client.ID()).Msg("resource resolver finished with primary key calculation error")

scheduler/resolvers/error_classifier_test.go

@@ -0,0 +1,125 @@
+package resolvers
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/apache/arrow-go/v18/arrow"
+	"github.com/cloudquery/plugin-sdk/v4/caser"
+	"github.com/cloudquery/plugin-sdk/v4/scheduler/metrics"
+	"github.com/cloudquery/plugin-sdk/v4/schema"
+	"github.com/rs/zerolog"
+	"github.com/stretchr/testify/require"
+)
+
+var errResolverBoom = errors.New("resolver boom")
+
+// TestResolveResourcesChunk_ErrorClassifier verifies that the error classifier controls
+// whether resolver errors at each phase are counted as errors, and that it receives an
+// ErrorEvent describing the phase that failed.
+func TestResolveResourcesChunk_ErrorClassifier(t *testing.T) {
+	for _, tc := range []struct {
+		name       string
+		table      func() *schema.Table
+		wantPhase  schema.ErrorPhase
+		wantColumn string
+	}{
+		{
+			name: "column resolver",
+			table: func() *schema.Table {
+				return &schema.Table{
+					Name: "test_table",
+					Columns: []schema.Column{
+						{
+							Name: "test_column",
+							Type: arrow.PrimitiveTypes.Int64,
+							Resolver: func(_ context.Context, _ schema.ClientMeta, _ *schema.Resource, _ schema.Column) error {
+								return errResolverBoom
+							},
+						},
+					},
+				}
+			},
+			wantPhase:  schema.ErrorPhaseColumnResolver,
+			wantColumn: "test_column",
+		},
+		{
+			name: "pre resource chunk resolver",
+			table: func() *schema.Table {
+				return &schema.Table{
+					Name: "test_table",
+					PreResourceChunkResolver: &schema.RowsChunkResolver{
+						ChunkSize: 10,
+						RowsResolver: func(_ context.Context, _ schema.ClientMeta, _ []*schema.Resource) error {
+							return errResolverBoom
+						},
+					},
+					Columns: []schema.Column{{Name: "test_column", Type: arrow.PrimitiveTypes.Int64}},
+				}
+			},
+			wantPhase: schema.ErrorPhasePreResourceChunkResolver,
+		},
+		{
+			name: "pre resource resolver",
+			table: func() *schema.Table {
+				return &schema.Table{
+					Name: "test_table",
+					PreResourceResolver: func(_ context.Context, _ schema.ClientMeta, _ *schema.Resource) error {
+						return errResolverBoom
+					},
+					Columns: []schema.Column{{Name: "test_column", Type: arrow.PrimitiveTypes.Int64}},
+				}
+			},
+			wantPhase: schema.ErrorPhasePreResourceResolver,
+		},
+		{
+			name: "post resource resolver",
+			table: func() *schema.Table {
+				return &schema.Table{
+					Name: "test_table",
+					PostResourceResolver: func(_ context.Context, _ schema.ClientMeta, _ *schema.Resource) error {
+						return errResolverBoom
+					},
+					Columns: []schema.Column{{Name: "test_column", Type: arrow.PrimitiveTypes.Int64}},
+				}
+			},
+			wantPhase: schema.ErrorPhasePostResourceResolver,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			client := testClient{}
+			chunk := []any{0}
+			logger := zerolog.New(zerolog.NewTestWriter(t))
+
+			t.Run("raised without classifier", func(t *testing.T) {
+				table := tc.table()
+				m := metrics.NewMetrics()
+				m.InitWithClients(table, []schema.ClientMeta{client})
+				ResolveResourcesChunkWithClassifier(context.Background(), logger, m, table, client, nil, chunk, caser.New(), nil)
+				require.Equal(t, uint64(1), m.GetErrors(m.NewSelector(client.ID(), table.Name)))
+			})
+
+			t.Run("suppressed by classifier", func(t *testing.T) {
+				table := tc.table()
+				m := metrics.NewMetrics()
+				m.InitWithClients(table, []schema.ClientMeta{client})
+				var gotEvent schema.ErrorEvent
+				classifier := func(_ context.Context, _ error, event schema.ErrorEvent) bool {
+					gotEvent = event
+					return true
+				}
+				ResolveResourcesChunkWithClassifier(context.Background(), logger, m, table, client, nil, chunk, caser.New(), classifier)
+				require.Equal(t, uint64(0), m.GetErrors(m.NewSelector(client.ID(), table.Name)), "suppressed error should not be counted")
+				require.Equal(t, tc.wantPhase, gotEvent.Phase)
+				require.Equal(t, table, gotEvent.Table)
+				if tc.wantColumn != "" {
+					require.NotNil(t, gotEvent.Column)
+					require.Equal(t, tc.wantColumn, gotEvent.Column.Name)
+				} else {
+					require.Nil(t, gotEvent.Column)
+				}
+			})
+		})
+	}
+}

scheduler/resolvers/resolvers.go

@@ -14,7 +14,7 @@ import (
 	"github.com/thoas/go-funk"
 )
 
-func resolveColumn(ctx context.Context, logger zerolog.Logger, m *metrics.Metrics, selector metrics.Selector, client schema.ClientMeta, resource *schema.Resource, column schema.Column, c *caser.Caser) {
+func resolveColumn(ctx context.Context, logger zerolog.Logger, m *metrics.Metrics, selector metrics.Selector, client schema.ClientMeta, resource *schema.Resource, column schema.Column, c *caser.Caser, classifier schema.ErrorClassifier) {
 	columnStartTime := time.Now()
 	defer func() {
 		if err := recover(); err != nil {
@@ -29,25 +29,38 @@ func resolveColumn(ctx context.Context, logger zerolog.Logger, m *metrics.Metric
 		}
 	}()
 
+	handleErr := func(err error) {
+		event := schema.ErrorEvent{Table: resource.Table, Client: client, Phase: schema.ErrorPhaseColumnResolver, Column: &column}
+		if classifier.Suppress(ctx, err, event) {
+			logger.Debug().Str("column", column.Name).Err(err).Msg("column resolver finished with suppressed error")
+			return
+		}
+		logger.Error().Err(err).Msg("column resolver finished with error")
+		m.AddErrors(ctx, 1, selector)
+	}
+
 	if column.Resolver != nil {
 		if err := column.Resolver(ctx, client, resource, column); err != nil {
-			logger.Error().Err(err).Msg("column resolver finished with error")
-			m.AddErrors(ctx, 1, selector)
+			handleErr(err)
 		}
 	} else {
 		// base use case: try to get column with CamelCase name
 		v := funk.Get(resource.GetItem(), c.ToPascal(column.Name), funk.WithAllowZero())
 		if v != nil {
-			err := resource.Set(column.Name, v)
-			if err != nil {
-				logger.Error().Err(err).Msg("column resolver finished with error")
-				m.AddErrors(ctx, 1, selector)
+			if err := resource.Set(column.Name, v); err != nil {
+				handleErr(err)
 			}
 		}
 	}
 }
 
+// Deprecated: use ResolveResourcesChunkWithClassifier. This retains the original
+// signature and resolves with a nil classifier, so every error is raised.
 func ResolveResourcesChunk(ctx context.Context, logger zerolog.Logger, m *metrics.Metrics, table *schema.Table, client schema.ClientMeta, parent *schema.Resource, chunk []any, c *caser.Caser) []*schema.Resource {
+	return ResolveResourcesChunkWithClassifier(ctx, logger, m, table, client, parent, chunk, c, nil)
+}
+
+func ResolveResourcesChunkWithClassifier(ctx context.Context, logger zerolog.Logger, m *metrics.Metrics, table *schema.Table, client schema.ClientMeta, parent *schema.Resource, chunk []any, c *caser.Caser, classifier schema.ErrorClassifier) []*schema.Resource {
 	ctx, cancel := context.WithTimeout(ctx, 10*time.Minute)
 	defer cancel()
 
@@ -72,8 +85,13 @@ func ResolveResourcesChunk(ctx context.Context, logger zerolog.Logger, m *metric
 
 	if table.PreResourceChunkResolver != nil {
 		if err := table.PreResourceChunkResolver.RowsResolver(ctx, client, resources); err != nil {
-			tableLogger.Error().Stack().Err(err).Msg("pre resource chunk resolver finished with error")
-			m.AddErrors(ctx, 1, selector)
+			event := schema.ErrorEvent{Table: table, Client: client, Phase: schema.ErrorPhasePreResourceChunkResolver}
+			if classifier.Suppress(ctx, err, event) {
+				tableLogger.Debug().Err(err).Msg("pre resource chunk resolver finished with suppressed error")
+			} else {
+				tableLogger.Error().Stack().Err(err).Msg("pre resource chunk resolver finished with error")
+				m.AddErrors(ctx, 1, selector)
+			}
 			return nil
 		}
 	}
@@ -82,30 +100,45 @@ func ResolveResourcesChunk(ctx context.Context, logger zerolog.Logger, m *metric
 		filtered := resources[:0]
 		for _, resource := range resources {
 			if err := table.PreResourceResolver(ctx, client, resource); err != nil {
-				if ctx.Err() != nil {
+				event := schema.ErrorEvent{Table: table, Client: client, Phase: schema.ErrorPhasePreResourceResolver}
+				suppress := classifier.Suppress(ctx, err, event)
+				switch {
+				case suppress && ctx.Err() != nil:
+					tableLogger.Debug().Err(err).Msg("pre resource resolver failed, context cancelled (suppressed)")
+					return nil
+				case suppress:
+					tableLogger.Debug().Err(err).Msg("pre resource resolver failed (suppressed)")
+					continue
+				case ctx.Err() != nil:
 					tableLogger.Error().Err(err).Msg("pre resource resolver failed, context cancelled")
 					m.AddErrors(ctx, 1, selector)
 					return nil
+				default:
+					tableLogger.Error().Err(err).Msg("pre resource resolver failed")
+					m.AddErrors(ctx, 1, selector)
+					continue
 				}
-				tableLogger.Error().Err(err).Msg("pre resource resolver failed")
-				m.AddErrors(ctx, 1, selector)
-				continue
 			}
 			filtered = append(filtered, resource)
 		}
 		resources = filtered
 	}
 	for _, resource := range resources {
 		for _, column := range table.Columns {
-			resolveColumn(ctx, tableLogger, m, selector, client, resource, column, c)
+			resolveColumn(ctx, tableLogger, m, selector, client, resource, column, c, classifier)
 		}
 	}
 
 	if table.PostResourceResolver != nil {
 		for _, resource := range resources {
 			if err := table.PostResourceResolver(ctx, client, resource); err != nil {
-				tableLogger.Error().Stack().Err(err).Msg("post resource resolver finished with error")
-				m.AddErrors(ctx, 1, selector)
+				event := schema.ErrorEvent{Table: table, Client: client, Phase: schema.ErrorPhasePostResourceResolver}
+				if classifier.Suppress(ctx, err, event) {
+					tableLogger.Debug().Err(err).Msg("post resource resolver finished with suppressed error")
+				} else {
+					tableLogger.Error().Stack().Err(err).Msg("post resource resolver finished with error")
+					m.AddErrors(ctx, 1, selector)
+				}
 			}
 		}
 	}

scheduler/scheduler.go

@@ -37,6 +37,21 @@ const (
 	StrategyShuffleQueue
 )
 
+// Re-exported from the schema package for use with WithErrorClassifier.
+type (
+	ErrorClassifier = schema.ErrorClassifier
+	ErrorEvent      = schema.ErrorEvent
+	ErrorPhase      = schema.ErrorPhase
+)
+
+const (
+	ErrorPhaseTableResolver            = schema.ErrorPhaseTableResolver
+	ErrorPhasePreResourceChunkResolver = schema.ErrorPhasePreResourceChunkResolver
+	ErrorPhasePreResourceResolver      = schema.ErrorPhasePreResourceResolver
+	ErrorPhaseColumnResolver           = schema.ErrorPhaseColumnResolver
+	ErrorPhasePostResourceResolver     = schema.ErrorPhasePostResourceResolver
+)
+
 type Option func(*Scheduler)
 
 func WithLogger(logger zerolog.Logger) Option {
@@ -45,6 +60,14 @@ func WithLogger(logger zerolog.Logger) Option {
 	}
 }
 
+// WithErrorClassifier sets the classifier consulted for every resolver error. See
+// schema.ErrorClassifier.
+func WithErrorClassifier(classifier ErrorClassifier) Option {
+	return func(s *Scheduler) {
+		s.errorClassifier = classifier
+	}
+}
+
 func WithConcurrency(concurrency int) Option {
 	return func(s *Scheduler) {
 		s.concurrency = concurrency
@@ -122,6 +145,8 @@ type Scheduler struct {
 	batchSettings *BatchSettings
 
 	invocationID string
+
+	errorClassifier schema.ErrorClassifier
 }
 
 type shard struct {

scheduler/scheduler_dfs.go

@@ -128,6 +128,11 @@ func (s *syncClient) resolveTableDfs(ctx context.Context, table *schema.Table, c
 			close(res)
 		}()
 		if err := table.Resolver(ctx, client, parent, res); err != nil {
+			event := schema.ErrorEvent{Table: table, Client: client, Phase: schema.ErrorPhaseTableResolver}
+			if s.scheduler.errorClassifier.Suppress(ctx, err, event) {
+				logger.Debug().Err(err).Msg("table resolver finished with suppressed error")
+				return
+			}
 			logger.Error().Err(err).Msg("table resolver finished with error")
 			s.metrics.AddErrors(ctx, 1, selector)
 			// Send SyncError message
@@ -192,7 +197,7 @@ func (s *syncClient) resolveResourcesDfs(ctx context.Context, table *schema.Tabl
 				defer resourceSem.Release(1)
 				defer s.scheduler.resourceSem.Release(1)
 				defer wg.Done()
-				resolvedResources := resolvers.ResolveResourcesChunk(ctx, s.logger, s.metrics, table, client, parent, chunks[i], s.scheduler.caser)
+				resolvedResources := resolvers.ResolveResourcesChunkWithClassifier(ctx, s.logger, s.metrics, table, client, parent, chunks[i], s.scheduler.caser, s.scheduler.errorClassifier)
 				if len(resolvedResources) == 0 {
 					return
 				}