fix(writers): Move to sub packages (#1011)

For convenient `Option`s etc.

Also adds default batch size values to the streamingbatchwriter and some compile-time interface checks.

Author: disq

writers/batch.go renamed to writers/batchwriter/batchwriter.go

@@ -1,4 +1,4 @@
-package writers
+package batchwriter
 
 import (
 	"context"
@@ -9,27 +9,18 @@ import (
 	"github.com/apache/arrow/go/v13/arrow/util"
 	"github.com/cloudquery/plugin-sdk/v4/message"
 	"github.com/cloudquery/plugin-sdk/v4/schema"
+	"github.com/cloudquery/plugin-sdk/v4/writers"
 	"github.com/rs/zerolog"
 )
 
-type Writer interface {
-	Write(ctx context.Context, res <-chan message.WriteMessage) error
-}
-
-const (
-	DefaultBatchTimeoutSeconds = 20
-	DefaultBatchSize           = 10000
-	DefaultBatchSizeBytes      = 5 * 1024 * 1024 // 5 MiB
-)
-
-type BatchWriterClient interface {
+type Client interface {
 	MigrateTables(context.Context, []*message.WriteMigrateTable) error
 	WriteTableBatch(ctx context.Context, name string, msgs []*message.WriteInsert) error
 	DeleteStale(context.Context, []*message.WriteDeleteStale) error
 }
 
 type BatchWriter struct {
-	client           BatchWriterClient
+	client           Client
 	workers          map[string]*worker
 	workersLock      sync.RWMutex
 	workersWaitGroup sync.WaitGroup
@@ -45,6 +36,9 @@ type BatchWriter struct {
 	batchSizeBytes int
 }
 
+// Assert at compile-time that BatchWriter implements the Writer interface
+var _ writers.Writer = (*BatchWriter)(nil)
+
 type Option func(*BatchWriter)
 
 func WithLogger(logger zerolog.Logger) Option {
@@ -77,14 +71,14 @@ type worker struct {
 	flush chan chan bool
 }
 
-func NewBatchWriter(client BatchWriterClient, opts ...Option) (*BatchWriter, error) {
+func New(client Client, opts ...Option) (*BatchWriter, error) {
 	c := &BatchWriter{
 		client:         client,
 		workers:        make(map[string]*worker),
 		logger:         zerolog.Nop(),
-		batchTimeout:   DefaultBatchTimeoutSeconds * time.Second,
-		batchSize:      DefaultBatchSize,
-		batchSizeBytes: DefaultBatchSizeBytes,
+		batchTimeout:   writers.DefaultBatchTimeoutSeconds * time.Second,
+		batchSize:      writers.DefaultBatchSize,
+		batchSizeBytes: writers.DefaultBatchSizeBytes,
 	}
 	for _, opt := range opts {
 		opt(c)

writers/batch_test.go renamed to writers/batchwriter/batchwriter_test.go

@@ -1,4 +1,4 @@
-package writers
+package batchwriter
 
 import (
 	"context"
@@ -76,7 +76,7 @@ func TestBatchFlushDifferentMessages(t *testing.T) {
 	ctx := context.Background()
 
 	testClient := &testBatchClient{}
-	wr, err := NewBatchWriter(testClient)
+	wr, err := New(testClient)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -117,7 +117,7 @@ func TestBatchSize(t *testing.T) {
 	ctx := context.Background()
 
 	testClient := &testBatchClient{}
-	wr, err := NewBatchWriter(testClient, WithBatchSize(2))
+	wr, err := New(testClient, WithBatchSize(2))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -155,7 +155,7 @@ func TestBatchTimeout(t *testing.T) {
 	ctx := context.Background()
 
 	testClient := &testBatchClient{}
-	wr, err := NewBatchWriter(testClient, WithBatchTimeout(time.Second))
+	wr, err := New(testClient, WithBatchTimeout(time.Second))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -190,7 +190,7 @@ func TestBatchUpserts(t *testing.T) {
 	ctx := context.Background()
 
 	testClient := &testBatchClient{}
-	wr, err := NewBatchWriter(testClient, WithBatchSize(2), WithBatchTimeout(time.Second))
+	wr, err := New(testClient, WithBatchSize(2), WithBatchTimeout(time.Second))
 	if err != nil {
 		t.Fatal(err)
 	}

writers/mixed_batch.go renamed to writers/mixedbatchwriter/mixedbatchwriter.go

@@ -1,65 +1,66 @@
-package writers
+package mixedbatchwriter
 
 import (
 	"context"
 	"time"
 
 	"github.com/apache/arrow/go/v13/arrow/util"
 	"github.com/cloudquery/plugin-sdk/v4/message"
+	"github.com/cloudquery/plugin-sdk/v4/writers"
 	"github.com/rs/zerolog"
 )
 
-// MixedBatchClient is a client that will receive batches of messages with a mixture of tables.
-type MixedBatchClient interface {
+// Client is a client that will receive batches of messages with a mixture of tables.
+type Client interface {
 	MigrateTableBatch(ctx context.Context, messages []*message.WriteMigrateTable) error
 	InsertBatch(ctx context.Context, messages []*message.WriteInsert) error
 	DeleteStaleBatch(ctx context.Context, messages []*message.WriteDeleteStale) error
 }
 
 type MixedBatchWriter struct {
-	client         MixedBatchClient
+	client         Client
 	logger         zerolog.Logger
 	batchTimeout   time.Duration
 	batchSize      int
 	batchSizeBytes int
 }
 
 // Assert at compile-time that MixedBatchWriter implements the Writer interface
-var _ Writer = (*MixedBatchWriter)(nil)
+var _ writers.Writer = (*MixedBatchWriter)(nil)
 
-type MixedBatchWriterOption func(writer *MixedBatchWriter)
+type Option func(writer *MixedBatchWriter)
 
-func WithMixedBatchWriterLogger(logger zerolog.Logger) MixedBatchWriterOption {
+func WithLogger(logger zerolog.Logger) Option {
 	return func(p *MixedBatchWriter) {
 		p.logger = logger
 	}
 }
 
-func WithMixedBatchWriterBatchTimeout(timeout time.Duration) MixedBatchWriterOption {
+func WithBatchTimeout(timeout time.Duration) Option {
 	return func(p *MixedBatchWriter) {
 		p.batchTimeout = timeout
 	}
 }
 
-func WithMixedBatchWriterBatchSize(size int) MixedBatchWriterOption {
+func WithBatchSize(size int) Option {
 	return func(p *MixedBatchWriter) {
 		p.batchSize = size
 	}
 }
 
-func WithMixedBatchWriterBatchSizeBytes(size int) MixedBatchWriterOption {
+func WithBatchSizeBytes(size int) Option {
 	return func(p *MixedBatchWriter) {
 		p.batchSizeBytes = size
 	}
 }
 
-func NewMixedBatchWriter(client MixedBatchClient, opts ...MixedBatchWriterOption) (*MixedBatchWriter, error) {
+func New(client Client, opts ...Option) (*MixedBatchWriter, error) {
 	c := &MixedBatchWriter{
 		client:         client,
 		logger:         zerolog.Nop(),
-		batchTimeout:   DefaultBatchTimeoutSeconds * time.Second,
-		batchSize:      DefaultBatchSize,
-		batchSizeBytes: DefaultBatchSizeBytes,
+		batchTimeout:   writers.DefaultBatchTimeoutSeconds * time.Second,
+		batchSize:      writers.DefaultBatchSize,
+		batchSizeBytes: writers.DefaultBatchSizeBytes,
 	}
 	for _, opt := range opts {
 		opt(c)
@@ -82,23 +83,23 @@ func (w *MixedBatchWriter) Write(ctx context.Context, msgChan <-chan message.Wri
 		batch:     make([]*message.WriteDeleteStale, 0, w.batchSize),
 		writeFunc: w.client.DeleteStaleBatch,
 	}
-	flush := func(msgType msgType) error {
+	flush := func(msgType writers.MsgType) error {
 		switch msgType {
-		case msgTypeMigrateTable:
+		case writers.MsgTypeMigrateTable:
 			return migrateTable.flush(ctx)
-		case msgTypeInsert:
+		case writers.MsgTypeInsert:
 			return insert.flush(ctx)
-		case msgTypeDeleteStale:
+		case writers.MsgTypeDeleteStale:
 			return deleteStale.flush(ctx)
 		default:
 			panic("unknown message type")
 		}
 	}
-	prevMsgType := msgTypeUnset
+	prevMsgType := writers.MsgTypeUnset
 	var err error
 	for msg := range msgChan {
-		msgType := msgID(msg)
-		if prevMsgType != msgTypeUnset && prevMsgType != msgType {
+		msgType := writers.MsgID(msg)
+		if prevMsgType != writers.MsgTypeUnset && prevMsgType != msgType {
 			if err := flush(prevMsgType); err != nil {
 				return err
 			}
@@ -118,7 +119,7 @@ func (w *MixedBatchWriter) Write(ctx context.Context, msgChan <-chan message.Wri
 			return err
 		}
 	}
-	if prevMsgType == msgTypeUnset {
+	if prevMsgType == writers.MsgTypeUnset {
 		return nil
 	}
 	return flush(prevMsgType)

writers/mixed_batch_test.go renamed to writers/mixedbatchwriter/mixedbatchwriter_test.go

@@ -1,4 +1,4 @@
-package writers
+package mixedbatchwriter_test
 
 import (
 	"context"
@@ -10,6 +10,7 @@ import (
 	"github.com/apache/arrow/go/v13/arrow/memory"
 	"github.com/cloudquery/plugin-sdk/v4/message"
 	"github.com/cloudquery/plugin-sdk/v4/schema"
+	"github.com/cloudquery/plugin-sdk/v4/writers/mixedbatchwriter"
 )
 
 type testMixedBatchClient struct {
@@ -43,7 +44,7 @@ func (c *testMixedBatchClient) DeleteStaleBatch(_ context.Context, msgs []*messa
 	return nil
 }
 
-var _ MixedBatchClient = (*testMixedBatchClient)(nil)
+var _ mixedbatchwriter.Client = (*testMixedBatchClient)(nil)
 
 func TestMixedBatchWriter(t *testing.T) {
 	ctx := context.Background()
@@ -169,7 +170,7 @@ func TestMixedBatchWriter(t *testing.T) {
 			client := &testMixedBatchClient{
 				receivedBatches: make([][]message.WriteMessage, 0),
 			}
-			wr, err := NewMixedBatchWriter(client)
+			wr, err := mixedbatchwriter.New(client)
 			if err != nil {
 				t.Fatal(err)
 			}

writers/msgtype.go

@@ -6,23 +6,23 @@ import (
 	"github.com/cloudquery/plugin-sdk/v4/message"
 )
 
-type msgType int
+type MsgType int
 
 const (
-	msgTypeUnset msgType = iota
-	msgTypeMigrateTable
-	msgTypeInsert
-	msgTypeDeleteStale
+	MsgTypeUnset MsgType = iota
+	MsgTypeMigrateTable
+	MsgTypeInsert
+	MsgTypeDeleteStale
 )
 
-func msgID(msg message.WriteMessage) msgType {
+func MsgID(msg message.WriteMessage) MsgType {
 	switch msg.(type) {
 	case *message.WriteMigrateTable:
-		return msgTypeMigrateTable
+		return MsgTypeMigrateTable
 	case *message.WriteInsert:
-		return msgTypeInsert
+		return MsgTypeInsert
 	case *message.WriteDeleteStale:
-		return msgTypeDeleteStale
+		return MsgTypeDeleteStale
 	}
 	panic("unknown message type: " + reflect.TypeOf(msg).Name())
 }