Skip to content

update mongo driver #376

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
adiom-mark merged 3 commits into from
Mar 31, 2026
Merged

update mongo driver #376

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
56d0f76
update mongo driver
Mar 30, 2026
6c88896
cosmos fixes
Mar 30, 2026
77112c6
small fix
Mar 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion connectors/airbyte/helpers.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ import (
"strings"

adiomv1 "github.com/adiom-data/dsync/gen/adiom/v1"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/v2/bson"
"golang.org/x/sync/errgroup"
)

Expand Down
15 changes: 7 additions & 8 deletions connectors/common/base.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,7 @@ import (
"github.com/cespare/xxhash"
"github.com/google/uuid"
"go.akshayshah.org/memhttp"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/bson/primitive"
"go.mongodb.org/mongo-driver/v2/bson"
"golang.org/x/time/rate"
)

Expand Down Expand Up @@ -102,11 +101,11 @@ func (c *connector) GetConnectorStatus(flowId iface.FlowID) iface.ConnectorStatu

func IDPartToString(a any) string {
switch t := a.(type) {
case primitive.ObjectID:
case bson.ObjectID:
return t.Hex()
case string:
return t
case primitive.Binary:
case bson.Binary:
if t.Subtype == bson.TypeBinaryUUID {
id, err := uuid.FromBytes(t.Data)
if err == nil {
Expand Down Expand Up @@ -181,7 +180,7 @@ func HashBson(hasher hash.Hash64, b bson.Raw, arr bool, projection map[string]in
if _, err := hasher.Write([]byte(e.Key())); err != nil {
return err
}
if err := HashBson(hasher, v.Array(), true, innerProjection); err != nil {
if err := HashBson(hasher, bson.Raw(v.Array()), true, innerProjection); err != nil {
return err
}
} else {
Expand Down Expand Up @@ -212,7 +211,7 @@ func (c *connector) IntegrityCheck(ctx context.Context, task iface.IntegrityChec

var pCursor []byte
if task.Low != nil {
pCursor = task.Low.(primitive.Binary).Data
pCursor = task.Low.(bson.Binary).Data
}
namespace := task.Namespace
for {
Expand Down Expand Up @@ -328,7 +327,7 @@ func (c *connector) RequestCreateReadPlan(flowId iface.FlowID, options iface.Con
Id: iface.ReadPlanTaskID(curID),
}
task.Def.Col = partition.GetNamespace()
task.Def.Low = primitive.Binary{Data: partition.GetCursor()}
task.Def.Low = bson.Binary{Data: partition.GetCursor()}
task.EstimatedDocCount = int64(partition.GetEstimatedCount())
tasks = append(tasks, task)
}
Expand Down Expand Up @@ -570,7 +569,7 @@ func (c *connector) StartReadToChannel(flowId iface.FlowID, options iface.Connec
var cursor []byte
var pCursor []byte
if task.Def.Low != nil {
pCursor = task.Def.Low.(primitive.Binary).Data
pCursor = task.Def.Low.(bson.Binary).Data
}
for {
res, err := c.maybeOptimizedImpl.ListData(c.flowCtx, connect.NewRequest(&adiomv1.ListDataRequest{
Expand Down
2 changes: 1 addition & 1 deletion connectors/common/hash_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import (
"github.com/adiom-data/dsync/connectors/common"
"github.com/cespare/xxhash"
"github.com/stretchr/testify/assert"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/v2/bson"
)

func compareHash(t *testing.T, hasher hash.Hash64, left []byte, right []byte, projection map[string]interface{}) bool {
Expand Down
8 changes: 4 additions & 4 deletions connectors/cosmos/change_streams.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,9 @@ import (
"context"

"github.com/adiom-data/dsync/protocol/iface"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo"
moptions "go.mongodb.org/mongo-driver/mongo/options"
"go.mongodb.org/mongo-driver/v2/bson"
"go.mongodb.org/mongo-driver/v2/mongo"
moptions "go.mongodb.org/mongo-driver/v2/mongo/options"
)

var streamPipeline = mongo.Pipeline{
Expand All @@ -27,7 +27,7 @@ var streamPipelineWithDelete = mongo.Pipeline{
}

// Creates a single changestream compatible with CosmosDB with the provided options
func createChangeStream(ctx context.Context, client *mongo.Client, namespace iface.Location, opts *moptions.ChangeStreamOptions, withDelete bool) (*mongo.ChangeStream, error) {
func createChangeStream(ctx context.Context, client *mongo.Client, namespace iface.Location, opts moptions.Lister[moptions.ChangeStreamOptions], withDelete bool) (*mongo.ChangeStream, error) {
db := namespace.Database
col := namespace.Collection
collection := client.Database(db).Collection(col)
Expand Down
9 changes: 4 additions & 5 deletions connectors/cosmos/check_deletes.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,9 @@ import (
"sync"

"github.com/adiom-data/dsync/protocol/iface"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/bson/primitive"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
"go.mongodb.org/mongo-driver/v2/bson"
"go.mongodb.org/mongo-driver/v2/mongo"
"go.mongodb.org/mongo-driver/v2/mongo/options"
)

const (
Expand Down Expand Up @@ -179,7 +178,7 @@ func checkSourceIdsAndGenerateDeletesWorker(ctx context.Context, client *mongo.C
}

// convert result to array
missingIds = []interface{}(res["missingIds"].(primitive.A))
missingIds = []interface{}(res["missingIds"].(bson.A))
} else {
slog.Debug(fmt.Sprintf("Missing ids source query returned nothing for %v", idsWithLoc.loc))
//this means all are missing on the source!
Expand Down
63 changes: 28 additions & 35 deletions connectors/cosmos/conn.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,9 @@ import (
adiomv1 "github.com/adiom-data/dsync/gen/adiom/v1"
"github.com/adiom-data/dsync/gen/adiom/v1/adiomv1connect"
"github.com/adiom-data/dsync/protocol/iface"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/bson/primitive"
"go.mongodb.org/mongo-driver/mongo"
moptions "go.mongodb.org/mongo-driver/mongo/options"
"go.mongodb.org/mongo-driver/v2/bson"
"go.mongodb.org/mongo-driver/v2/mongo"
moptions "go.mongodb.org/mongo-driver/v2/mongo/options"
)

const (
Expand Down Expand Up @@ -64,7 +63,7 @@ func encodeResumeToken(epoch int64, token []byte) []byte {
func decodeResumeToken(input []byte) (int64, []byte) {
var s bson.M
_ = bson.Unmarshal(input, &s)
return s["a"].(int64), s["b"].(primitive.Binary).Data
return s["a"].(int64), s["b"].(bson.Binary).Data
}

// GeneratePlan implements adiomv1connect.ConnectorServiceHandler.
Expand Down Expand Up @@ -209,13 +208,13 @@ func (c *conn) StreamLSN(ctx context.Context, r *connect.Request[adiomv1.StreamL
if err != nil {
slog.Error(fmt.Sprintf("Failed to get resume token for location %v: %v", loc, err))
}
var opts *moptions.ChangeStreamOptions
var opts *moptions.ChangeStreamOptionsBuilder
if token != nil {
//set the change stream options to start from the resume token
opts = moptions.ChangeStream().SetResumeAfter(token).SetFullDocument(moptions.UpdateLookup)
} else { //we need to start from the read plan creation time to be safe
// create timestamp from read plan start time
ts := primitive.Timestamp{T: uint32(readPlanStartAt)}
ts := bson.Timestamp{T: uint32(readPlanStartAt)}
slog.Debug(fmt.Sprintf("Starting change stream for %v at timestamp %v", ns, ts))
opts = moptions.ChangeStream().SetStartAtOperationTime(&ts).SetFullDocument(moptions.UpdateLookup)
}
Expand All @@ -231,12 +230,6 @@ func (c *conn) StreamLSN(ctx context.Context, r *connect.Request[adiomv1.StreamL
defer changeStream.Close(ctx)

for changeStream.Next(ctx) {
var change bson.M
if err := changeStream.Decode(&change); err != nil {
slog.Error(fmt.Sprintf("Failed to decode change stream event: %v", err))
continue
}

lsnTracker.IncrementLSN(ns)

if changeStream.RemainingBatchLength() == 0 {
Expand All @@ -262,19 +255,27 @@ func (c *conn) StreamLSN(ctx context.Context, r *connect.Request[adiomv1.StreamL
return nil
}

func convertChangeStreamEventToUpdate(change bson.M) (*adiomv1.Update, error) {
func convertChangeStreamEventToUpdate(change mongoconn.MongoUpdate) (*adiomv1.Update, error) {
//slog.Debug(fmt.Sprintf("Converting change stream event %v", change))

// treat all change stream events as updates
// get the id of the document that was changed
id := change["documentKey"].(bson.M)["_id"]
// convert id to raw bson
idType, idVal, err := bson.MarshalValue(id)
if err != nil {
return nil, fmt.Errorf("failed to marshal _id: %v", err)
var idType bson.Type
var idVal []byte
for _, k := range change.DocumentKey {
if k.Key == "_id" {
var err error
idType, idVal, err = bson.MarshalValue(k.Value)
if err != nil {
return nil, fmt.Errorf("failed to marshal _id: %v", err)
}
break
}
}
if idVal == nil {
return nil, fmt.Errorf("_id not found in documentKey")
}

if change["operationType"] == "delete" {
if change.OperationType == "delete" {
return &adiomv1.Update{
Id: []*adiomv1.BsonValue{{
Data: idVal,
Expand All @@ -286,24 +287,18 @@ func convertChangeStreamEventToUpdate(change bson.M) (*adiomv1.Update, error) {
}

// get the full state of the document after the change
if change["fullDocument"] == nil {
if len(change.FullDocument) == 0 {
//TODO (AK, 6/2024): find a better way to report that we need to ignore this event
return nil, nil // no full document, nothing to do (probably got deleted before we got to the event in the change stream)
}
fullDocument := change["fullDocument"].(bson.M)
// convert fulldocument to BSON.Raw
fullDocumentRaw, err := bson.Marshal(fullDocument)
if err != nil {
return nil, fmt.Errorf("failed to marshal full document: %v", err)
}
update := &adiomv1.Update{
Id: []*adiomv1.BsonValue{{
Data: idVal,
Type: uint32(idType),
Name: "_id",
}},
Type: adiomv1.UpdateType_UPDATE_TYPE_UPDATE,
Data: fullDocumentRaw,
Data: change.FullDocument,
}
return update, nil
}
Expand Down Expand Up @@ -430,13 +425,13 @@ func (c *conn) StreamUpdates(ctx context.Context, r *connect.Request[adiomv1.Str
if err != nil {
slog.Error(fmt.Sprintf("Failed to get resume token for location %v: %v", loc, err))
}
var opts *moptions.ChangeStreamOptions
var opts *moptions.ChangeStreamOptionsBuilder
if token != nil {
//set the change stream options to start from the resume token
opts = moptions.ChangeStream().SetResumeAfter(token).SetFullDocument(moptions.UpdateLookup)
} else { //we need to start from the read plan creation time to be safe
// create timestamp from read plan start time
ts := primitive.Timestamp{T: uint32(readPlanStartAt)}
ts := bson.Timestamp{T: uint32(readPlanStartAt)}
slog.Debug(fmt.Sprintf("Starting change stream for %v at timestamp %v", ns, ts))
opts = moptions.ChangeStream().SetStartAtOperationTime(&ts).SetFullDocument(moptions.UpdateLookup)
}
Expand All @@ -454,7 +449,7 @@ func (c *conn) StreamUpdates(ctx context.Context, r *connect.Request[adiomv1.Str
var updates []*adiomv1.Update

for changeStream.Next(ctx) {
var change bson.M
var change mongoconn.MongoUpdate
if err := changeStream.Decode(&change); err != nil {
slog.Error(fmt.Sprintf("Failed to decode change stream event: %v", err))
continue
Expand Down Expand Up @@ -528,10 +523,8 @@ func NewConn(settings ConnectorSettings) adiomv1connect.ConnectorServiceHandler
var witnessMongoClient *mongo.Client
// Connect to the witness MongoDB instance
if settings.EmulateDeletes {
ctxConnect, cancel := context.WithTimeout(context.Background(), settings.ServerConnectTimeout)
defer cancel()
clientOptions := moptions.Client().ApplyURI(settings.WitnessMongoConnString).SetConnectTimeout(settings.ServerConnectTimeout)
client, err := mongo.Connect(ctxConnect, clientOptions)
client, err := mongo.Connect(clientOptions)
if err != nil {
panic(err)
}
Expand Down
11 changes: 5 additions & 6 deletions connectors/cosmos/connector_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,10 +27,9 @@ import (
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"github.com/tryvium-travels/memongo"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/bson/bsontype"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
"go.mongodb.org/mongo-driver/v2/bson"
"go.mongodb.org/mongo-driver/v2/mongo"
"go.mongodb.org/mongo-driver/v2/mongo/options"
)

const (
Expand Down Expand Up @@ -120,7 +119,7 @@ func TestCosmosConnectorSuite2(t *testing.T) {

tSuite.AssertExists = func(ctx context.Context, a *assert.Assertions, id []*adiomv1.BsonValue, exists bool) error {
mongoID := bson.RawValue{
Type: bsontype.Type(id[0].GetType()),
Type: bson.Type(id[0].GetType()),
Value: id[0].GetData(),
}
idFilter := bson.D{{Key: "_id", Value: mongoID}}
Expand Down Expand Up @@ -460,7 +459,7 @@ type CosmosTestDataStore struct {
func (c *CosmosTestDataStore) Setup() error {
// connect to the underlying database
clientOptions := options.Client().ApplyURI(c.ConnectionString)
client, err := mongo.Connect(context.TODO(), clientOptions)
client, err := mongo.Connect(clientOptions)
if err != nil {
return err
}
Expand Down
2 changes: 1 addition & 1 deletion connectors/cosmos/planner.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import (
"sync"

"github.com/adiom-data/dsync/protocol/iface"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/v2/mongo"
"golang.org/x/exp/rand"
)

Expand Down
9 changes: 4 additions & 5 deletions connectors/cosmos/range_partitioning.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,9 @@ import (
"time"

"github.com/adiom-data/dsync/protocol/iface"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/bson/primitive"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
"go.mongodb.org/mongo-driver/v2/bson"
"go.mongodb.org/mongo-driver/v2/mongo"
"go.mongodb.org/mongo-driver/v2/mongo/options"
)

// split the range between two boundaries into equal parts
Expand Down Expand Up @@ -65,7 +64,7 @@ func splitRangeObjectId(value1 bson.RawValue, value2 bson.RawValue, numParts int
boundaries := make([]bson.RawValue, numParts+1)
for i := 1; i < numParts; i++ {
t := minT.Add(time.Duration(i) * partSize)
val := primitive.NewObjectIDFromTimestamp(t)
val := bson.NewObjectIDFromTimestamp(t)
boundaries[i] = bson.RawValue{
Type: bson.TypeObjectID,
Value: val[:],
Expand Down
2 changes: 1 addition & 1 deletion connectors/cosmos/token_map.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ import (
"sync"

"github.com/adiom-data/dsync/protocol/iface"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/v2/bson"
)

type TokenMap struct {
Expand Down
6 changes: 3 additions & 3 deletions connectors/cosmos/util.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,9 @@ import (

"github.com/adiom-data/dsync/protocol/iface"
"github.com/mitchellh/hashstructure"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo"
moptions "go.mongodb.org/mongo-driver/mongo/options"
"go.mongodb.org/mongo-driver/v2/bson"
"go.mongodb.org/mongo-driver/v2/mongo"
moptions "go.mongodb.org/mongo-driver/v2/mongo/options"
)

const (
Expand Down
Loading
Loading