feat: Add handling of context cancellation

Author: murarustefaan

scheduler/resolvers/resolvers.go

@@ -9,9 +9,7 @@ import (
 	"github.com/cloudquery/plugin-sdk/v4/caser"
 	"github.com/cloudquery/plugin-sdk/v4/scheduler/metrics"
 	"github.com/cloudquery/plugin-sdk/v4/schema"
-	"github.com/getsentry/sentry-go"
 	"github.com/rs/zerolog"
-	"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) {
@@ -82,6 +80,11 @@ func ResolveResourcesChunk(ctx context.Context, logger zerolog.Logger, m *metric
 		filtered := make([]*schema.Resource, 0, len(resources))
 		for _, resource := range resources {
 			if err := table.PreResourceResolver(ctx, client, resource); err != nil {
+				if ctx.Err() != nil {
+					tableLogger.Error().Err(err).Msg("pre resource resolver failed, context cancelled")
+					m.AddErrors(ctx, 1, selector)
+					return nil
+				}
 				tableLogger.Error().Err(err).Msg("pre resource resolver failed")
 				m.AddErrors(ctx, 1, selector)
 				continue

scheduler/resolvers/resolvers_test.go

@@ -94,3 +94,49 @@ func TestResolveResourcesChunk_PreResourceResolverPartialFailure(t *testing.T) {
 		})
 	}
 }
+
+// TestResolveResourcesChunk_PreResourceResolverContextCancelled verifies that
+// once the context is cancelled, the chunk is dropped immediately with a single
+// error instead of emitting one error per remaining resource.
+func TestResolveResourcesChunk_PreResourceResolverContextCancelled(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	calls := 0
+	table := &schema.Table{
+		Name: "test_table",
+		PreResourceResolver: func(_ context.Context, _ schema.ClientMeta, _ *schema.Resource) error {
+			calls++
+			if calls == 2 {
+				cancel()
+				return errors.New("pre resource resolver boom")
+			}
+			return nil
+		},
+		Columns: []schema.Column{
+			{
+				Name: "test_column",
+				Type: arrow.PrimitiveTypes.Int64,
+				Resolver: func(_ context.Context, _ schema.ClientMeta, resource *schema.Resource, c schema.Column) error {
+					return resource.Set(c.Name, int64(resource.Item.(int)))
+				},
+			},
+		},
+	}
+
+	client := testClient{}
+	m := metrics.NewMetrics()
+	m.InitWithClients(table, []schema.ClientMeta{client})
+
+	chunk := []any{0, 1, 2, 3, 4}
+	logger := zerolog.New(zerolog.NewTestWriter(t))
+
+	resources := ResolveResourcesChunk(ctx, logger, m, table, client, nil, chunk, caser.New())
+
+	require.Empty(t, resources, "cancelled chunk should not return resources")
+	require.Equal(t, 2, calls, "resolver should not be called for resources after cancellation")
+
+	selector := m.NewSelector(client.ID(), table.Name)
+	require.Equal(t, uint64(1), m.GetErrors(selector), "cancellation should be counted as a single error, not one per remaining resource")
+	require.Equal(t, uint64(0), m.GetResources(selector), "no resources should be counted for a cancelled chunk")
+}