Support separate streams for mongo. - #314
Conversation
WalkthroughAdds a per-namespace streaming mode for MongoDB change streams via a new Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor User as Runner
participant Conn as Mongo Connector
participant Plan as GeneratePlan
participant Watch as Watchable
participant CS as ChangeStream
Note over Conn,Plan: Plan generation (updates)
User->>Conn: Start planning
Conn->>Plan: GeneratePlan(request)
alt PerNamespaceStreams = true
loop For each namespace
Plan->>Watch: Watch(ns pipeline)
Watch-->>Plan: ChangeStream
Plan->>CS: TryNext()
CS-->>Plan: ResumeToken (ns)
Plan-->>Conn: Add UpdatesPartition(ns, token)
end
else PerNamespaceStreams = false
Plan->>Watch: Watch(cluster pipeline)
Watch-->>Plan: ChangeStream
Plan->>CS: TryNext()
CS-->>Plan: ResumeToken (cluster)
Plan-->>Conn: One UpdatesPartition(namespaces, token)
end
Conn-->>User: Plan with UpdatesPartitions
sequenceDiagram
autonumber
participant Util as getLatestResumeToken
participant W as Watchable
participant CS as ChangeStream
Util->>W: Watch(pipeline, opts)
W-->>Util: CS
Util->>CS: TryNext(ctx)
alt Token available
CS-->>Util: ResumeToken
Util-->>Caller: token
else No token / error
CS-->>Util: Err()
Util-->>Caller: error (no resume token)
end
Note over Util,CS: ChangeStream closed via defer Close(ctx)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests
Tip 👮 Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Please see the documentation for more information. Example: reviews:
pre_merge_checks:
custom_checks:
- name: "Undocumented Breaking Changes"
mode: "warning"
instructions: |
Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
connectors/mongo/conn.go (1)
26-27: Duplicate import of the same package breaks compilation.You import go.mongodb.org/mongo-driver/mongo/options twice (once as options and once as moptions). Consolidate to a single import and update references.
- "go.mongodb.org/mongo-driver/mongo/options" - moptions "go.mongodb.org/mongo-driver/mongo/options" + "go.mongodb.org/mongo-driver/mongo/options"- opts := moptions.ChangeStream().SetStartAfter(bson.Raw(r.Msg.GetCursor())) + opts := options.ChangeStream().SetStartAfter(bson.Raw(r.Msg.GetCursor()))- opts := moptions.ChangeStream().SetStartAfter(bson.Raw(r.Msg.GetCursor())).SetFullDocument("updateLookup") + opts := options.ChangeStream().SetStartAfter(bson.Raw(r.Msg.GetCursor())).SetFullDocument("updateLookup")- clientOptions := moptions.Client().SetAppName("dsync").ApplyURI(settings.ConnectionString).SetConnectTimeout(settings.ServerConnectTimeout) + clientOptions := options.Client().SetAppName("dsync").ApplyURI(settings.ConnectionString).SetConnectTimeout(settings.ServerConnectTimeout)
🧹 Nitpick comments (4)
internal/app/options/connectorflags.go (1)
354-360: Flag wiring looks correct; consider grouping/help polish.Add a Category (e.g., "MongoDB-specific Options") and an env alias to ease discovery and config (altsrc supports both).
- altsrc.NewBoolFlag(&cli.BoolFlag{ + altsrc.NewBoolFlag(&cli.BoolFlag{ Name: "per-namespace-streams", Usage: "Each namespace has a separate stream", Destination: &settings.PerNamespaceStreams, + Category: "MongoDB-specific Options", + EnvVars: []string{"DSYNC_MONGO_PER_NAMESPACE_STREAMS"}, }),connectors/mongo/util.go (1)
62-74: Make resume-token acquisition robust and return a clear error when token is nil.TryNext may yield no error yet still no token; wrap with explicit nil check and avoid %w with a nil error.
func getLatestResumeToken(ctx context.Context, client Watchable) (bson.Raw, error) { slog.Debug("Getting latest resume token...") changeStream, err := client.Watch(ctx, mongo.Pipeline{}) if err != nil { return nil, fmt.Errorf("failed to open change stream: %v", err) } defer changeStream.Close(ctx) - _ = changeStream.TryNext(ctx) - if token := changeStream.ResumeToken(); token != nil { - return token, nil - } - return nil, fmt.Errorf("failed to get a resume token: %w", changeStream.Err()) + _ = changeStream.TryNext(ctx) + if token := changeStream.ResumeToken(); token != nil { + return token, nil + } + if err := changeStream.Err(); err != nil { + return nil, fmt.Errorf("failed to get a resume token: %w", err) + } + return nil, fmt.Errorf("failed to get a resume token: stream returned no token") }connectors/mongo/conn.go (2)
171-189: Per-namespace resume tokens fetched serially; parallelize with bounded concurrency.Fetching tokens per collection sequentially will be slow on large deployments.
- for _, partition := range partitions { - ns, _ := ToNS(partition.Namespace) - col := c.client.Database(ns.Db).Collection(ns.Col) - resumeToken, err := getLatestResumeToken(ctx, col) - if err != nil { ... } - updatesPartitions = append(updatesPartitions, &adiomv1.UpdatesPartition{ ... }) - } + g, gctx := errgroup.WithContext(ctx) + sem := make(chan struct{}, 8) // tune concurrency + mu := &sync.Mutex{} + for _, p := range partitions { + partition := p + g.Go(func() error { + sem <- struct{}{} + defer func(){ <-sem }() + ns, _ := ToNS(partition.Namespace) + col := c.client.Database(ns.Db).Collection(ns.Col) + resumeToken, err := getLatestResumeToken(gctx, col) + if err != nil { return fmt.Errorf("ns %s: %w", partition.GetNamespace(), err) } + mu.Lock() + updatesPartitions = append(updatesPartitions, &adiomv1.UpdatesPartition{ + Namespaces: []string{partition.GetNamespace()}, + Cursor: resumeToken, + }) + mu.Unlock() + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, connect.NewError(connect.CodeInternal, err) + }
658-665: Optional: watch collection when a single namespace is requested.When r.Msg.Namespaces has length 1, using collection.Watch reduces server-side filtering vs. client.Watch+$match.
- changeStream, err := c.client.Watch(ctx, mongo.Pipeline{ + watcher := any(c.client) + if len(namespaces) == 1 { + watcher = c.client.Database(namespaces[0].Db).Collection(namespaces[0].Col) + } + changeStream, err := watcher.(interface { + Watch(context.Context, interface{}, ...*options.ChangeStreamOptions) (*mongo.ChangeStream, error) + }).Watch(ctx, mongo.Pipeline{ {{"$match", nsFilter}}, }, opts)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
connectors/mongo/change_stream_handler.go(0 hunks)connectors/mongo/conn.go(4 hunks)connectors/mongo/util.go(2 hunks)internal/app/options/connectorflags.go(1 hunks)
💤 Files with no reviewable changes (1)
- connectors/mongo/change_stream_handler.go
🧰 Additional context used
🧬 Code graph analysis (1)
connectors/mongo/conn.go (2)
gen/adiom/v1/messages.pb.go (3)
UpdatesPartition(185-195)UpdatesPartition(208-208)UpdatesPartition(223-225)protocol/iface/connector.go (1)
Namespace(103-106)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: build
🔇 Additional comments (5)
connectors/mongo/util.go (3)
58-60: Good abstraction for testability.Exported Watchable cleanly decouples from concrete *mongo.Client/*mongo.Collection.
20-20: Import is appropriate.options import is needed for the Watchable signature.
53-56: Do not remove dummyDB/dummyCol — they are still referenced in planner.go.connectors/mongo/planner.go uses these constants (ExcludedDBListForIC and the change-stream filter — connectors/mongo/planner.go lines ~16 and ~32–33); delete only after updating/removing those usages or keep the constants.
Likely an incorrect or invalid review comment.
connectors/mongo/conn.go (2)
281-292: Capability DefaultPlan ties to PerNamespaceStreams: verify UX.DefaultPlan = !PerNamespaceStreams changes behavior surfaced to callers. Confirm docs/CLI help convey that enabling per-namespace streams disables the “single default plan”.
551-561: No action required — createChangeStreamNamespaceFilter exists.
Found at connectors/mongo/planner.go:38; keep the call as-is.
1450876 to
fdaa749
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
connectors/mongo/util.go (1)
58-60: Clarify Watchable’s purpose and scope.Add a short doc comment explaining Watchable is satisfied by *mongo.Client and *mongo.Collection and is used for resume-token acquisition; consider unexporting if it’s package-internal only.
connectors/mongo/conn.go (2)
41-41: Document PerNamespaceStreams behavior.Add a brief comment on how this toggles plan shape and stream scope so downstream users understand capability changes and defaulting.
171-190: Per-namespace resume-token fetch runs serially; parallelize with a limit.For large deployments (hundreds/thousands of collections), serial getLatestResumeToken calls will add noticeable latency. Use errgroup with a concurrency cap to fetch tokens in parallel.
Apply this refactor within the PerNamespaceStreams branch:
- // TODO: maybe parallelize - for _, partition := range partitions { - ns, _ := ToNS(partition.Namespace) - col := c.client.Database(ns.Db).Collection(ns.Col) - resumeToken, err := getLatestResumeToken(ctx, col) - if err != nil { - slog.Error(fmt.Sprintf("Failed to get latest resume token for ns %v: %v", partition.GetNamespace(), err)) - return nil, connect.NewError(connect.CodeInternal, err) - } - - updatesPartitions = append(updatesPartitions, &adiomv1.UpdatesPartition{ - Namespaces: []string{partition.GetNamespace()}, - Cursor: resumeToken, - }) - } + // Parallelize with bounded concurrency. + { + var mu sync.Mutex + eg, gctx := errgroup.WithContext(ctx) + // tune as needed + eg.SetLimit(16) + for _, p := range partitions { + part := p + eg.Go(func() error { + ns, _ := ToNS(part.Namespace) + col := c.client.Database(ns.Db).Collection(ns.Col) + resumeToken, err := getLatestResumeToken(gctx, col) + if err != nil { + slog.Error(fmt.Sprintf("Failed to get latest resume token for ns %v: %v", part.GetNamespace(), err)) + return err + } + mu.Lock() + updatesPartitions = append(updatesPartitions, &adiomv1.UpdatesPartition{ + Namespaces: []string{part.GetNamespace()}, + Cursor: resumeToken, + }) + mu.Unlock() + return nil + }) + } + if err := eg.Wait(); err != nil { + return nil, connect.NewError(connect.CodeInternal, err) + } + }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
connectors/mongo/change_stream_handler.go(0 hunks)connectors/mongo/conn.go(4 hunks)connectors/mongo/util.go(2 hunks)internal/app/options/connectorflags.go(1 hunks)
💤 Files with no reviewable changes (1)
- connectors/mongo/change_stream_handler.go
🚧 Files skipped from review as they are similar to previous changes (1)
- internal/app/options/connectorflags.go
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: build
🔇 Additional comments (3)
connectors/mongo/conn.go (3)
191-204: Confirm semantics when no namespaces are requested.When r.Msg.GetNamespaces() is empty, you return a single UpdatesPartition with an empty Namespaces list. Verify the runner interprets “empty” as “all namespaces,” not “none.” If ambiguity exists, explicitly omit the field or add a flag.
267-267: LGTM: UpdatesPartitions is now correctly derived from mode.Response wiring matches the newly computed updatesPartitions slice.
292-292: Capability DefaultPlan toggled by PerNamespaceStreams.Looks right. Please ensure client UX reflects the new default when the flag is set.
| func getLatestResumeToken(ctx context.Context, client Watchable) (bson.Raw, error) { | ||
| slog.Debug("Getting latest resume token...") | ||
| changeStream, err := client.Watch(ctx, mongo.Pipeline{}) //TODO (AK, 6/2024): We should limit this to just the dummy collection or we can catch something that we don't want :) | ||
| changeStream, err := client.Watch(ctx, mongo.Pipeline{}) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to open change stream: %v", err) | ||
| } | ||
| defer changeStream.Close(ctx) | ||
|
|
||
| done := make(chan struct{}) | ||
| // we need ANY event to get the resume token that we can use to extract the cluster time | ||
| go func() { | ||
| if err := insertDummyRecord(ctx, client); err != nil { | ||
| slog.Error(fmt.Sprintf("Error inserting dummy record: %v", err.Error())) | ||
| } | ||
| close(done) | ||
| }() | ||
|
|
||
| changeStream.Next(ctx) | ||
| resumeToken := changeStream.ResumeToken() | ||
| if resumeToken == nil { | ||
| return nil, fmt.Errorf("failed to get resume token from change stream") | ||
| _ = changeStream.TryNext(ctx) | ||
| if token := changeStream.ResumeToken(); token != nil { | ||
| return token, nil | ||
| } | ||
|
|
||
| <-done | ||
| return resumeToken, nil | ||
| if changeStream.Err() != nil { | ||
| return nil, fmt.Errorf("failed to get a resume token: %w", changeStream.Err()) | ||
| } | ||
| return nil, fmt.Errorf("failed to get a resume token") | ||
| } |
There was a problem hiding this comment.
Resume-token retrieval can spuriously fail; use Next with a bounded wait and MaxAwaitTime instead of a single TryNext.
Calling TryNext once often returns before the driver performs a server round-trip, so ResumeToken may remain nil even though a postBatchResumeToken is available. This will intermittently fail plan generation. Use Next in a short, bounded loop and set ChangeStream MaxAwaitTime to force the server to return an empty batch (with PBRT) if there are no changes.
Apply these diffs:
Imports (add errors/time):
-import (
+import (
"context"
"fmt"
"log/slog"
"regexp"
"strconv"
"strings"
+ "errors"
+ "time"Function:
-func getLatestResumeToken(ctx context.Context, client Watchable) (bson.Raw, error) {
- slog.Debug("Getting latest resume token...")
- changeStream, err := client.Watch(ctx, mongo.Pipeline{})
- if err != nil {
- return nil, fmt.Errorf("failed to open change stream: %v", err)
- }
- defer changeStream.Close(ctx)
- _ = changeStream.TryNext(ctx)
- if token := changeStream.ResumeToken(); token != nil {
- return token, nil
- }
- if changeStream.Err() != nil {
- return nil, fmt.Errorf("failed to get a resume token: %w", changeStream.Err())
- }
- return nil, fmt.Errorf("failed to get a resume token")
-}
+func getLatestResumeToken(ctx context.Context, client Watchable) (bson.Raw, error) {
+ slog.Debug("Getting latest resume token...")
+ // Bound the wait and force server round-trips to obtain postBatchResumeToken.
+ waitCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
+ defer cancel()
+
+ cs, err := client.Watch(waitCtx, mongo.Pipeline{}, options.ChangeStream().SetMaxAwaitTime(2*time.Second))
+ if err != nil {
+ return nil, fmt.Errorf("failed to open change stream: %w", err)
+ }
+ defer cs.Close(waitCtx)
+
+ for {
+ _ = cs.Next(waitCtx) // blocks until event or empty batch arrives (bounded by MaxAwaitTime)
+ if tok := cs.ResumeToken(); tok != nil {
+ return tok, nil
+ }
+ if err := cs.Err(); err != nil {
+ // Treat context expiry separately to return a clearer message.
+ if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
+ break
+ }
+ return nil, fmt.Errorf("failed to get a resume token: %w", err)
+ }
+ if waitCtx.Err() != nil {
+ break
+ }
+ }
+ return nil, fmt.Errorf("failed to get a resume token: timed out waiting for postBatchResumeToken")
+}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| func getLatestResumeToken(ctx context.Context, client Watchable) (bson.Raw, error) { | |
| slog.Debug("Getting latest resume token...") | |
| changeStream, err := client.Watch(ctx, mongo.Pipeline{}) //TODO (AK, 6/2024): We should limit this to just the dummy collection or we can catch something that we don't want :) | |
| changeStream, err := client.Watch(ctx, mongo.Pipeline{}) | |
| if err != nil { | |
| return nil, fmt.Errorf("failed to open change stream: %v", err) | |
| } | |
| defer changeStream.Close(ctx) | |
| done := make(chan struct{}) | |
| // we need ANY event to get the resume token that we can use to extract the cluster time | |
| go func() { | |
| if err := insertDummyRecord(ctx, client); err != nil { | |
| slog.Error(fmt.Sprintf("Error inserting dummy record: %v", err.Error())) | |
| } | |
| close(done) | |
| }() | |
| changeStream.Next(ctx) | |
| resumeToken := changeStream.ResumeToken() | |
| if resumeToken == nil { | |
| return nil, fmt.Errorf("failed to get resume token from change stream") | |
| _ = changeStream.TryNext(ctx) | |
| if token := changeStream.ResumeToken(); token != nil { | |
| return token, nil | |
| } | |
| <-done | |
| return resumeToken, nil | |
| if changeStream.Err() != nil { | |
| return nil, fmt.Errorf("failed to get a resume token: %w", changeStream.Err()) | |
| } | |
| return nil, fmt.Errorf("failed to get a resume token") | |
| } | |
| import ( | |
| "context" | |
| "fmt" | |
| "log/slog" | |
| "regexp" | |
| "strconv" | |
| "strings" | |
| "errors" | |
| "time" | |
| ) | |
| func getLatestResumeToken(ctx context.Context, client Watchable) (bson.Raw, error) { | |
| slog.Debug("Getting latest resume token...") | |
| // Bound the wait and force server round-trips to obtain postBatchResumeToken. | |
| waitCtx, cancel := context.WithTimeout(ctx, 5*time.Second) | |
| defer cancel() | |
| cs, err := client.Watch(waitCtx, mongo.Pipeline{}, options.ChangeStream().SetMaxAwaitTime(2*time.Second)) | |
| if err != nil { | |
| return nil, fmt.Errorf("failed to open change stream: %w", err) | |
| } | |
| defer cs.Close(waitCtx) | |
| for { | |
| _ = cs.Next(waitCtx) // blocks until event or empty batch arrives (bounded by MaxAwaitTime) | |
| if tok := cs.ResumeToken(); tok != nil { | |
| return tok, nil | |
| } | |
| if err := cs.Err(); err != nil { | |
| // Treat context expiry separately to return a clearer message. | |
| if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { | |
| break | |
| } | |
| return nil, fmt.Errorf("failed to get a resume token: %w", err) | |
| } | |
| if waitCtx.Err() != nil { | |
| break | |
| } | |
| } | |
| return nil, fmt.Errorf("failed to get a resume token: timed out waiting for postBatchResumeToken") | |
| } |
1 participant
Summary by CodeRabbit
New Features
Bug Fixes
Refactor