feat: Support DeleteRecord in all writers (#1295)

<!--
Explain what problem this PR addresses
-->

---

Author: bbernays

writers/batchwriter/batchwriter.go

@@ -17,6 +17,7 @@ type Client interface {
 	MigrateTables(context.Context, message.WriteMigrateTables) error
 	WriteTableBatch(ctx context.Context, name string, messages message.WriteInserts) error
 	DeleteStale(context.Context, message.WriteDeleteStales) error
+	DeleteRecord(context.Context, message.WriteDeleteRecords) error
 }
 
 type BatchWriter struct {
@@ -29,6 +30,8 @@ type BatchWriter struct {
 	migrateTableMessages message.WriteMigrateTables
 	deleteStaleLock      sync.Mutex
 	deleteStaleMessages  message.WriteDeleteStales
+	deleteRecordLock     sync.Mutex
+	deleteRecordMessages message.WriteDeleteRecords
 
 	logger         zerolog.Logger
 	batchTimeout   time.Duration
@@ -199,6 +202,19 @@ func (w *BatchWriter) flushDeleteStaleTables(ctx context.Context) error {
 	return nil
 }
 
+func (w *BatchWriter) flushDeleteRecordTables(ctx context.Context) error {
+	w.deleteRecordLock.Lock()
+	defer w.deleteRecordLock.Unlock()
+	if len(w.deleteRecordMessages) == 0 {
+		return nil
+	}
+	if err := w.client.DeleteRecord(ctx, w.deleteRecordMessages); err != nil {
+		return err
+	}
+	w.deleteRecordMessages = w.deleteRecordMessages[:0]
+	return nil
+}
+
 func (w *BatchWriter) flushInsert(tableName string) {
 	w.workersLock.RLock()
 	worker, ok := w.workers[tableName]
@@ -239,6 +255,26 @@ func (w *BatchWriter) Write(ctx context.Context, msgs <-chan message.WriteMessag
 					return err
 				}
 			}
+		case *message.WriteDeleteRecord:
+			if err := w.flushMigrateTables(ctx); err != nil {
+				return err
+			}
+			if err := w.flushDeleteStaleTables(ctx); err != nil {
+				return err
+			}
+			// Ensure all related workers are flushed
+			for _, rel := range m.TableRelations {
+				w.flushInsert(rel.TableName)
+			}
+			w.deleteRecordLock.Lock()
+			w.deleteRecordMessages = append(w.deleteRecordMessages, m)
+			l := len(w.deleteRecordMessages)
+			w.deleteRecordLock.Unlock()
+			if w.batchSize > 0 && l > w.batchSize {
+				if err := w.flushDeleteRecordTables(ctx); err != nil {
+					return err
+				}
+			}
 		case *message.WriteInsert:
 			if err := w.flushMigrateTables(ctx); err != nil {
 				return err

writers/batchwriter/batchwriter_test.go

@@ -18,6 +18,7 @@ type testBatchClient struct {
 	migrateTables message.WriteMigrateTables
 	inserts       message.WriteInserts
 	deleteStales  message.WriteDeleteStales
+	deleteRecords message.WriteDeleteRecords
 }
 
 func (c *testBatchClient) MigrateTablesLen() int {
@@ -58,6 +59,13 @@ func (c *testBatchClient) DeleteStale(_ context.Context, messages message.WriteD
 	return nil
 }
 
+func (c *testBatchClient) DeleteRecord(_ context.Context, messages message.WriteDeleteRecords) error {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+	c.deleteRecords = append(c.deleteRecords, messages...)
+	return nil
+}
+
 var batchTestTables = schema.Tables{
 	{
 		Name: "table1",

writers/batchwriter/unimplemented.go

@@ -19,3 +19,9 @@ type UnimplementedDeleteStale struct{}
 func (UnimplementedDeleteStale) DeleteStale(context.Context, message.WriteDeleteStales) error {
 	return fmt.Errorf("DeleteStale: %w", plugin.ErrNotImplemented)
 }
+
+type UnimplementedDeleteRecord struct{}
+
+func (UnimplementedDeleteRecord) DeleteRecord(context.Context, message.WriteDeleteRecords) error {
+	return fmt.Errorf("DeleteRecord: %w", plugin.ErrNotImplemented)
+}

writers/batchwriter/unimplemented_test.go

@@ -10,6 +10,7 @@ import (
 type testDummyClient struct {
 	batchwriter.IgnoreMigrateTables
 	batchwriter.UnimplementedDeleteStale
+	batchwriter.UnimplementedDeleteRecord
 }
 
 func (testDummyClient) WriteTableBatch(context.Context, string, message.WriteInserts) error {

writers/streamingbatchwriter/streamingbatchwriter.go

@@ -40,6 +40,9 @@ type Client interface {
 	// DeleteStale should block and handle WriteDeleteStale messages until the channel is closed.
 	DeleteStale(context.Context, <-chan *message.WriteDeleteStale) error
 
+	// DeleteRecords should block and handle WriteDeleteRecord messages until the channel is closed.
+	DeleteRecords(context.Context, <-chan *message.WriteDeleteRecord) error
+
 	// WriteTable should block and handle writes to a single table until the channel is closed. Table metadata can be found in the first WriteInsert message.
 	// The channel is closed when all inserts in the batch have been sent. New batches, if any, will be sent on a new call to WriteTable.
 	WriteTable(context.Context, <-chan *message.WriteInsert) error
@@ -48,9 +51,11 @@ type Client interface {
 type StreamingBatchWriter struct {
 	client Client
 
-	insertWorkers    map[string]*streamingWorkerManager[*message.WriteInsert]
-	migrateWorker    *streamingWorkerManager[*message.WriteMigrateTable]
-	deleteWorker     *streamingWorkerManager[*message.WriteDeleteStale]
+	insertWorkers      map[string]*streamingWorkerManager[*message.WriteInsert]
+	migrateWorker      *streamingWorkerManager[*message.WriteMigrateTable]
+	deleteStaleWorker  *streamingWorkerManager[*message.WriteDeleteStale]
+	deleteRecordWorker *streamingWorkerManager[*message.WriteDeleteRecord]
+
 	workersLock      sync.RWMutex
 	workersWaitGroup sync.WaitGroup
 
@@ -128,9 +133,14 @@ func (w *StreamingBatchWriter) Flush(_ context.Context) error {
 		w.migrateWorker.flush <- done
 		<-done
 	}
-	if w.deleteWorker != nil {
+	if w.deleteStaleWorker != nil {
+		done := make(chan bool)
+		w.deleteStaleWorker.flush <- done
+		<-done
+	}
+	if w.deleteRecordWorker != nil {
 		done := make(chan bool)
-		w.deleteWorker.flush <- done
+		w.deleteRecordWorker.flush <- done
 		<-done
 	}
 	for _, worker := range w.insertWorkers {
@@ -151,14 +161,18 @@ func (w *StreamingBatchWriter) Close(context.Context) error {
 	if w.migrateWorker != nil {
 		close(w.migrateWorker.ch)
 	}
-	if w.deleteWorker != nil {
-		close(w.deleteWorker.ch)
+	if w.deleteStaleWorker != nil {
+		close(w.deleteStaleWorker.ch)
+	}
+	if w.deleteRecordWorker != nil {
+		close(w.deleteRecordWorker.ch)
 	}
 	w.workersWaitGroup.Wait()
 
 	w.insertWorkers = make(map[string]*streamingWorkerManager[*message.WriteInsert])
 	w.migrateWorker = nil
-	w.deleteWorker = nil
+	w.deleteStaleWorker = nil
+	w.deleteRecordWorker = nil
 	w.lastMsgType = writers.MsgTypeUnset
 
 	return nil
@@ -232,13 +246,13 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 	case *message.WriteDeleteStale:
 		w.workersLock.Lock()
 		defer w.workersLock.Unlock()
-		if w.deleteWorker != nil {
-			w.deleteWorker.ch <- m
+		if w.deleteStaleWorker != nil {
+			w.deleteStaleWorker.ch <- m
 			return nil
 		}
 		ch := make(chan *message.WriteDeleteStale)
 		flush := make(chan chan bool)
-		w.deleteWorker = &streamingWorkerManager[*message.WriteDeleteStale]{
+		w.deleteStaleWorker = &streamingWorkerManager[*message.WriteDeleteStale]{
 			ch:        ch,
 			writeFunc: w.client.DeleteStale,
 
@@ -251,8 +265,8 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 		}
 
 		w.workersWaitGroup.Add(1)
-		go w.deleteWorker.run(ctx, &w.workersWaitGroup, tableName)
-		w.deleteWorker.ch <- m
+		go w.deleteStaleWorker.run(ctx, &w.workersWaitGroup, tableName)
+		w.deleteStaleWorker.ch <- m
 		return nil
 	case *message.WriteInsert:
 		w.workersLock.RLock()
@@ -285,7 +299,32 @@ func (w *StreamingBatchWriter) startWorker(ctx context.Context, errCh chan<- err
 		go wr.run(ctx, &w.workersWaitGroup, tableName)
 		ch <- m
 		return nil
+	case *message.WriteDeleteRecord:
+		w.workersLock.Lock()
+		defer w.workersLock.Unlock()
+		if w.deleteRecordWorker != nil {
+			w.deleteRecordWorker.ch <- m
+			return nil
+		}
+		ch := make(chan *message.WriteDeleteRecord)
+		flush := make(chan chan bool)
+		// TODO: flush all workers for nested tables as well (See https://github.com/cloudquery/plugin-sdk/issues/1296)
+		w.deleteRecordWorker = &streamingWorkerManager[*message.WriteDeleteRecord]{
+			ch:        ch,
+			writeFunc: w.client.DeleteRecords,
+
+			flush: flush,
+			errCh: errCh,
 
+			batchSizeRows: w.batchSizeRows,
+			batchTimeout:  w.batchTimeout,
+			tickerFn:      w.tickerFn,
+		}
+
+		w.workersWaitGroup.Add(1)
+		go w.deleteRecordWorker.run(ctx, &w.workersWaitGroup, tableName)
+		w.deleteRecordWorker.ch <- m
+		return nil
 	default:
 		return fmt.Errorf("unhandled message type: %T", msg)
 	}

writers/streamingbatchwriter/streamingbatchwriter_test.go

@@ -19,6 +19,7 @@ const (
 	messageTypeMigrateTable messageType = iota
 	messageTypeInsert
 	messageTypeDeleteStale
+	messageTypeDeleteRecord
 )
 
 type testStreamingBatchClient struct {
@@ -79,6 +80,14 @@ func (c *testStreamingBatchClient) DeleteStale(ctx context.Context, msgs <-chan
 	return c.handleTypeCommit(ctx, messageTypeDeleteStale, key)
 }
 
+func (c *testStreamingBatchClient) DeleteRecords(ctx context.Context, msgs <-chan *message.WriteDeleteRecord) error {
+	key := ""
+	for m := range msgs {
+		key = c.handleTypeMessage(ctx, messageTypeDeleteRecord, m, key)
+	}
+	return c.handleTypeCommit(ctx, messageTypeDeleteRecord, key)
+}
+
 func (c *testStreamingBatchClient) handleTypeMessage(_ context.Context, t messageType, msg message.WriteMessage, key string) string {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()

writers/streamingbatchwriter/unimplemented.go

@@ -27,3 +27,12 @@ func (UnimplementedDeleteStale) DeleteStale(_ context.Context, ch <-chan *messag
 	}
 	return fmt.Errorf("DeleteStale: %w", plugin.ErrNotImplemented)
 }
+
+type UnimplementedDeleteRecords struct{}
+
+func (UnimplementedDeleteRecords) DeleteRecords(_ context.Context, ch <-chan *message.WriteDeleteRecord) error {
+	// nolint:revive
+	for range ch {
+	}
+	return fmt.Errorf("DeleteRecords: %w", plugin.ErrNotImplemented)
+}

writers/streamingbatchwriter/unimplemented_test.go

@@ -10,6 +10,7 @@ import (
 type testDummyClient struct {
 	streamingbatchwriter.IgnoreMigrateTable
 	streamingbatchwriter.UnimplementedDeleteStale
+	streamingbatchwriter.UnimplementedDeleteRecords
 }
 
 func (testDummyClient) WriteTable(context.Context, <-chan *message.WriteInsert) error {