fix(writers): StreamingBatchWriter hangs with non-append mode

writers/streamingbatchwriter/mocktimer_test.go

@@ -1,17 +1,25 @@
 package streamingbatchwriter
 
 import (
+	"sync"
 	"time"
 
 	"github.com/cloudquery/plugin-sdk/v4/writers"
 )
 
 type mockTicker struct {
-	expire chan time.Time
+	expire  chan time.Time
+	stopped sync.Once
 }
 
 func (t *mockTicker) Stop() {
-	close(t.expire)
+	t.stopped.Do(func() {
+		close(t.expire)
+	})
+}
+
+func (t *mockTicker) Tick() {
+	t.expire <- time.Now()
 }
 
 func (*mockTicker) Reset(time.Duration) {}
@@ -20,12 +28,12 @@ func (t *mockTicker) Chan() <-chan time.Time {
 	return t.expire
 }
 
-func newMockTicker() (writers.TickerFunc, chan<- time.Time) {
+func newMockTicker() (writers.TickerFunc, func()) {
 	expire := make(chan time.Time)
 	t := &mockTicker{
 		expire: expire,
 	}
 	return func(time.Duration) writers.Ticker {
 		return t
-	}, expire
+	}, t.Tick
 }

writers/streamingbatchwriter/streamingbatchwriter.go

@@ -156,13 +156,17 @@ func (w *StreamingBatchWriter) Close(context.Context) error {
 	}
 	w.workersWaitGroup.Wait()
 
-	w.insertWorkers = nil
+	w.insertWorkers = make(map[string]*streamingWorkerManager[*message.WriteInsert])
+	w.migrateWorker = nil
+	w.deleteWorker = nil
+	w.lastMsgType = writers.MsgTypeUnset
 
 	return nil
 }
 
 func (w *StreamingBatchWriter) Write(ctx context.Context, msgs <-chan message.WriteMessage) error {
 	errCh := make(chan error)
+	defer close(errCh)
 
 	go func() {
 		for err := range errCh {
@@ -172,7 +176,7 @@ func (w *StreamingBatchWriter) Write(ctx context.Context, msgs <-chan message.Wr
 
 	for msg := range msgs {
 		msgType := writers.MsgID(msg)
-		if w.lastMsgType != msgType {
+		if w.lastMsgType != writers.MsgTypeUnset && w.lastMsgType != msgType {
 			if err := w.Flush(ctx); err != nil {
 				return err
 			}
@@ -183,12 +187,7 @@ func (w *StreamingBatchWriter) Write(ctx context.Context, msgs <-chan message.Wr
 		}
 	}
 
-	if err := w.Flush(ctx); err != nil {
-		return err
-	}
-
-	close(errCh)
-	return nil
+	return w.Close(ctx)
 }
 
 func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- error, msg message.WriteMessage) error {

writers/streamingbatchwriter/streamingbatchwriter_test.go

@@ -227,7 +227,7 @@ func TestStreamingBatchTimeout(t *testing.T) {
 	ch := make(chan message.WriteMessage)
 
 	testClient := newClient()
-	tickerFn, expire := newMockTicker()
+	tickerFn, tickFn := newMockTicker()
 
 	wr, err := New(testClient, withTickerFn(tickerFn))
 	if err != nil {
@@ -258,7 +258,7 @@ func TestStreamingBatchTimeout(t *testing.T) {
 	}
 
 	// flush
-	close(expire)
+	tickFn()
 	waitForLength(t, testClient.MessageLen, messageTypeInsert, 1)
 
 	close(ch)
@@ -332,7 +332,7 @@ func TestStreamingBatchUpserts(t *testing.T) {
 	ch := make(chan message.WriteMessage)
 
 	testClient := newClient()
-	tickerFn, expire := newMockTicker()
+	tickerFn, tickFn := newMockTicker()
 	wr, err := New(testClient, WithBatchSizeRows(2), withTickerFn(tickerFn))
 	if err != nil {
 		t.Fatal(err)
@@ -363,7 +363,7 @@ func TestStreamingBatchUpserts(t *testing.T) {
 	time.Sleep(50 * time.Millisecond)
 
 	// flush the batch
-	close(expire)
+	tickFn()
 	waitForLength(t, testClient.MessageLen, messageTypeInsert, 2)
 
 	close(ch)

writers/streamingbatchwriter/unimplemented.go

@@ -18,9 +18,12 @@ func (IgnoreMigrateTable) MigrateTable(_ context.Context, ch <-chan *message.Wri
 	return nil
 }
 
-// UnimplementedDeleteStale is a dummy handler to error on DeleteStale messages
+// UnimplementedDeleteStale is a dummy handler to consume and error on DeleteStale messages
 type UnimplementedDeleteStale struct{}
 
-func (UnimplementedDeleteStale) DeleteStale(_ context.Context, _ <-chan *message.WriteDeleteStale) error {
+func (UnimplementedDeleteStale) DeleteStale(_ context.Context, ch <-chan *message.WriteDeleteStale) error {
+	// nolint:revive
+	for range ch {
+	}
 	return fmt.Errorf("DeleteStale: %w", plugin.ErrNotImplemented)
 }