feat: Allow for DeleteRecord tests to work with batch writers

Author: maaarcelino

plugin/testing_write.go

@@ -83,6 +83,8 @@ type WriterTestSuiteTests struct {
 	SkipSpecificMigrations Migrations
 
 	SkipSpecificWriteTests WriteTests
+
+	DeleteRecordBatchSize int
 }
 
 type NewPluginFunc func() *Plugin

plugin/testing_write_delete.go

@@ -177,70 +177,58 @@ func (s *WriterTestSuite) testDeleteRecordBasic(ctx context.Context, t *testing.
 	r.NoErrorf(err, "failed to read")
 	r.EqualValuesf(1, TotalRows(records), "unexpected amount of items")
 
-	// create value for delete statement but nothing will be deleted because ID value isn't present
-	bldrDeleteNoMatch := array.NewRecordBuilder(memory.DefaultAllocator, (&schema.Table{
-		Name: tableName,
-		Columns: schema.ColumnList{
-			schema.Column{Name: "id", Type: arrow.PrimitiveTypes.Int64},
-		},
-	}).ToArrowSchema())
-	bldrDeleteNoMatch.Field(0).(*array.Int64Builder).Append(1)
-	deleteValue := bldrDeleteNoMatch.NewRecord()
-
-	r.NoErrorf(s.plugin.writeOne(ctx, &message.WriteDeleteRecord{
-		DeleteRecord: message.DeleteRecord{
-			TableName: table.Name,
-			WhereClause: message.PredicateGroups{
-				{
-					GroupingType: "AND",
-					Predicates: []message.Predicate{
-						{
-							Operator: "eq",
-							Column:   "id",
-							Record:   deleteValue,
-						},
-					},
-				},
-			},
-		},
-	}), "failed to delete record no match")
+	r.NoError(s.plugin.WriteAll(ctx, s.toMessageBatch(table, 1)), "failed to delete record")
 
 	records, err = s.plugin.readAll(ctx, table)
 	r.NoErrorf(err, "failed to read after delete with no match")
 	r.EqualValuesf(1, TotalRows(records), "unexpected amount of items after delete with no match")
 	r.Emptyf(RecordsDiff(table.ToArrowSchema(), records, []arrow.Record{record1}), "record differs after delete with no match")
 
-	// create value for delete statement will be delete One record
-	bldrDeleteMatch := array.NewRecordBuilder(memory.DefaultAllocator, (&schema.Table{
-		Name: tableName,
+	r.NoError(s.plugin.WriteAll(ctx, s.toMessageBatch(table, 0)), "failed to delete record")
+
+	records, err = s.plugin.readAll(ctx, table)
+	r.NoErrorf(err, "failed to read after delete with match")
+	r.EqualValuesf(0, TotalRows(records), "unexpected amount of items after delete with match")
+}
+
+func (s *WriterTestSuite) toMessageBatch(table *schema.Table, val int64) []message.WriteMessage {
+	batchSize := 1
+	if s.tests.DeleteRecordBatchSize > 0 {
+		batchSize = s.tests.DeleteRecordBatchSize
+	}
+
+	messages := make([]message.WriteMessage, batchSize)
+
+	builder := array.NewRecordBuilder(memory.DefaultAllocator, (&schema.Table{
+		Name: table.Name,
 		Columns: schema.ColumnList{
 			schema.Column{Name: "id", Type: arrow.PrimitiveTypes.Int64},
 		},
 	}).ToArrowSchema())
-	bldrDeleteMatch.Field(0).(*array.Int64Builder).Append(0)
-	deleteValue = bldrDeleteMatch.NewRecord()
-
-	r.NoErrorf(s.plugin.writeOne(ctx, &message.WriteDeleteRecord{
-		DeleteRecord: message.DeleteRecord{
-			TableName: table.Name,
-			WhereClause: message.PredicateGroups{
-				{
-					GroupingType: "AND",
-					Predicates: []message.Predicate{
-						{
-							Operator: "eq",
-							Column:   "id",
-							Record:   deleteValue,
+	builder.Field(0).(*array.Int64Builder).Append(val)
+	deleteValue := builder.NewRecord()
+
+	for i, _ := range messages {
+		messages[i] = &message.WriteDeleteRecord{
+			DeleteRecord: message.DeleteRecord{
+				TableName: table.Name,
+				WhereClause: message.PredicateGroups{
+					{
+						GroupingType: "AND",
+						Predicates: []message.Predicate{
+							{
+								Operator: "eq",
+								Column:   "id",
+								Record:   deleteValue,
+							},
 						},
 					},
 				},
 			},
-		},
-	}), "failed to delete record no match")
+		}
+	}
 
-	records, err = s.plugin.readAll(ctx, table)
-	r.NoErrorf(err, "failed to read after delete with match")
-	r.EqualValuesf(0, TotalRows(records), "unexpected amount of items after delete with match")
+	return messages
 }
 
 func (s *WriterTestSuite) testDeleteAllRecords(ctx context.Context, t *testing.T) {