internal/app/options/connectorflags.go (1)

604-608: Clarify prefix flag behavior in usage text.

The usage says "Override or append to" but the implementation only overrides (via Destination). If the CLI flag is provided, it completely replaces the prefix from the connection string.

Consider updating the usage to reflect actual behavior:

 		altsrc.NewStringFlag(&cli.StringFlag{
 			Name:        "prefix",
-			Usage:       "Override or append to the key prefix derived from the connection string",
+			Usage:       "Override the key prefix derived from the connection string",
 			Destination: &settings.Prefix,
 		}),

connectors/s3/connector.go (2)

418-432: Consider accepting context parameter.

flushBatch uses context.Background() which won't respect cancellation signals from the caller. If the application is shutting down, this S3 put operation could block indefinitely.

-func (c *connector) flushBatch(namespace string, taskID uint, docs [][]byte) error {
+func (c *connector) flushBatch(ctx context.Context, namespace string, taskID uint, docs [][]byte) error {
 	payload := buildJSONArray(docs)
 	key := c.objectKey(namespace, taskID)
-	_, err := c.client.PutObject(context.Background(), &s3.PutObjectInput{
+	_, err := c.client.PutObject(ctx, &s3.PutObjectInput{

---

462-474: Buffer pre-allocation is underestimated.

buf.Grow(len(docs) * 2) significantly underestimates the required capacity since each document could be much larger than 2 bytes. While this won't cause correctness issues (the buffer grows automatically), it defeats the purpose of pre-allocation.

Consider estimating based on actual document sizes:

 func buildJSONArray(docs [][]byte) []byte {
 	var buf bytes.Buffer
-	buf.Grow(len(docs) * 2)
+	totalSize := 2 + len(docs) // brackets + commas
+	for _, doc := range docs {
+		totalSize += len(doc)
+	}
+	buf.Grow(totalSize)
 	buf.WriteByte('[')

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

connectors/s3/connector.go (2)

373-395: Empty files still created for tasks with no data.

When batch is nil (no data was written for this task), the code creates an empty storedBatch{} and proceeds to call flushBatch, which writes an empty JSON array to S3. This pollutes the bucket with unnecessary files.

Return early when there's no data to flush:

 func (c *connector) OnTaskCompletionBarrierHandler(taskID uint) error {
 	batch := c.detachBatch(taskID)
 	if batch == nil {
 		slog.Debug("s3 connector received barrier with no data", "taskId", taskID)
-		batch = &storedBatch{}
+		return nil
 	}
+	if len(batch.docs) == 0 {
+		slog.Debug("s3 connector received barrier with empty batch", "taskId", taskID)
+		return nil
+	}
 	if err := c.flushBatch(batch.namespace, taskID, batch.docs); err != nil {

---

560-587: Race condition in concurrent metadata updates.

updateMetadataAfterFlush performs an unsynchronized read-modify-write. When multiple tasks for the same namespace complete concurrently, metadata updates can be lost because each goroutine reads, modifies, and writes independently.

Example race:

1. Task A reads metadata {file1: 100}
2. Task B reads metadata {file1: 100}
3. Task A writes {file1: 100, task-2.json: 50}
4. Task B writes {file1: 100, task-3.json: 75} — task-2.json entry is lost

Add per-namespace mutex to serialize metadata updates:

 type connector struct {
 	adiomv1connect.UnimplementedConnectorServiceHandler

 	client       *s3.Client
 	settings     ConnectorSettings
 	batchesMutex sync.Mutex
 	batches      map[taskKey]*storedBatch

+	metadataMutex sync.Mutex // Serialize metadata updates

 	errMutex sync.RWMutex
 	err      error
 }

Then protect the read-modify-write in updateMetadataAfterFlush:

 func (c *connector) updateMetadataAfterFlush(ctx context.Context, namespace string, taskID uint, recordCount uint64) error {
+	c.metadataMutex.Lock()
+	defer c.metadataMutex.Unlock()
+
 	// Read current metadata (or create empty map if doesn't exist)
 	metadata, err := c.readMetadata(ctx, namespace)

For better concurrency with many namespaces, consider per-namespace locks using sync.Map.

connectors/common/base.go (1)

1098-1098: TaskId type narrowing from uint to uint32.

The conversion from uint (which is 64-bit on 64-bit systems) to uint32 could theoretically overflow if task IDs exceed 2³²-1. While task IDs are typically small sequential numbers making overflow unlikely in practice, consider documenting this limitation or adding validation if task ID limits are a concern.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

connectors/s3/connector.go (5)

68-68: Consider per-namespace mutexes for better concurrency.

While the single metadataMutex correctly prevents race conditions (addressing the past review comment), it serializes metadata updates across all namespaces. For workloads with many concurrent tasks updating different namespaces, this global lock can become a bottleneck.

Consider using a sync.Map to hold per-namespace mutexes:

 type connector struct {
 	adiomv1connect.UnimplementedConnectorServiceHandler
 
 	client       *s3.Client
 	settings     ConnectorSettings
 	batchesMutex sync.Mutex
 	batches      map[taskKey]*storedBatch
 
-	metadataMutex sync.Mutex // Serialize metadata updates
+	namespaceMutexes sync.Map // map[string]*sync.Mutex for per-namespace locking
 
 	errMutex sync.RWMutex
 	err      error
 }

Then in updateMetadataAfterFlush, acquire the namespace-specific mutex:

func (c *connector) updateMetadataAfterFlush(ctx context.Context, namespace string, taskID uint, recordCount uint64) error {
	// Get or create mutex for this namespace
	mu, _ := c.namespaceMutexes.LoadOrStore(namespace, &sync.Mutex{})
	namespaceMu := mu.(*sync.Mutex)
	
	namespaceMu.Lock()
	defer namespaceMu.Unlock()
	
	// ... rest of the function unchanged
}

---

287-329: Consider memory implications for large files.

The current implementation loads the entire JSON array into memory at once. For very large S3 files (e.g., hundreds of MB or GB), this could cause memory pressure or OOM issues.

If you expect to handle large files, consider:

1. Implementing pagination/chunking within the file (using the NextCursor field)
2. Setting size limits on individual S3 files
3. Streaming the JSON array parsing rather than decoding all at once
4. Monitoring memory usage in production for files above a certain size threshold

---

393-398: Redundant check: batch is already confirmed non-empty.

The check if len(batch.docs) > 0 is unnecessary because the function already returns early at lines 382-385 if the batch is empty.

Apply this diff to remove the redundant check:

 	// Update metadata file atomically after successful flush
-	if len(batch.docs) > 0 {
-		if err := c.updateMetadataAfterFlush(context.Background(), batch.namespace, taskID, uint64(len(batch.docs))); err != nil {
-			slog.Error("failed to update metadata", "namespace", batch.namespace, "taskId", taskID, "err", err)
-			// Log error but don't fail the barrier - the data was successfully flushed
-		}
+	if err := c.updateMetadataAfterFlush(context.Background(), batch.namespace, taskID, uint64(len(batch.docs))); err != nil {
+		slog.Error("failed to update metadata", "namespace", batch.namespace, "taskId", taskID, "err", err)
+		// Log error but don't fail the barrier - the data was successfully flushed
 	}

---

424-438: Accept context parameter for better cancellation control.

The function uses context.Background() for the S3 PutObject call, which means the flush operation cannot be cancelled or timeout-controlled by the caller. If OnTaskCompletionBarrierHandler needs to respect deadlines or cancellation, this will be problematic.

Apply this diff to thread the context through:

-func (c *connector) flushBatch(namespace string, taskID uint, docs [][]byte) error {
+func (c *connector) flushBatch(ctx context.Context, namespace string, taskID uint, docs [][]byte) error {
 	payload := buildJSONArray(docs)
 	key := c.objectKey(namespace, taskID)
-	_, err := c.client.PutObject(context.Background(), &s3.PutObjectInput{
+	_, err := c.client.PutObject(ctx, &s3.PutObjectInput{

And update the call site in OnTaskCompletionBarrierHandler:

-	if err := c.flushBatch(batch.namespace, taskID, batch.docs); err != nil {
+	if err := c.flushBatch(context.Background(), batch.namespace, taskID, batch.docs); err != nil {

(Or pass a proper context if available from the barrier handler)

---

468-480: Inefficient buffer pre-allocation.

The buf.Grow(len(docs) * 2) estimates only 2 bytes per document, which is far too small for typical JSON documents. This will cause multiple reallocations as the buffer grows. For example, 1000 documents of 100 bytes each would need ~100KB but only 2KB is pre-allocated.

Consider a more realistic estimate:

 func buildJSONArray(docs [][]byte) []byte {
 	var buf bytes.Buffer
-	buf.Grow(len(docs) * 2)
+	// Estimate: sum of document sizes + commas + brackets + 10% overhead
+	totalSize := 2 // '[' and ']'
+	for _, doc := range docs {
+		totalSize += len(doc) + 1 // +1 for comma (or close bracket)
+	}
+	buf.Grow(totalSize)
 	buf.WriteByte('[')

Or use a simple heuristic like len(docs) * 100 if you expect documents to average ~100 bytes.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

connectors/s3/connector.go (2)

290-345: ListData behavior is reasonable; BSON round‑trip goes through JSON and may not preserve exact BSON types.

The method validates the S3 key from the partition cursor, reads a JSON array, and:

• For DATA_TYPE_JSON_ID, returns each element as validated JSON bytes.
• For DATA_TYPE_MONGO_BSON, unmarshals each JSON element into map[string]any and then bson.Marshal’s it back to BSON.

This is functionally fine if S3 is primarily a JSON export format. However, the BSON branch does a JSON→generic‑map→BSON conversion, which can change certain BSON types on round‑trip (e.g., extended JSON representations, numeric precision nuances). This was already called out in a previous review; if exact BSON semantics matter, you may want to adopt an Extended JSON strategy instead or clearly document that the S3 connector is JSON-oriented and BSON support is best-effort.

Check the `go.mongodb.org/mongo-driver/bson` documentation for how `bson.Unmarshal` into `map[string]any` followed by `bson.Marshal` treats ObjectID, DateTime, and Decimal128 types when starting from JSON, and whether `MarshalExtJSON` is preferable for preserving BSON semantics.

---

512-532: convertToJSON validates JSON and handles BSON input; BSON path is lossy and double‑marshals.

convertToJSON:

• Validates raw JSON for DATA_TYPE_JSON_ID and returns a copy.
• For DATA_TYPE_MONGO_BSON, unmarshals BSON into map[string]any and re‑marshals to JSON, which is straightforward but does a full decode/encode.
• Returns ErrUnsupportedType for other data types, aligning with ListData.

As with ListData, the BSON path may lose some type fidelity relative to raw BSON, and it does a full double marshal; if performance or exact BSON preservation becomes important, you might want to revisit this (e.g., Extended JSON, or deferring conversion until needed).

Look up best practices for converting MongoDB BSON to JSON in Go using `go.mongodb.org/mongo-driver/bson` and whether `MarshalExtJSON` is recommended when round-tripping data between BSON and JSON.

internal/app/options/connectorflags.go (1)

596-652: S3Flags align with ConnectorSettings; flag semantics vs description could be clarified.

The flags correctly populate all fields on s3connector.ConnectorSettings (region, prefix, output-format, profile, endpoint, credentials, path-style, pretty-json). Two small nits:

• The prefix flag description says “Override or append…”, but the current implementation only overwrites settings.Prefix. If append semantics are desired (e.g. connection-string prefix + --prefix subdir), you’d need a bit of logic in the S3 Create action to combine them.
• pretty-json defaults to true via Value: true, which is good; just ensure this matches your intended default for programmatic callers that bypass CLI.

Both are behavioral clarifications rather than blockers; the current implementation is otherwise consistent.

connectors/s3/connector.go (5)

146-288: GeneratePlan / GetNamespaceMetadata and metadata caching are coherent; watch out for very large buckets.

The plan generation loop:

• Lists objects under the optional configured prefix.
• Filters to .json objects while excluding .metadata.json.
• Derives namespaces from the key path (falling back to "default"), applies optional namespace filters, and looks up per-file counts from a cached readMetadata result per namespace.

GetNamespaceMetadata reuses the same metadata format and simply sums counts, which is consistent.

Only potential concern is scalability: for buckets with very large object counts, doing a full ListObjectsV2 across the prefix to build the plan could be expensive. If you expect such buckets, consider an optional limit or a naming convention that lets you narrow the listing further.

---

391-438: Barrier handler and batch bookkeeping correctly avoid empty flushes; consider using request context for metadata.

OnTaskCompletionBarrierHandler cleanly:

• Detaches and discards missing/empty batches (avoiding creation of empty files).
• Flushes non-empty batches to S3, sets a sticky error on failure, and returns the error.
• Updates namespace metadata after a successful flush, logging but not failing the barrier on metadata errors.

appendBatch/detachBatch under batchesMutex ensure per-task batches are consistent.

One small improvement would be to thread a context into OnTaskCompletionBarrierHandler (or at least use a shared cancellable context instead of context.Background() inside updateMetadataAfterFlush), so metadata writes can respect shutdown/timeouts like the flush itself.

---

484-510: JSON array builder works; preallocation heuristic is very low and may cause extra allocations.

buildJSONArray correctly builds a JSON array from individual documents and optionally pretty‑prints each element with json.Indent, falling back gracefully and logging when indentation fails. The newline placement also produces valid JSON.

The only nit is buf.Grow(len(docs) * 2), which significantly underestimates the needed capacity for realistic document sizes and will lead to repeated reallocations for large batches. If you care about performance here, consider summing len(doc) across docs (plus a small overhead) and using that to initialize the buffer capacity.

---

534-626: Namespace metadata helpers are consistent; global mutex serializes updates safely but could be narrowed later.

metadataKey, readMetadata, writeMetadata, and updateMetadataAfterFlush together implement per-namespace metadata as a JSON map of filename -> recordCount:

• readMetadata gracefully treats NoSuchKey as “no metadata yet” and decodes the JSON map otherwise.
• writeMetadata writes the map back as JSON with appropriate content type.
• updateMetadataAfterFlush locks metadataMutex, performs a read-modify-write on the metadata for the file just flushed, and uses objectKey/path.Base to keep keys consistent with GeneratePlan.

Using a single metadataMutex on the connector is safe and avoids lost updates, matching past review guidance. If you later run into contention with many namespaces flushing concurrently, you could evolve this into a per-namespace lock map, but that’s not required for correctness.

---

37-120: Consider passing a cancellable context to LoadDefaultConfig instead of context.Background().

The AWS SDK v2 best practices recommend passing a caller-owned context (rather than context.Background()) to LoadDefaultConfig so the config loading can be cancelled or respect process-wide timeouts. This enables graceful shutdown and timeout handling during initialization. If constructor timeouts are not a concern, context.Background() is acceptable, but using a cancellable context (or one with a timeout) is preferred for production deployments with strict shutdown requirements.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

internal/app/options/connectorflags.go (1)

300-316: S3 connector registration looks correct, but help-by-name invocation issue remains.

The past review comment about help invocation (e.g., dsync S3 --help) causing an "invalid s3 connection string" error is still applicable. While the urfave/cli library typically handles --help before executing the action, the current implementation sets settings.Uri = args[0] before calling CreateHelper, which could cause issues if the action is invoked.

connectors/s3/connector.go (1)

316-371: ListData correctly implements data retrieval.

The implementation properly reads JSON arrays from S3 and converts them based on the requested data type. For DATA_TYPE_MONGO_BSON, the JSON-to-BSON conversion (lines 352-361) reconstructs BSON documents from the JSON stored in S3.

Note that BSON-specific types (ObjectID, DateTime, Decimal128, etc.) are stored as their JSON representations, not as extended JSON with type markers. This means round-tripping BSON → JSON (in WriteData) → BSON (in ListData) loses type fidelity. If this is intentional for an analytics/export use case, consider documenting this behavior in the package comment or GetInfo response.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

connectors/s3/connector.go (2)

440-445: Optional: Remove redundant check.

The condition if len(batch.docs) > 0 at line 440 is always true because lines 429-432 already return early when batch.docs is empty. This check can be safely removed.

 	}
 
 	// Update metadata file atomically after successful flush
-	if len(batch.docs) > 0 {
-		if err := c.updateMetadataAfterFlush(context.Background(), batch.namespace, taskID, uint64(len(batch.docs))); err != nil {
-			slog.Error("failed to update metadata", "namespace", batch.namespace, "taskId", taskID, "err", err)
-			// Log error but don't fail the barrier - the data was successfully flushed
-		}
+	if err := c.updateMetadataAfterFlush(context.Background(), batch.namespace, taskID, uint64(len(batch.docs))); err != nil {
+		slog.Error("failed to update metadata", "namespace", batch.namespace, "taskId", taskID, "err", err)
+		// Log error but don't fail the barrier - the data was successfully flushed
 	}

---

627-657: Recommended: Consider per-namespace mutex for better concurrency.

The global metadataMutex (line 630) serializes metadata updates across all namespaces. With many concurrent namespaces, this could become a bottleneck since updates to different namespaces don't actually conflict.

Consider using per-namespace locking for better concurrency:

 type connector struct {
 	adiomv1connect.UnimplementedConnectorServiceHandler
 
 	client       *s3.Client
 	settings     ConnectorSettings
 	batchesMutex sync.Mutex
 	batches      map[taskKey]*storedBatch
 
-	metadataMutex sync.Mutex // Serialize metadata updates
+	namespaceLocks sync.Map // map[string]*sync.Mutex for per-namespace locking
 
 	errMutex sync.RWMutex
 	err      error
 }

Then in updateMetadataAfterFlush:

 func (c *connector) updateMetadataAfterFlush(ctx context.Context, namespace string, taskID uint, recordCount uint64) error {
-	c.metadataMutex.Lock()
-	defer c.metadataMutex.Unlock()
+	// Get or create mutex for this namespace
+	lockVal, _ := c.namespaceLocks.LoadOrStore(namespace, &sync.Mutex{})
+	mu := lockVal.(*sync.Mutex)
+	mu.Lock()
+	defer mu.Unlock()

This allows concurrent metadata updates for different namespaces while still preventing races within the same namespace.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro