Implement batched audit log export utilities (#54684)

* Fix typo

* Implement audit log exporting utilities

* Improve batch export

api/client/client.go

@@ -2616,7 +2616,7 @@ func (c *Client) SearchEvents(ctx context.Context, fromUTC, toUTC time.Time, nam
 // SearchUnstructuredEvents allows searching for events with a full pagination support
 // and returns events in an unstructured format (json like).
 // This method is used by the Teleport event-handler plugin to receive events
-// from the auth server wihout having to support the Protobuf event schema.
+// from the auth server without having to support the Protobuf event schema.
 func (c *Client) SearchUnstructuredEvents(ctx context.Context, fromUTC, toUTC time.Time, namespace string, eventTypes []string, limit int, order types.EventOrder, startKey string) ([]*auditlogpb.EventUnstructured, string, error) {
 	request := &auditlogpb.GetUnstructuredEventsRequest{
 		Namespace:  namespace,

lib/config/configuration.go

@@ -515,6 +515,7 @@ func ApplyFileConfig(fc *FileConfig, cfg *servicecfg.Config) error {
 		cfg.AccessGraph.CA = fc.AccessGraph.CA
 		// TODO(tigrato): change this behavior when we drop support for plain text connections
 		cfg.AccessGraph.Insecure = fc.AccessGraph.Insecure
+		cfg.AccessGraph.AuditLog = servicecfg.AuditLogConfig(fc.AccessGraph.AuditLog)
 	}
 
 	applyString(fc.NodeName, &cfg.Hostname)

lib/config/fileconf.go

@@ -832,6 +832,16 @@ type AccessGraph struct {
 	CA string `yaml:"ca"`
 	// Insecure is true if the AccessGraph service should not verify the CA.
 	Insecure bool `yaml:"insecure"`
+	// AuditLog contains audit log export details.
+	AuditLog AuditLogConfig `yaml:"audit_log"`
+}
+
+// AuditLogConfig specifies the audit log event export setup.
+type AuditLogConfig struct {
+	// Enabled indicates if Audit Log event exporting is enabled.
+	Enabled bool `yaml:"enabled"`
+	// StartDate is the start date for exporting audit logs. It defaults to 90 days ago on the first export.
+	StartDate time.Time `yaml:"start_date"`
 }
 
 // Opsgenie represents the configuration for the Opsgenie plugin.

lib/events/export/date_exporter.go

@@ -19,6 +19,7 @@
 package export
 
 import (
+	"cmp"
 	"context"
 	"log/slog"
 	"sync"
@@ -27,6 +28,7 @@ import (
 
 	"github.com/gravitational/trace"
 	"golang.org/x/time/rate"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	auditlogpb "github.com/gravitational/teleport/api/gen/proto/go/teleport/auditlog/v1"
@@ -50,6 +52,10 @@ type DateExporterConfig struct {
 	// Export is the callback used to export events. Must be safe for concurrent use if
 	// the Concurrency parameter is greater than 1.
 	Export func(ctx context.Context, event *auditlogpb.ExportEventUnstructured) error
+	// BatchExport is the callback with configuration used to export multiple
+	// events in batches.
+	BatchExport *BatchExportConfig
+
 	// OnIdle is an optional callback that gets invoked periodically when the exporter is idle. Note that it is
 	// safe to close the exporter or inspect its state from within this callback, but waiting on the exporter's
 	// Done channel within this callback will deadlock.
@@ -69,20 +75,24 @@ func (cfg *DateExporterConfig) CheckAndSetDefaults() error {
 	if cfg.Client == nil {
 		return trace.BadParameter("missing required parameter Client in DateExporterConfig")
 	}
-	if cfg.Export == nil {
-		return trace.BadParameter("missing required parameter Export in DateExporterConfig")
+	if cfg.Export == nil && cfg.BatchExport == nil {
+		return trace.BadParameter("missing required parameter Export or BatchExport in DateExporterConfig")
+	}
+	if cfg.BatchExport != nil && cfg.BatchExport.Callback == nil {
+		return trace.BadParameter("missing parameter BatchExport.Callback in DateExporterConfig")
+	}
+	if cfg.Export != nil && cfg.BatchExport != nil {
+		return trace.BadParameter("only one of Export or BatchExport may be set in ExporterConfig")
 	}
 	if cfg.Date.IsZero() {
 		return trace.BadParameter("missing required parameter Date in DateExporterConfig")
 	}
-	if cfg.Concurrency == 0 {
-		cfg.Concurrency = 1
-	}
-	if cfg.MaxBackoff == 0 {
-		cfg.MaxBackoff = 90 * time.Second
-	}
-	if cfg.PollInterval == 0 {
-		cfg.PollInterval = 16 * time.Second
+	cfg.Concurrency = cmp.Or(cfg.Concurrency, 1)
+	cfg.MaxBackoff = cmp.Or(cfg.MaxBackoff, 90*time.Second)
+	cfg.PollInterval = cmp.Or(cfg.PollInterval, 16*time.Second)
+	if cfg.BatchExport != nil {
+		cfg.BatchExport.MaxDelay = cmp.Or(cfg.BatchExport.MaxDelay, 5*time.Second)
+		cfg.BatchExport.MaxSize = cmp.Or(cfg.BatchExport.MaxSize, 2*1024*1024 /* 2MiB */)
 	}
 	return nil
 }
@@ -437,7 +447,13 @@ Outer:
 			Cursor: entry.getCursor(),
 		})
 
-		if err := e.exportEvents(ctx, events, entry); err != nil {
+		var err error
+		if e.cfg.Export != nil {
+			err = e.exportEvents(ctx, events, entry)
+		} else {
+			err = e.batchExportEvents(ctx, events, entry, chunk)
+		}
+		if err != nil {
 			failures++
 
 			if e.chunkLogLimiter.Allow() {
@@ -458,6 +474,121 @@ Outer:
 	}
 }
 
+type batch struct {
+	maxSize int
+	chunk   string
+	entry   *chunkEntry
+
+	size   int
+	events []*auditlogpb.EventUnstructured
+	cursor string
+}
+
+func (b *batch) addEvent(event *auditlogpb.ExportEventUnstructured) bool {
+	e := event.GetEvent()
+	size := proto.Size(e)
+	if b.size+size <= b.maxSize || len(b.events) == 0 {
+		b.size += size
+		b.events = append(b.events, e)
+		b.cursor = event.GetCursor()
+		return true
+	}
+	return false
+}
+
+func (b *batch) reset() {
+	b.events = nil
+	b.size = 0
+	b.cursor = ""
+}
+
+func (b *batch) isEmpty() bool {
+	return len(b.events) == 0
+}
+
+// batchExportEvents reads events from the provided stream and exports them in batches.
+// Batching adheres to the export configuration: batches are sent either when
+// they reach the configured maximum size (MaxSize) or after the maximum delay
+// (MaxDelay) has passed with pending events. Processing continues until the
+// stream closes or an error occurs.
+func (e *DateExporter) batchExportEvents(ctx context.Context, stream stream.Stream[*auditlogpb.ExportEventUnstructured], entry *chunkEntry, chunk string) error {
+	exportEventCh := make(chan *auditlogpb.ExportEventUnstructured, 1)
+	go func() {
+		for stream.Next() {
+			exportEventCh <- stream.Item()
+		}
+		close(exportEventCh)
+	}()
+
+	batch := &batch{
+		maxSize: e.cfg.BatchExport.MaxSize,
+		chunk:   chunk,
+		entry:   entry,
+	}
+	timer := time.NewTimer(e.cfg.BatchExport.MaxDelay)
+	defer timer.Stop()
+loop:
+	for {
+		var unprocessedEvent *auditlogpb.ExportEventUnstructured
+		select {
+		case exportEvent, ok := <-exportEventCh:
+			if !ok {
+				// all events have been processed
+				break loop
+			}
+			if batch.addEvent(exportEvent) {
+				continue
+			}
+			unprocessedEvent = exportEvent
+		case <-timer.C:
+			if batch.isEmpty() {
+				timer.Reset(e.cfg.BatchExport.MaxDelay)
+				continue
+			}
+		}
+		err := e.batchExport(ctx, batch, false /*completed*/)
+		if err != nil {
+			stream.Done()
+			return trace.Wrap(err)
+		}
+
+		batch.reset()
+		timer.Reset(e.cfg.BatchExport.MaxDelay)
+
+		if unprocessedEvent != nil {
+			batch.addEvent(unprocessedEvent)
+		}
+	}
+
+	if err := stream.Done(); err != nil {
+		return trace.Wrap(err)
+	}
+
+	// One final call back with completed set to true, even if batch is empty.
+	err := e.batchExport(ctx, batch, true /*completed*/)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	return nil
+}
+
+// batchExport calls the batch export callback with the given batched events and
+// the derived resume state.
+func (e *DateExporter) batchExport(ctx context.Context, batch *batch, completed bool) error {
+	cursor := batch.cursor
+	if completed {
+		cursor = ""
+	}
+	resumeState := BulkExportResumeState{
+		Chunk:     batch.chunk,
+		Cursor:    cursor,
+		Date:      e.cfg.Date,
+		Completed: completed,
+	}
+	batch.entry.setCursor(cursor)
+	return e.cfg.BatchExport.Callback(ctx, batch.events, resumeState)
+}
+
 // exportEvents exports all events from the provided stream, updating the supplied entry on each successful export.
 func (e *DateExporter) exportEvents(ctx context.Context, events stream.Stream[*auditlogpb.ExportEventUnstructured], entry *chunkEntry) error {
 	for events.Next() {

lib/events/export/exporter.go

@@ -19,6 +19,7 @@
 package export
 
 import (
+	"cmp"
 	"context"
 	"log/slog"
 	"slices"
@@ -57,22 +58,64 @@ func (s *ExporterState) Clone() ExporterState {
 	return out
 }
 
+// BulkExportResumeState contains the information used to update the resume state
+// (specifically the Cursor and Completed flag) for a given Date and Chunk
+// on an external system managing the overall bulk export process.
+type BulkExportResumeState struct {
+	// Date is the specific day for which the resume state is being updated,
+	// expected to be normalized to UTC midnight (00:00:00).
+	Date time.Time
+	// Chunk identifies the specific data segment or partition (within the Date)
+	// for which this resume state applies.
+	Chunk string
+	// Cursor marks the position (e.g., the next event ID or offset) within
+	// the Chunk from where the export should resume.
+	Cursor string
+	// Completed indicates whether processing for this specific Date and Chunk
+	// has been fully and successfully exported.
+	Completed bool
+}
+
+// BatchExportConfig instructs the Exporter to use the
+// BatchExportConfig.Callback with events batched up to MaxSize and a maximum
+// delay of BatchExportConfig.MaxDelay between calls.
+type BatchExportConfig struct {
+	// Callback is used to export batched events. It must be safe for concurrent
+	// use if the ExportConfig.Concurrency parameter is greater than 1.
+	Callback func(ctx context.Context, events []*auditlogpb.EventUnstructured, resumeState BulkExportResumeState) error
+	// MaxDelay specifies the maximum duration to wait before invoking the
+	// Callback with a batch of events, even if MaxSize has not been reached.
+	// If the zero value is provided for this field, a default duration of 5 *
+	// time.Second will be used by the exporter.
+	MaxDelay time.Duration
+	// MaxSize defines the maximum total size in bytes of the events accumulated
+	// in a single batch before the Callback is invoked. If the zero value is
+	// provided for this field, a default size of 2 MB will be used by the
+	// exporter.
+	MaxSize int
+}
+
 // ExporterConfig configured an exporter.
 type ExporterConfig struct {
 	// Client is the audit event client used to fetch and export events.
 	Client Client
 	// StartDate is the date from which to start exporting events.
 	StartDate time.Time
+
 	// Export is the callback used to export events. Must be safe for concurrent use if
 	// the Concurrency parameter is greater than 1.
 	Export func(ctx context.Context, event *auditlogpb.ExportEventUnstructured) error
+	// BatchExport is the callback with configuration used to export bulk events in batches.
+	BatchExport *BatchExportConfig
 	// OnIdle is an optional callback that gets invoked periodically when the exporter is idle. Note that it is
 	// safe to close the exporter or inspect its state from within this callback, but waiting on the exporter's
 	// Done channel within this callback will deadlock. This callback is an asynchronous signal and additional
 	// events may be discovered concurrently with its invocation.
 	OnIdle func(ctx context.Context)
+
 	// PreviousState is an optional parameter used to resume from a previous date export run.
 	PreviousState ExporterState
+
 	// Concurrency sets the maximum number of event chunks that will be processed concurrently
 	// for a given date (defaults to 1). Note that the total number of inflight chunk processing
 	// may be up to Conurrency * (BacklogSize + 1).
@@ -94,20 +137,21 @@ func (cfg *ExporterConfig) CheckAndSetDefaults() error {
 	if cfg.StartDate.IsZero() {
 		return trace.BadParameter("missing required parameter StartDate in ExporterConfig")
 	}
-	if cfg.Export == nil {
-		return trace.BadParameter("missing required parameter Export in ExporterConfig")
+	if cfg.Export == nil && cfg.BatchExport == nil {
+		return trace.BadParameter("missing required parameter Export or BatchExport in ExporterConfig")
 	}
-	if cfg.Concurrency == 0 {
-		cfg.Concurrency = 1
+	if cfg.Export != nil && cfg.BatchExport != nil {
+		return trace.BadParameter("only one of Export or BatchExport may be set in ExporterConfig")
 	}
-	if cfg.BacklogSize == 0 {
-		cfg.BacklogSize = 1
+	if cfg.BatchExport != nil && cfg.BatchExport.Callback == nil {
+		return trace.BadParameter("missing parameter BatchExport.Callback in ExporterConfig")
 	}
-	if cfg.MaxBackoff == 0 {
-		cfg.MaxBackoff = 90 * time.Second
-	}
-	if cfg.PollInterval == 0 {
-		cfg.PollInterval = 16 * time.Second
+	cfg.Concurrency = cmp.Or(cfg.Concurrency, 1)
+	cfg.MaxBackoff = cmp.Or(cfg.MaxBackoff, 90*time.Second)
+	cfg.PollInterval = cmp.Or(cfg.PollInterval, 16*time.Second)
+	if cfg.BatchExport != nil {
+		cfg.BatchExport.MaxDelay = cmp.Or(cfg.BatchExport.MaxDelay, 5*time.Second)
+		cfg.BatchExport.MaxSize = cmp.Or(cfg.BatchExport.MaxSize, 2*1024*1024 /* 2MiB */)
 	}
 	return nil
 }
@@ -189,7 +233,7 @@ func (e *Exporter) Close() {
 
 // Done provides a channel that will be closed when the exporter has completed processing all inflight dates. When saving the
 // final state of the exporter for future resumption, this channel must be waited upon before state is loaded. Note that the date
-// exporter never termiantes unless Close is called, so waiting on Done is only meaningful after Close has been called.
+// exporter never terminates unless Close is called, so waiting on Done is only meaningful after Close has been called.
 func (e *Exporter) Done() <-chan struct{} {
 	return e.done
 }
@@ -405,6 +449,7 @@ func (e *Exporter) resumeExportLocked(ctx context.Context, date time.Time, state
 		Client:        e.cfg.Client,
 		Date:          date,
 		Export:        e.cfg.Export,
+		BatchExport:   e.cfg.BatchExport,
 		OnIdle:        onIdle,
 		PreviousState: state,
 		Concurrency:   e.cfg.Concurrency,

lib/service/servicecfg/config.go

@@ -329,6 +329,17 @@ type AccessGraphConfig struct {
 
 	// Insecure is true if the connection to the Access Graph service should be insecure
 	Insecure bool
+
+	// AuditLog contains audit log export details.
+	AuditLog AuditLogConfig
+}
+
+// AuditLogConfig specifies the audit log event export setup.
+type AuditLogConfig struct {
+	// Enabled indicates if Audit Log event exporting is enabled.
+	Enabled bool
+	// StartDate is the start date for exporting audit logs. It defaults to 90 days ago on the first export.
+	StartDate time.Time
 }
 
 // RoleAndIdentityEvent is a role and its corresponding identity event.