chore: Make experimental, add tests

Author: erezrokah

scheduler/scheduler_test.go

@@ -3,15 +3,18 @@ package scheduler
 import (
 	"context"
 	"fmt"
+	"strconv"
 	"testing"
 	"time"
 
 	"github.com/apache/arrow-go/v18/arrow"
 	"github.com/apache/arrow-go/v18/arrow/array"
+	"github.com/apache/arrow-go/v18/arrow/memory"
 	"github.com/cloudquery/plugin-sdk/v4/message"
 	"github.com/cloudquery/plugin-sdk/v4/scalar"
 	"github.com/cloudquery/plugin-sdk/v4/schema"
 	"github.com/rs/zerolog"
+	"github.com/samber/lo"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -184,9 +187,82 @@ func testTableRelationSuccess() *schema.Table {
 	}
 }
 
+const chunkSize = 200
+
+func testTableSuccessWithRowsChunkResolverSendSingleItemToResChan() *schema.Table {
+	return &schema.Table{
+		Name: "test_table_success_with_rows_chunk_resolver",
+		Resolver: func(ctx context.Context, meta schema.ClientMeta, parent *schema.Resource, res chan<- any) error {
+			for i := range chunkSize {
+				res <- i
+			}
+			return nil
+		},
+		PreResourceChunkResolver: &schema.RowsChunkResolver{
+			ChunkSize: chunkSize,
+			RowsResolver: func(ctx context.Context, meta schema.ClientMeta, resourcesChunk []*schema.Resource) error {
+				for _, resource := range resourcesChunk {
+					resource.Set("test_column", strconv.Itoa(resource.Item.(int)))
+				}
+				return nil
+			},
+		},
+		Columns: []schema.Column{
+			{
+				Name: "test_column",
+				Type: arrow.BinaryTypes.String,
+			},
+		},
+	}
+}
+
+func testTableSuccessWithRowsChunkResolverSendSliceToResChan() *schema.Table {
+	return &schema.Table{
+		Name: "test_table_success_with_rows_chunk_resolver",
+		Resolver: func(ctx context.Context, meta schema.ClientMeta, parent *schema.Resource, res chan<- any) error {
+			data := make([]int, chunkSize)
+			for i := range chunkSize {
+				data[i] = i
+			}
+			res <- data
+			return nil
+		},
+		PreResourceChunkResolver: &schema.RowsChunkResolver{
+			ChunkSize: chunkSize,
+			RowsResolver: func(ctx context.Context, meta schema.ClientMeta, resourcesChunk []*schema.Resource) error {
+				for _, resource := range resourcesChunk {
+					resource.Set("test_column", strconv.Itoa(resource.Item.(int)))
+				}
+				return nil
+			},
+		},
+		Columns: []schema.Column{
+			{
+				Name: "test_column",
+				Type: arrow.BinaryTypes.String,
+			},
+		},
+	}
+}
+
+func expectedChunkedResolverData(schema *arrow.Schema) []arrow.Record {
+	const rowsPerRecord = 50
+	data := make([]arrow.Record, chunkSize/rowsPerRecord)
+	for i := range data {
+		builder := array.NewRecordBuilder(memory.DefaultAllocator, schema)
+		for j := range rowsPerRecord {
+			builder.Field(0).(*array.StringBuilder).Append(strconv.Itoa(i*rowsPerRecord + j))
+		}
+		record := builder.NewRecord()
+		data[i] = record
+	}
+	return data
+}
+
 type syncTestCase struct {
 	table             *schema.Table
 	data              []scalar.Vector
+	dataAsRecords     []arrow.Record
 	deterministicCQID bool
 	err               error
 }
@@ -298,6 +374,14 @@ var syncTestCases = []syncTestCase{
 			},
 		},
 	},
+	{
+		table:         testTableSuccessWithRowsChunkResolverSendSingleItemToResChan(),
+		dataAsRecords: expectedChunkedResolverData(testTableSuccessWithRowsChunkResolverSendSingleItemToResChan().ToArrowSchema()),
+	},
+	{
+		table:         testTableSuccessWithRowsChunkResolverSendSliceToResChan(),
+		dataAsRecords: expectedChunkedResolverData(testTableSuccessWithRowsChunkResolverSendSliceToResChan().ToArrowSchema()),
+	},
 }
 
 type optionsTestCase struct {
@@ -348,36 +432,24 @@ func testSyncTable(t *testing.T, tc syncTestCase, strategy Strategy, determinist
 		WithStrategy(strategy),
 	}, extra...)
 	sc := NewScheduler(opts...)
