feat(testdata): Let destinations bound generated 64-bit integers

internal/memdb/memdb.go

@@ -7,6 +7,7 @@ import (
 
 	"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/plugin"
 	"github.com/cloudquery/plugin-sdk/v4/schema"
@@ -21,6 +22,7 @@ type client struct {
 	memoryDBLock  sync.RWMutex
 	errOnWrite    bool
 	blockingWrite bool
+	float64Ints   bool
 }
 
 type Option func(*client)
@@ -40,6 +42,47 @@ func WithBlockingWrite() Option {
 	}
 }
 
+// WithFloat64Ints makes writes round 64-bit integers through float64, mimicking destinations
+// whose storage or wire format cannot represent them exactly.
+func WithFloat64Ints() Option {
+	return func(c *client) {
+		c.float64Ints = true
+	}
+}
+
+func roundIntsThroughFloat64(record arrow.RecordBatch) arrow.RecordBatch {
+	columns := make([]arrow.Array, record.NumCols())
+	for i := range columns {
+		switch col := record.Column(i).(type) {
+		case *array.Int64:
+			bldr := array.NewInt64Builder(memory.DefaultAllocator)
+			for j := 0; j < col.Len(); j++ {
+				if col.IsNull(j) {
+					bldr.AppendNull()
+					continue
+				}
+				bldr.Append(int64(float64(col.Value(j))))
+			}
+			columns[i] = bldr.NewArray()
+			bldr.Release()
+		case *array.Uint64:
+			bldr := array.NewUint64Builder(memory.DefaultAllocator)
+			for j := 0; j < col.Len(); j++ {
+				if col.IsNull(j) {
+					bldr.AppendNull()
+					continue
+				}
+				bldr.Append(uint64(float64(col.Value(j))))
+			}
+			columns[i] = bldr.NewArray()
+			bldr.Release()
+		default:
+			columns[i] = record.Column(i)
+		}
+	}
+	return array.NewRecordBatch(record.Schema(), columns, record.NumRows())
+}
+
 func GetNewClient(options ...Option) plugin.NewClientFunc {
 	c := &client{
 		memoryDB:     make(map[string][]arrow.RecordBatch),
@@ -239,7 +282,11 @@ func (c *client) Write(ctx context.Context, msgs <-chan message.WriteMessage) er
 				return errors.New("table name not found in schema metadata")
 			}
 			table := c.tables[tableName]
-			c.overwrite(table, msg.Record)
+			record := msg.Record
+			if c.float64Ints {
+				record = roundIntsThroughFloat64(record)
+			}
+			c.overwrite(table, record)
 		}
 
 		c.memoryDBLock.Unlock()

internal/memdb/memdb_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/cloudquery/plugin-sdk/v4/plugin"
+	"github.com/cloudquery/plugin-sdk/v4/schema"
 )
 
 func TestPlugin(t *testing.T) {
@@ -22,6 +23,22 @@ func TestPlugin(t *testing.T) {
 	)
 }
 
