fix: Change logic for batch writing to write when batch size is reached, not exceeded (#2153)

The current logic for batch writing prevents us from cleanly testing the `DeleteRecord` handling in the plugin. Currently we have to set `BatchSize: 1` and send two `DeleteRecord` messages in the test, which is not the cleanest logic. This PR changes the logic so the batch is flushed when the batch size is reached, not exceeded (so for `BatchSize: 1` it will flush after one `DeleteRecord` is received).

This also matches the logic we have in https://github.com/cloudquery/plugin-sdk/blob/main/internal/batch/cap.go#L7.

Author: maaarcelino

writers/batchwriter/batchwriter.go

@@ -281,7 +281,7 @@ func (w *BatchWriter) Write(ctx context.Context, msgs <-chan message.WriteMessag
 			w.deleteStaleMessages = append(w.deleteStaleMessages, m)
 			l := int64(len(w.deleteStaleMessages))
 			w.deleteStaleLock.Unlock()
-			if w.batchSize > 0 && l > w.batchSize {
+			if w.isLimitReached(l) {
 				if err := w.flushDeleteStaleTables(ctx); err != nil {
 					return err
 				}
@@ -301,7 +301,7 @@ func (w *BatchWriter) Write(ctx context.Context, msgs <-chan message.WriteMessag
 			w.deleteRecordMessages = append(w.deleteRecordMessages, m)
 			l := int64(len(w.deleteRecordMessages))
 			w.deleteRecordLock.Unlock()
-			if w.batchSize > 0 && l > w.batchSize {
+			if w.isLimitReached(l) {
 				if err := w.flushDeleteRecordTables(ctx); err != nil {
 					return err
 				}
@@ -328,7 +328,7 @@ func (w *BatchWriter) Write(ctx context.Context, msgs <-chan message.WriteMessag
 			w.migrateTableMessages = append(w.migrateTableMessages, m)
 			l := int64(len(w.migrateTableMessages))
 			w.migrateTableLock.Unlock()
-			if w.batchSize > 0 && l > w.batchSize {
+			if w.isLimitReached(l) {
 				if err := w.flushMigrateTables(ctx); err != nil {
 					return err
 				}
@@ -338,6 +338,12 @@ func (w *BatchWriter) Write(ctx context.Context, msgs <-chan message.WriteMessag
 	return nil
 }
 
+func (w *BatchWriter) isLimitReached(rowCount int64) bool {
+	limit := batch.CappedAt(0, w.batchSize)
+	limit.AddRows(rowCount)
+	return limit.ReachedLimit()
+}
+
 func (w *BatchWriter) startWorker(_ context.Context, msg *message.WriteInsert) error {
 	w.workersLock.RLock()
 	md := msg.Record.Schema().Metadata()