-	msgs := make(chan message.SyncMessage, 10)
+	// We use a buffer channel to avoid the complexity of testing using a Go routine, we just need the buffer to be big enough to contain all sync messages
+	msgs := make(chan message.SyncMessage, 500)
 	err := sc.Sync(ctx, &c, tables, msgs, WithSyncDeterministicCQID(deterministicCQID))
 	require.ErrorIs(t, err, tc.err)
 	close(msgs)
 
-	var i int
+	dataAsRecords := tc.dataAsRecords
+	if dataAsRecords == nil {
+		dataAsRecords = lo.Map(tc.data, func(item scalar.Vector, _ int) arrow.Record {
+			return item.ToArrowRecord(tc.table.ToArrowSchema())
+		})
+	}
+
+	gotRecords := make([]arrow.Record, 0)
 	for msg := range msgs {
 		switch v := msg.(type) {
 		case *message.SyncInsert:
-			record := v.Record
-			rec := tc.data[i].ToArrowRecord(record.Schema())
-			if !array.RecordEqual(rec, record) {
-				// For records that include CqIDColumn, we can't verify equality because it is generated by the scheduler, unless deterministicCQID is true
-				onlyCqIDInequality := false
-				for col := range rec.Columns() {
-					if !deterministicCQID && rec.ColumnName(col) == schema.CqIDColumn.Name {
-						onlyCqIDInequality = true
-						continue
-					}
-					lc := rec.Column(col)
-					rc := record.Column(col)
-					if !array.Equal(lc, rc) {
-						onlyCqIDInequality = false
-					}
-				}
-				if !onlyCqIDInequality {
-					t.Fatalf("expected at i=%d: %v. got %v", i, tc.data[i], record)
-				}
-			}
-			i++
+			gotRecords = append(gotRecords, v.Record)
 		case *message.SyncMigrateTable:
 			migratedTable := v.Table
 
@@ -402,8 +474,47 @@ func testSyncTable(t *testing.T, tc syncTestCase, strategy Strategy, determinist
 			t.Fatalf("expected insert message. got %T", msg)
 		}
 	}
-	if len(tc.data) != i {
-		t.Fatalf("expected %d resources. got %d", len(tc.data), i)
+
+	// We do this since the SDK can batch rows into a single record, so we need to compare them as single row records
+	slicedExpectedRecords := make([]arrow.Record, 0)
+	for _, record := range dataAsRecords {
+		for j := int64(0); j < record.NumRows(); j++ {
+			slicedRecord := record.NewSlice(j, j+1)
+			slicedExpectedRecords = append(slicedExpectedRecords, slicedRecord)
+		}
+	}
+	gotSlicedRecords := make([]arrow.Record, 0)
+	for _, record := range gotRecords {
+		for j := int64(0); j < record.NumRows(); j++ {
+			slicedRecord := record.NewSlice(j, j+1)
+			gotSlicedRecords = append(gotSlicedRecords, slicedRecord)
+		}
+	}
+	if len(slicedExpectedRecords) != len(gotSlicedRecords) {
+		t.Fatalf("expected %d rows. got %d", len(slicedExpectedRecords), len(gotSlicedRecords))
+	}
+
+	for _, expectedRecord := range slicedExpectedRecords {
+		// Records can be returned in any order, so we need to find the matching record
+		_, found := lo.Find(gotSlicedRecords, func(gotRecord arrow.Record) bool {
+			if deterministicCQID {
+				return array.RecordEqual(gotRecord, expectedRecord)
+			}
+			for col := range gotRecord.Columns() {
+				// skip equality check for random CQID values
+				if gotRecord.ColumnName(col) == schema.CqIDColumn.Name {
+					continue
+				}
+				if !array.Equal(gotRecord.Column(col), expectedRecord.Column(col)) {
+					return false
+				}
+			}
+			return true
+		})
+
+		if !found {
+			t.Fatalf("expected record %v not found", expectedRecord)
+		}
 	}
 }
 

schema/table.go

@@ -25,6 +25,7 @@ type TableResolver func(ctx context.Context, meta ClientMeta, parent *Resource,
 
 type RowResolver func(ctx context.Context, meta ClientMeta, resource *Resource) error
 
+// EXPERIMENTAL: RowsChunkResolver API might change in future versions of the SDK
 type RowsChunkResolver struct {
 	ChunkSize    int
 	RowsResolver func(ctx context.Context, meta ClientMeta, resourcesChunk []*Resource) error