+func TestPluginFloat64Ints(t *testing.T) {
+	ctx := context.Background()
+	p := plugin.NewPlugin("test", "development", GetNewClient(WithFloat64Ints()))
+	if err := p.Init(ctx, nil, plugin.NewClientOptions{}); err != nil {
+		t.Fatal(err)
+	}
+	plugin.TestWriterSuiteRunner(
+		t,
+		p,
+		plugin.WriterTestSuiteTests{
+			SafeMigrations: plugin.SafeMigrations{},
+		},
+		plugin.WithTestDataOptions(schema.TestSourceOptions{MaxIntegerBits: schema.Float64SafeIntegerBits}),
+	)
+}
+
 func TestPluginOnNewError(t *testing.T) {
 	ctx := context.Background()
 	p := plugin.NewPlugin("test", "development", NewMemDBClientErrOnNew)

plugin/diff.go

@@ -13,6 +13,99 @@ func RecordsDiff(sc *arrow.Schema, have, want []arrow.RecordBatch) string {
 	return TableDiff(array.NewTableFromRecords(sc, have), array.NewTableFromRecords(sc, want))
 }
 
+func containsInt64(dt arrow.DataType) bool {
+	switch dt.ID() {
+	case arrow.INT64, arrow.UINT64:
+		return true
+	}
+	if nested, ok := dt.(arrow.NestedType); ok {
+		for _, field := range nested.Fields() {
+			if containsInt64(field.Type) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+func equalAsFloat64(want, have arrow.Array) bool {
+	if want.Len() != have.Len() || !arrow.TypeEqual(want.DataType(), have.DataType()) {
+		return false
+	}
+	for i := 0; i < want.Len(); i++ {
+		if want.IsNull(i) != have.IsNull(i) {
+			return false
+		}
+	}
+
+	switch wantCol := want.(type) {
+	case *array.Int64:
+		haveCol := have.(*array.Int64)
+		for i := 0; i < wantCol.Len(); i++ {
+			if !wantCol.IsNull(i) && float64(wantCol.Value(i)) != float64(haveCol.Value(i)) {
+				return false
+			}
+		}
+		return true
+	case *array.Uint64:
+		haveCol := have.(*array.Uint64)
+		for i := 0; i < wantCol.Len(); i++ {
+			if !wantCol.IsNull(i) && float64(wantCol.Value(i)) != float64(haveCol.Value(i)) {
+				return false
+			}
+		}
+		return true
+	case *array.Struct:
+		haveCol := have.(*array.Struct)
+		for i := 0; i < wantCol.NumField(); i++ {
+			if !equalAsFloat64(wantCol.Field(i), haveCol.Field(i)) {
+				return false
+			}
+		}
+		return true
+	case *array.List:
+		haveCol := have.(*array.List)
+		return sameOffsets(wantCol.Offsets(), haveCol.Offsets()) && equalAsFloat64(wantCol.ListValues(), haveCol.ListValues())
+	case *array.LargeList:
+		haveCol := have.(*array.LargeList)
+		return sameLargeOffsets(wantCol.Offsets(), haveCol.Offsets()) && equalAsFloat64(wantCol.ListValues(), haveCol.ListValues())
+	case *array.Map:
+		haveCol := have.(*array.Map)
+		return sameOffsets(wantCol.Offsets(), haveCol.Offsets()) &&
+			equalAsFloat64(wantCol.Keys(), haveCol.Keys()) &&
+			equalAsFloat64(wantCol.Items(), haveCol.Items())
+	}
+
+	if containsInt64(want.DataType()) {
+		return false
+	}
+	return array.Equal(want, have)
+}
+
+func sameOffsets(want, have []int32) bool {
+	if len(want) != len(have) {
+		return false
+	}
+	for i := range want {
+		if want[i] != have[i] {
+			return false
+		}
+	}
+	return true
+}
+
+func sameLargeOffsets(want, have []int64) bool {
+	if len(want) != len(have) {
+		return false
+	}
+	for i := range want {
+		if want[i] != have[i] {
+			return false
+		}
+	}
+	return true
+}
+
 func TableDiff(have, want arrow.Table) string {
 	if array.TableApproxEqual(have, want, array.WithUnorderedMapKeys(true)) {
 		return ""
@@ -44,6 +137,9 @@ func TableDiff(have, want arrow.Table) string {
 			sb.WriteString(have.Schema().Field(i).Name)
 			sb.WriteString(": ")
 			sb.WriteString(diff)
+			if equalAsFloat64(wantCol, haveCol) {
+				sb.WriteString("values are equal after a float64 round-trip. If this destination cannot store 64-bit integers exactly, set schema.TestSourceOptions.MaxIntegerBits to schema.Float64SafeIntegerBits\n")
+			}
 			sb.WriteString("\n")
 		}
 	}

plugin/diff_test.go

@@ -0,0 +1,47 @@
+package plugin
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/apache/arrow-go/v18/arrow"
+	"github.com/apache/arrow-go/v18/arrow/array"
+	"github.com/apache/arrow-go/v18/arrow/memory"
+)
+
+func int64Record(t *testing.T, values []int64) arrow.RecordBatch {
+	t.Helper()
+	sc := arrow.NewSchema([]arrow.Field{{Name: "int64", Type: arrow.PrimitiveTypes.Int64}}, nil)
+	bldr := array.NewRecordBuilder(memory.DefaultAllocator, sc)
+	defer bldr.Release()
+	bldr.Field(0).(*array.Int64Builder).AppendValues(values, nil)
+	return bldr.NewRecordBatch()
+}
+
+func TestRecordsDiffHintsAtMaxIntegerBits(t *testing.T) {
+	sc := arrow.NewSchema([]arrow.Field{{Name: "int64", Type: arrow.PrimitiveTypes.Int64}}, nil)
+	exact := int64Record(t, []int64{-8717895732742165505})
+	rounded := int64Record(t, []int64{-8717895732742165504})
+
+	diff := RecordsDiff(sc, []arrow.RecordBatch{rounded}, []arrow.RecordBatch{exact})
+	if diff == "" {
+		t.Fatal("expected a diff between an exact and a float64-rounded int64")
+	}
+	if !strings.Contains(diff, "MaxIntegerBits") {
+		t.Errorf("expected the diff to point at MaxIntegerBits, got: %s", diff)
+	}
+}
+
+func TestRecordsDiffOmitsHintForUnrelatedDifference(t *testing.T) {
+	sc := arrow.NewSchema([]arrow.Field{{Name: "int64", Type: arrow.PrimitiveTypes.Int64}}, nil)
+	want := int64Record(t, []int64{1})
+	have := int64Record(t, []int64{2})
+
+	diff := RecordsDiff(sc, []arrow.RecordBatch{have}, []arrow.RecordBatch{want})
+	if diff == "" {
+		t.Fatal("expected a diff between 1 and 2")
+	}
+	if strings.Contains(diff, "MaxIntegerBits") {
+		t.Errorf("did not expect a MaxIntegerBits hint for a difference float64 cannot explain, got: %s", diff)
+	}
+}

plugin/testing_write_delete.go

@@ -96,6 +96,7 @@ func (s *WriterTestSuite) testDeleteStaleAll(ctx context.Context, t *testing.T)
 	normalRecord := tg.Generate(table, schema.GenTestDataOptions{
 		MaxRows:            rowsPerRecord,
 		TimePrecision:      s.genDatOptions.TimePrecision,
+		MaxIntegerBits:     s.genDatOptions.MaxIntegerBits,
 		SourceName:         "test",
 		SyncTime:           syncTime, // Generate call may truncate the value further based on the options
 		UseHomogeneousType: s.useHomogeneousTypes,
@@ -122,6 +123,7 @@ func (s *WriterTestSuite) testDeleteStaleAll(ctx context.Context, t *testing.T)
 	nullRecord := tg.Generate(table, schema.GenTestDataOptions{
 		MaxRows:            rowsPerRecord,
 		TimePrecision:      s.genDatOptions.TimePrecision,
+		MaxIntegerBits:     s.genDatOptions.MaxIntegerBits,
 		NullRows:           true,
 		SourceName:         "test",
 		SyncTime:           syncTime, // Generate call may truncate the value further based on the options

plugin/testing_write_insert.go

@@ -93,6 +93,7 @@ func (s *WriterTestSuite) testInsertAll(ctx context.Context) error {
 	normalRecord := tg.Generate(table, schema.GenTestDataOptions{
 		MaxRows:            rowsPerRecord,
 		TimePrecision:      s.genDatOptions.TimePrecision,
+		MaxIntegerBits:     s.genDatOptions.MaxIntegerBits,
 		UseHomogeneousType: s.useHomogeneousTypes,
 	})
 	if err := s.plugin.writeOne(ctx, &message.WriteInsert{
@@ -113,9 +114,10 @@ func (s *WriterTestSuite) testInsertAll(ctx context.Context) error {
 	}
 
 	nullRecord := tg.Generate(table, schema.GenTestDataOptions{
-		MaxRows:       rowsPerRecord,
-		TimePrecision: s.genDatOptions.TimePrecision,
-		NullRows:      true,
+		MaxRows:        rowsPerRecord,
+		TimePrecision:  s.genDatOptions.TimePrecision,
+		MaxIntegerBits: s.genDatOptions.MaxIntegerBits,
+		NullRows:       true,
 	})
 	if err := s.plugin.writeOne(ctx, &message.WriteInsert{
 		Record: nullRecord,

plugin/testing_write_migrate.go

@@ -42,6 +42,7 @@ func (s *WriterTestSuite) migrate(ctx context.Context, target *schema.Table, sou
 		SyncTime:           syncTime,
 		MaxRows:            rowsPerRecord,
 		TimePrecision:      s.genDatOptions.TimePrecision,
+		MaxIntegerBits:     s.genDatOptions.MaxIntegerBits,
 		UseHomogeneousType: s.useHomogeneousTypes,
 	}
 	// Test Generator should be initialized with the current number of items in the destination

plugin/testing_write_upsert.go

@@ -82,6 +82,7 @@ func (s *WriterTestSuite) testUpsertAll(ctx context.Context) error {
 	normalRecord := tg.Generate(table, schema.GenTestDataOptions{
 		MaxRows:            rowsPerRecord,
 		TimePrecision:      s.genDatOptions.TimePrecision,
+		MaxIntegerBits:     s.genDatOptions.MaxIntegerBits,
 		UseHomogeneousType: s.useHomogeneousTypes,
 	})
 	if err := s.plugin.writeOne(ctx, &message.WriteInsert{
@@ -107,7 +108,12 @@ func (s *WriterTestSuite) testUpsertAll(ctx context.Context) error {
 	}
 
 	tg.Reset()
-	nullRecord := tg.Generate(table, schema.GenTestDataOptions{MaxRows: rowsPerRecord, TimePrecision: s.genDatOptions.TimePrecision, NullRows: true})
+	nullRecord := tg.Generate(table, schema.GenTestDataOptions{
+		MaxRows:        rowsPerRecord,
+		TimePrecision:  s.genDatOptions.TimePrecision,
+		MaxIntegerBits: s.genDatOptions.MaxIntegerBits,
+		NullRows:       true,
+	})
 	if err := s.plugin.writeOne(ctx, &message.WriteInsert{
 		Record: nullRecord,
 	}); err != nil {