revert writers to their original form

candiduslynx · candiduslynx · commit a06aacd63870 · 2024-06-13T11:18:57.000+03:00
diff --git a/writers/batchwriter/batchwriter.go b/writers/batchwriter/batchwriter.go
@@ -6,7 +6,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/cloudquery/plugin-sdk/v4/internal/batch"
+	"github.com/apache/arrow/go/v16/arrow/util"
 	"github.com/cloudquery/plugin-sdk/v4/message"
 	"github.com/cloudquery/plugin-sdk/v4/schema"
 	"github.com/cloudquery/plugin-sdk/v4/writers"
@@ -122,7 +122,7 @@ func (w *BatchWriter) Close(context.Context) error {
 }
 
 func (w *BatchWriter) worker(ctx context.Context, tableName string, ch <-chan *message.WriteInsert, flush <-chan chan bool) {
-	limit := batch.CappedAt(w.batchSizeBytes, w.batchSize)
+	var bytes, rows int64
 	resources := make([]*message.WriteInsert, 0, w.batchSize) // at least we have 1 row per record
 
 	ticker := writers.NewTicker(w.batchTimeout)
@@ -134,50 +134,47 @@ func (w *BatchWriter) worker(ctx context.Context, tableName string, ch <-chan *m
 		w.flushTable(ctx, tableName, resources)
 		clear(resources)
 		resources = resources[:0]
-		limit.Reset()
+		bytes, rows = 0, 0
 	}
-
 	for {
 		select {
 		case r, ok := <-ch:
 			if !ok {
-				if limit.Rows() > 0 {
+				if rows > 0 {
 					w.flushTable(ctx, tableName, resources)
 				}
 				return
 			}
 
-			if r.Record.NumRows() == 0 {
-				// skip empty ones
-				continue
-			}
-
-			add, toFlush, rest := batch.SliceRecord(r.Record, limit)
-			if add != nil {
-				resources = append(resources, &message.WriteInsert{Record: add.Record})
-			}
-			if len(toFlush) > 0 || rest != nil || limit.ReachedLimit() {
-				// flush current batch
-				send()
-				ticker.Reset(w.batchTimeout)
-			}
-			for _, sliceToFlush := range toFlush {
-				resources = append(resources, &message.WriteInsert{Record: sliceToFlush})
+			recordRows, recordBytes := r.Record.NumRows(), util.TotalRecordSize(r.Record)
+			if (w.batchSize > 0 && rows+recordRows > w.batchSize) ||
+				(w.batchSizeBytes > 0 && bytes+recordBytes > w.batchSizeBytes) {
+				if rows == 0 {
+					// New record overflows batch by itself.
+					// Flush right away.
+					// TODO: slice
+					resources = append(resources, r)
+					send()
+					ticker.Reset(w.batchTimeout)
+					continue
+				}
+				// rows > 0
 				send()
 				ticker.Reset(w.batchTimeout)
 			}
-
-			// set the remainder
-			if rest != nil {
-				resources = append(resources, &message.WriteInsert{Record: rest.Record})
+			if recordRows > 0 {
+				// only save records with rows
+				resources = append(resources, r)
+				rows += recordRows
+				bytes += recordBytes
 			}
 
 		case <-tickerCh:
-			if limit.Rows() > 0 {
+			if rows > 0 {
 				send()
 			}
 		case done := <-flush:
-			if limit.Rows() > 0 {
+			if rows > 0 {
 				send()
 				ticker.Reset(w.batchTimeout)
 			}
diff --git a/writers/mixedbatchwriter/mixedbatchwriter.go b/writers/mixedbatchwriter/mixedbatchwriter.go
@@ -4,7 +4,7 @@ import (
 	"context"
 	"time"
 
-	"github.com/cloudquery/plugin-sdk/v4/internal/batch"
+	"github.com/apache/arrow/go/v16/arrow/util"
 	"github.com/cloudquery/plugin-sdk/v4/message"
 	"github.com/cloudquery/plugin-sdk/v4/writers"
 	"github.com/rs/zerolog"
@@ -92,7 +92,8 @@ func (w *MixedBatchWriter) Write(ctx context.Context, msgChan <-chan message.Wri
 	insert := &insertBatchManager{
 		batch:     make([]*message.WriteInsert, 0, w.batchSize),
 		writeFunc: w.client.InsertBatch,
-		limit:     batch.CappedAt(w.batchSizeBytes, w.batchSize),
+		maxRows:   w.batchSize,
+		maxBytes:  w.batchSizeBytes,
 		logger:    w.logger,
 	}
 	deleteStale := &batchManager[message.WriteDeleteStales, *message.WriteDeleteStale]{
@@ -198,53 +199,47 @@ func (m *batchManager[A, T]) flush(ctx context.Context) error {
 
 // special batch manager for insert messages that also keeps track of the total size of the batch
 type insertBatchManager struct {
-	batch     []*message.WriteInsert
-	writeFunc func(ctx context.Context, messages message.WriteInserts) error
-	limit     *batch.Cap
-	logger    zerolog.Logger
+	batch              []*message.WriteInsert
+	writeFunc          func(ctx context.Context, messages message.WriteInserts) error
+	curRows, maxRows   int64
+	curBytes, maxBytes int64
+	logger             zerolog.Logger
 }
 
 func (m *insertBatchManager) append(ctx context.Context, msg *message.WriteInsert) error {
-	add, toFlush, rest := batch.SliceRecord(msg.Record, m.limit)
-	if add != nil {
-		m.batch = append(m.batch, &message.WriteInsert{Record: add.Record})
-	}
-	if len(toFlush) > 0 || rest != nil || m.limit.ReachedLimit() {
-		// flush current batch
-		if err := m.flush(ctx); err != nil {
-			return err
-		}
-	}
-	for _, sliceToFlush := range toFlush {
-		m.batch = append(m.batch, &message.WriteInsert{Record: sliceToFlush})
+	recordRows, recordBytes := msg.Record.NumRows(), util.TotalRecordSize(msg.Record)
+	if (m.maxRows > 0 && m.curRows+recordRows > m.maxRows) ||
+		(m.maxBytes > 0 && m.curBytes+recordBytes > m.maxBytes) {
 		if err := m.flush(ctx); err != nil {
 			return err
 		}
 	}
 
-	// set the remainder
-	if rest != nil {
-		m.batch = append(m.batch, &message.WriteInsert{Record: rest.Record})
+	if recordRows > 0 {
+		// only save records with rows
+		m.batch = append(m.batch, msg)
+		m.curRows += recordRows
+		m.curBytes += recordBytes
 	}
 
 	return nil
 }
 
 func (m *insertBatchManager) flush(ctx context.Context) error {
-	if m.limit.Rows() == 0 {
+	if m.curRows == 0 {
 		// no rows to insert
 		return nil
 	}
 	start := time.Now()
 	err := m.writeFunc(ctx, m.batch)
 	if err != nil {
-		m.logger.Err(err).Int64("len", m.limit.Rows()).Dur("duration", time.Since(start)).Msg("failed to write batch")
+		m.logger.Err(err).Int64("len", m.curRows).Dur("duration", time.Since(start)).Msg("failed to write batch")
 		return err
 	}
-	m.logger.Debug().Int64("len", m.limit.Rows()).Dur("duration", time.Since(start)).Msg("batch written successfully")
+	m.logger.Debug().Int64("len", m.curRows).Dur("duration", time.Since(start)).Msg("batch written successfully")
 
 	clear(m.batch) // GC can work
 	m.batch = m.batch[:0]
-	m.limit.Reset()
+	m.curRows, m.curBytes = 0, 0
 	return nil
 }
diff --git a/writers/streamingbatchwriter/streamingbatchwriter.go b/writers/streamingbatchwriter/streamingbatchwriter.go
@@ -25,7 +25,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/cloudquery/plugin-sdk/v4/internal/batch"
+	"github.com/apache/arrow/go/v16/arrow/util"
 	"github.com/cloudquery/plugin-sdk/v4/message"
 	"github.com/cloudquery/plugin-sdk/v4/schema"
 	"github.com/cloudquery/plugin-sdk/v4/writers"
@@ -233,9 +233,9 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 			flush: make(chan chan bool),
 			errCh: errCh,
 
-			limit:        batch.CappedAt(0, w.batchSizeRows),
-			batchTimeout: w.batchTimeout,
-			tickerFn:     w.tickerFn,
+			batchSizeRows: w.batchSizeRows,
+			batchTimeout:  w.batchTimeout,
+			tickerFn:      w.tickerFn,
 		}
 
 		w.workersWaitGroup.Add(1)
@@ -257,9 +257,9 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 			flush: make(chan chan bool),
 			errCh: errCh,
 
-			limit:        batch.CappedAt(0, w.batchSizeRows),
-			batchTimeout: w.batchTimeout,
-			tickerFn:     w.tickerFn,
+			batchSizeRows: w.batchSizeRows,
+			batchTimeout:  w.batchTimeout,
+			tickerFn:      w.tickerFn,
 		}
 
 		w.workersWaitGroup.Add(1)
@@ -283,9 +283,10 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 			flush: make(chan chan bool),
 			errCh: errCh,
 
-			limit:        batch.CappedAt(w.batchSizeBytes, w.batchSizeRows),
-			batchTimeout: w.batchTimeout,
-			tickerFn:     w.tickerFn,
+			batchSizeRows:  w.batchSizeRows,
+			batchSizeBytes: w.batchSizeBytes,
+			batchTimeout:   w.batchTimeout,
+			tickerFn:       w.tickerFn,
 		}
 		w.workersLock.Lock()
 		wrOld, ok := w.insertWorkers[tableName]
@@ -319,9 +320,9 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 			flush: make(chan chan bool),
 			errCh: errCh,
 
-			limit:        batch.CappedAt(w.batchSizeBytes, w.batchSizeRows),
-			batchTimeout: w.batchTimeout,
-			tickerFn:     w.tickerFn,
+			batchSizeRows: w.batchSizeRows,
+			batchTimeout:  w.batchTimeout,
+			tickerFn:      w.tickerFn,
 		}
 
 		w.workersWaitGroup.Add(1)
@@ -340,17 +341,19 @@ type streamingWorkerManager[T message.WriteMessage] struct {
 	flush chan chan bool
 	errCh chan<- error
 
-	limit        *batch.Cap
-	batchTimeout time.Duration
-	tickerFn     writers.TickerFunc
+	batchSizeRows  int64
+	batchSizeBytes int64
+	batchTimeout   time.Duration
+	tickerFn       writers.TickerFunc
 }
 
 func (s *streamingWorkerManager[T]) run(ctx context.Context, wg *sync.WaitGroup, tableName string) {
 	defer wg.Done()
 	var (
-		clientCh    chan T
-		clientErrCh chan error
-		open        bool
+		clientCh            chan T
+		clientErrCh         chan error
+		open                bool
+		sizeBytes, sizeRows int64
 	)
 
 	ensureOpened := func() {
@@ -379,7 +382,7 @@ func (s *streamingWorkerManager[T]) run(ctx context.Context, wg *sync.WaitGroup,
 			}
 		}
 		open = false
-		s.limit.Reset()
+		sizeBytes, sizeRows = 0, 0
 	}
 	defer closeFlush()
 
@@ -395,45 +398,44 @@ func (s *streamingWorkerManager[T]) run(ctx context.Context, wg *sync.WaitGroup,
 				return
 			}
 
+			recordRows := int64(1) // at least 1 row for messages without records
+			var recordBytes int64
 			if ins, ok := any(r).(*message.WriteInsert); ok {
-				add, toFlush, rest := batch.SliceRecord(ins.Record, s.limit)
-				if add != nil {
-					ensureOpened()
-					clientCh <- any(&message.WriteInsert{Record: add.Record}).(T)
-				}
-				if len(toFlush) > 0 || rest != nil || s.limit.ReachedLimit() {
-					// flush current batch
-					closeFlush()
-					ticker.Reset(s.batchTimeout)
-				}
-				for _, sliceToFlush := range toFlush {
+				recordBytes = util.TotalRecordSize(ins.Record)
+				recordRows = ins.Record.NumRows()
+			}
+
+			if (s.batchSizeRows > 0 && sizeRows+recordRows > s.batchSizeRows) ||
+				(s.batchSizeBytes > 0 && sizeBytes+recordBytes > s.batchSizeBytes) {
+				if sizeRows == 0 {
+					// New record overflows batch by itself.
+					// Flush right away.
+					// TODO: slice
 					ensureOpened()
-					clientCh <- any(&message.WriteInsert{Record: sliceToFlush}).(T)
+					clientCh <- r
 					closeFlush()
 					ticker.Reset(s.batchTimeout)
+					continue
 				}
+				// sizeRows > 0
+				closeFlush()
+				ticker.Reset(s.batchTimeout)
+			}
 
-				// set the remainder
-				if rest != nil {
-					ensureOpened()
-					clientCh <- any(&message.WriteInsert{Record: rest.Record}).(T)
-				}
-			} else {
+			if recordRows > 0 {
+				// only save records with rows
 				ensureOpened()
 				clientCh <- r
-				s.limit.AddRows(1)
-				if s.limit.ReachedLimit() {
-					closeFlush()
-					ticker.Reset(s.batchTimeout)
-				}
+				sizeRows += recordRows
+				sizeBytes += recordBytes
 			}
 
 		case <-tickerCh:
-			if s.limit.Rows() > 0 {
+			if sizeRows > 0 {
 				closeFlush()
 			}
 		case done := <-s.flush:
-			if s.limit.Rows() > 0 {
+			if sizeRows > 0 {
 				closeFlush()
 				ticker.Reset(s.batchTimeout)
 			}
