diff --git a/connectors/mongo/docdb.go b/connectors/mongo/docdb.go index 3602fb9..a2a1ed8 100644 --- a/connectors/mongo/docdb.go +++ b/connectors/mongo/docdb.go @@ -6,6 +6,8 @@ import ( "encoding/binary" "encoding/hex" "fmt" + "math" + "math/big" "sort" "context" @@ -16,19 +18,45 @@ import ( ) var supportedIDTypes = map[bson.Type]bool{ - bson.TypeObjectID: true, - bson.TypeString: true, - bson.TypeInt32: true, - bson.TypeInt64: true, - bson.TypeBinary: true, + bson.TypeObjectID: true, + bson.TypeString: true, + bson.TypeInt32: true, + bson.TypeInt64: true, + bson.TypeBinary: true, + bson.TypeEmbeddedDocument: true, } func compareBSONRawValues(a, b bson.RawValue) int { + if c := cmp.Compare(bsonTypeSortOrder(a.Type), bsonTypeSortOrder(b.Type)); c != 0 { + return c + } + + switch bsonTypeSortOrder(a.Type) { + case bsonTypeOrderNumber: + if c, ok := compareBSONRawNumbers(a, b); ok { + return c + } + case bsonTypeOrderString: + aString, aOK := bsonRawString(a) + bString, bOK := bsonRawString(b) + if aOK && bOK { + return cmp.Compare(aString, bString) + } + } + switch a.Type { + case bson.TypeDouble: + ai := math.Float64frombits(binary.LittleEndian.Uint64(a.Value)) + bi := math.Float64frombits(binary.LittleEndian.Uint64(b.Value)) + return cmp.Compare(ai, bi) case bson.TypeObjectID: return bytes.Compare(a.Value, b.Value) case bson.TypeString: return bytes.Compare(a.Value[4:len(a.Value)-1], b.Value[4:len(b.Value)-1]) + case bson.TypeEmbeddedDocument: + return compareBSONRawDocuments(bson.Raw(a.Value), bson.Raw(b.Value)) + case bson.TypeArray: + return compareBSONRawArrays(bson.RawArray(a.Value), bson.RawArray(b.Value)) case bson.TypeInt32: ai := int32(binary.LittleEndian.Uint32(a.Value)) bi := int32(binary.LittleEndian.Uint32(b.Value)) @@ -38,9 +66,233 @@ func compareBSONRawValues(a, b bson.RawValue) int { bi := int64(binary.LittleEndian.Uint64(b.Value)) return cmp.Compare(ai, bi) case bson.TypeBinary: - return bytes.Compare(a.Value[5:], b.Value[5:]) + if c, ok := compareBSONRawBinary(a, b); ok { + return c + } + case bson.TypeBoolean: + return cmp.Compare(a.Value[0], b.Value[0]) + case bson.TypeDateTime: + ai := int64(binary.LittleEndian.Uint64(a.Value)) + bi := int64(binary.LittleEndian.Uint64(b.Value)) + return cmp.Compare(ai, bi) + case bson.TypeTimestamp: + at, ai := a.Timestamp() + bt, bi := b.Timestamp() + if c := cmp.Compare(at, bt); c != 0 { + return c + } + return cmp.Compare(ai, bi) + case bson.TypeRegex: + ap, ao := a.Regex() + bp, bo := b.Regex() + if c := cmp.Compare(ap, bp); c != 0 { + return c + } + return cmp.Compare(ao, bo) + case bson.TypeJavaScript: + return cmp.Compare(a.JavaScript(), b.JavaScript()) + case bson.TypeCodeWithScope: + ac, as := a.CodeWithScope() + bc, bs := b.CodeWithScope() + if c := cmp.Compare(ac, bc); c != 0 { + return c + } + return compareBSONRawDocuments(as, bs) + case bson.TypeNull, bson.TypeMinKey, bson.TypeMaxKey: + return 0 + } + return bytes.Compare(a.Value, b.Value) +} + +func compareBSONRawNumbers(a, b bson.RawValue) (int, bool) { + aRank, ar, ok := bsonRawNumber(a) + if !ok { + return 0, false + } + bRank, br, ok := bsonRawNumber(b) + if !ok { + return 0, false + } + if c := cmp.Compare(aRank, bRank); c != 0 { + return c, true + } + if ar == nil || br == nil { + return 0, true + } + return ar.Cmp(br), true +} + +func bsonRawNumber(v bson.RawValue) (int, *big.Rat, bool) { + switch v.Type { + case bson.TypeDouble: + f := math.Float64frombits(binary.LittleEndian.Uint64(v.Value)) + if math.IsNaN(f) { + return 0, nil, true + } + if math.IsInf(f, -1) { + return 1, nil, true + } + if math.IsInf(f, 1) { + return 3, nil, true + } + r := new(big.Rat) + if r.SetFloat64(f) == nil { + return 0, nil, false + } + return 2, r, true + case bson.TypeInt32: + i := int32(binary.LittleEndian.Uint32(v.Value)) + return 2, new(big.Rat).SetInt64(int64(i)), true + case bson.TypeInt64: + i := int64(binary.LittleEndian.Uint64(v.Value)) + return 2, new(big.Rat).SetInt64(i), true + case bson.TypeDecimal128: + d := v.Decimal128() + if d.IsNaN() { + return 0, nil, true + } + switch d.IsInf() { + case -1: + return 1, nil, true + case 1: + return 3, nil, true + } + r, ok := decimal128Rat(d) + return 2, r, ok + default: + return 0, nil, false + } +} + +func decimal128Rat(d bson.Decimal128) (*big.Rat, bool) { + bi, exp, err := d.BigInt() + if err != nil { + return nil, false + } + r := new(big.Rat).SetInt(bi) + if exp == 0 { + return r, true + } + scale := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(abs(exp))), nil) + if exp > 0 { + return r.Mul(r, new(big.Rat).SetInt(scale)), true + } + return r.Quo(r, new(big.Rat).SetInt(scale)), true +} + +func abs(i int) int { + if i < 0 { + return -i + } + return i +} + +func bsonRawString(v bson.RawValue) (string, bool) { + switch v.Type { + case bson.TypeString: + s, ok := v.StringValueOK() + return s, ok + case bson.TypeSymbol: + s, ok := v.SymbolOK() + return s, ok + default: + return "", false + } +} + +func compareBSONRawBinary(a, b bson.RawValue) (int, bool) { + if len(a.Value) < 5 || len(b.Value) < 5 { + return 0, false + } + aLen := int32(binary.LittleEndian.Uint32(a.Value[:4])) + bLen := int32(binary.LittleEndian.Uint32(b.Value[:4])) + if c := cmp.Compare(aLen, bLen); c != 0 { + return c, true + } + if c := cmp.Compare(a.Value[4], b.Value[4]); c != 0 { + return c, true + } + return bytes.Compare(a.Value[5:], b.Value[5:]), true +} + +func compareBSONRawDocuments(a, b bson.Raw) int { + aElements, aErr := a.Elements() + bElements, bErr := b.Elements() + if aErr != nil || bErr != nil { + return bytes.Compare(a, b) + } + + for i := 0; i < min(len(aElements), len(bElements)); i++ { + aValue := aElements[i].Value() + bValue := bElements[i].Value() + if c := cmp.Compare(bsonTypeSortOrder(aValue.Type), bsonTypeSortOrder(bValue.Type)); c != 0 { + return c + } + if c := cmp.Compare(aElements[i].Key(), bElements[i].Key()); c != 0 { + return c + } + if c := compareBSONRawValues(aValue, bValue); c != 0 { + return c + } + } + return cmp.Compare(len(aElements), len(bElements)) +} + +func compareBSONRawArrays(a, b bson.RawArray) int { + aValues, aErr := a.Values() + bValues, bErr := b.Values() + if aErr != nil || bErr != nil { + return bytes.Compare(a, b) + } + + for i := 0; i < min(len(aValues), len(bValues)); i++ { + if c := compareBSONRawValues(aValues[i], bValues[i]); c != 0 { + return c + } + } + return cmp.Compare(len(aValues), len(bValues)) +} + +const ( + bsonTypeOrderNumber = 3 + bsonTypeOrderString = 4 +) + +func bsonTypeSortOrder(t bson.Type) int { + switch t { + case bson.TypeMinKey: + return 1 + case bson.TypeNull, bson.TypeUndefined: + return 2 + case bson.TypeInt32, bson.TypeInt64, bson.TypeDouble, bson.TypeDecimal128: + return 3 + case bson.TypeSymbol, bson.TypeString: + return 4 + case bson.TypeEmbeddedDocument: + return 5 + case bson.TypeArray: + return 6 + case bson.TypeBinary: + return 7 + case bson.TypeObjectID: + return 8 + case bson.TypeBoolean: + return 9 + case bson.TypeDateTime: + return 10 + case bson.TypeTimestamp: + return 11 + case bson.TypeRegex: + return 12 + case bson.TypeJavaScript: + return 13 + case bson.TypeCodeWithScope: + return 14 + case bson.TypeMaxKey: + return 15 + default: + return 14 } - panic("compareBSONRawValues called with unsupported type") } func (c *conn) sampleIDs(ctx context.Context, col *mongo.Collection, numSamples int64) ([]bson.RawValue, error) { diff --git a/connectors/mongo/docdb_sort_external_test.go b/connectors/mongo/docdb_sort_external_test.go new file mode 100644 index 0000000..4d2d429 --- /dev/null +++ b/connectors/mongo/docdb_sort_external_test.go @@ -0,0 +1,122 @@ +package mongo + +import ( + "context" + "math" + "os" + "sort" + "testing" + "time" + + "github.com/stretchr/testify/require" + "go.mongodb.org/mongo-driver/v2/bson" + "go.mongodb.org/mongo-driver/v2/mongo" + "go.mongodb.org/mongo-driver/v2/mongo/options" +) + +func TestMongoServerEmbeddedIDSortMatchesLocalComparator(t *testing.T) { + uri := os.Getenv("MONGO_SORT_TEST_URI") + if uri == "" { + t.Skip("set MONGO_SORT_TEST_URI to run") + } + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + client, err := mongo.Connect(options.Client().ApplyURI(uri).SetServerSelectionTimeout(10 * time.Second)) + require.NoError(t, err) + defer client.Disconnect(context.Background()) + require.NoError(t, client.Ping(ctx, nil)) + + col := client.Database("dsync_sort_test").Collection("embedded_ids_" + bson.NewObjectID().Hex()) + defer func() { + require.NoError(t, col.Drop(context.Background())) + }() + + decimal25, err := bson.ParseDecimal128("2.5") + require.NoError(t, err) + decimal300, err := bson.ParseDecimal128("300") + require.NoError(t, err) + + type sample struct { + label string + id bson.D + } + samples := []sample{ + {"minkey-later-key", bson.D{{"z", bson.MinKey{}}}}, + {"null-earlier-key", bson.D{{"a", nil}}}, + {"double-neg-inf", bson.D{{"n", math.Inf(-1)}}}, + {"int32-2", bson.D{{"n", int32(2)}}}, + {"decimal-2.5", bson.D{{"n", decimal25}}}, + {"int64-10", bson.D{{"n", int64(10)}}}, + {"double-20.25", bson.D{{"n", 20.25}}}, + {"decimal-300", bson.D{{"n", decimal300}}}, + {"double-pos-inf", bson.D{{"n", math.Inf(1)}}}, + {"double-nan", bson.D{{"n", math.NaN()}}}, + {"string-a", bson.D{{"s", "a"}}}, + {"nested-a1", bson.D{{"tenant", "a"}, {"bucket", bson.D{{"region", "us"}, {"part", int32(1)}}}}}, + {"nested-a2", bson.D{{"tenant", "a"}, {"bucket", bson.D{{"region", "us"}, {"part", int32(2)}}}}}, + {"array-lex-1-2", bson.D{{"arr", bson.A{int32(1), int32(2)}}}}, + {"array-lex-1-3", bson.D{{"arr", bson.A{int32(1), int32(3)}}}}, + {"binary-short-z", bson.D{{"bin", bson.Binary{Subtype: 0x00, Data: []byte("z")}}}}, + {"binary-long-aa", bson.D{{"bin", bson.Binary{Subtype: 0x00, Data: []byte("aa")}}}}, + {"binary-subtype-4", bson.D{{"bin", bson.Binary{Subtype: 0x04, Data: []byte("a")}}}}, + {"objectid", bson.D{{"oid", bson.NewObjectID()}}}, + {"bool-false", bson.D{{"b", false}}}, + {"date", bson.D{{"d", time.Unix(100, 0)}}}, + {"timestamp", bson.D{{"ts", bson.Timestamp{T: 100, I: 1}}}}, + {"regex", bson.D{{"r", bson.Regex{Pattern: "a", Options: "i"}}}}, + {"maxkey", bson.D{{"m", bson.MaxKey{}}}}, + } + + rawIDs := make(map[string]bson.RawValue, len(samples)) + for _, s := range samples { + _, err := col.InsertOne(ctx, bson.D{{"_id", s.id}, {"label", s.label}}) + require.NoError(t, err, s.label) + rawIDs[s.label] = bsonRawValue(t, s.id) + } + + cursor, err := col.Find(ctx, bson.D{}, options.Find().SetSort(bson.D{{"_id", 1}})) + require.NoError(t, err) + defer cursor.Close(ctx) + + var serverOrder []string + for cursor.Next(ctx) { + serverOrder = append(serverOrder, cursor.Current.Lookup("label").StringValue()) + } + require.NoError(t, cursor.Err()) + + var localOrder []string + for i := len(samples) - 1; i >= 0; i-- { + localOrder = append(localOrder, samples[i].label) + } + sort.Slice(localOrder, func(i, j int) bool { + return compareBSONRawValues(rawIDs[localOrder[i]], rawIDs[localOrder[j]]) < 0 + }) + + require.Equal(t, serverOrder, localOrder) + t.Logf("server/local order: %v", serverOrder) + + serverRank := make(map[string]int, len(serverOrder)) + for i, label := range serverOrder { + serverRank[label] = i + } + for _, a := range samples { + for _, b := range samples { + want := cmpSign(serverRank[a.label] - serverRank[b.label]) + got := cmpSign(compareBSONRawValues(rawIDs[a.label], rawIDs[b.label])) + require.Equal(t, want, got, "%s vs %s", a.label, b.label) + } + } +} + +func cmpSign(i int) int { + switch { + case i < 0: + return -1 + case i > 0: + return 1 + default: + return 0 + } +} diff --git a/connectors/mongo/docdb_test.go b/connectors/mongo/docdb_test.go new file mode 100644 index 0000000..e6fb939 --- /dev/null +++ b/connectors/mongo/docdb_test.go @@ -0,0 +1,98 @@ +package mongo + +import ( + "math" + "sort" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.mongodb.org/mongo-driver/v2/bson" +) + +func bsonRawValue(t *testing.T, v any) bson.RawValue { + t.Helper() + typ, data, err := bson.MarshalValue(v) + require.NoError(t, err) + return bson.RawValue{Type: typ, Value: data} +} + +func TestSupportedIDTypesIncludesEmbeddedDocument(t *testing.T) { + assert.True(t, supportedIDTypes[bson.TypeEmbeddedDocument]) +} + +func TestCompareBSONRawValuesEmbeddedDocumentSortsCompositeIDs(t *testing.T) { + ids := []bson.RawValue{ + bsonRawValue(t, bson.D{{"tenant", "b"}, {"seq", int32(1)}}), + bsonRawValue(t, bson.D{{"tenant", "a"}, {"seq", int32(2)}}), + bsonRawValue(t, bson.D{{"tenant", "a"}, {"seq", int32(1)}}), + } + + sort.Slice(ids, func(i, j int) bool { + return compareBSONRawValues(ids[i], ids[j]) < 0 + }) + + assert.True(t, ids[0].Equal(bsonRawValue(t, bson.D{{"tenant", "a"}, {"seq", int32(1)}}))) + assert.True(t, ids[1].Equal(bsonRawValue(t, bson.D{{"tenant", "a"}, {"seq", int32(2)}}))) + assert.True(t, ids[2].Equal(bsonRawValue(t, bson.D{{"tenant", "b"}, {"seq", int32(1)}}))) +} + +func TestCompareBSONRawValuesEmbeddedDocumentSortsNestedDocuments(t *testing.T) { + low := bsonRawValue(t, bson.D{ + {"tenant", "a"}, + {"bucket", bson.D{{"region", "us"}, {"part", int32(1)}}}, + }) + high := bsonRawValue(t, bson.D{ + {"tenant", "a"}, + {"bucket", bson.D{{"region", "us"}, {"part", int32(2)}}}, + }) + + assert.Negative(t, compareBSONRawValues(low, high)) + assert.Positive(t, compareBSONRawValues(high, low)) + assert.Zero(t, compareBSONRawValues(low, low)) +} + +func TestCompareBSONRawValuesEmbeddedDocumentComparesMixedNumericTypes(t *testing.T) { + oneInt32 := bsonRawValue(t, bson.D{{"seq", int32(1)}}) + oneInt64 := bsonRawValue(t, bson.D{{"seq", int64(1)}}) + twoInt32 := bsonRawValue(t, bson.D{{"seq", int32(2)}}) + tenInt64 := bsonRawValue(t, bson.D{{"seq", int64(10)}}) + decimal, err := bson.ParseDecimal128("2.5") + require.NoError(t, err) + twoPointFiveDecimal := bsonRawValue(t, bson.D{{"seq", decimal}}) + + assert.Zero(t, compareBSONRawValues(oneInt32, oneInt64)) + assert.Negative(t, compareBSONRawValues(twoInt32, tenInt64)) + assert.Positive(t, compareBSONRawValues(twoPointFiveDecimal, twoInt32)) +} + +func TestCompareBSONRawValuesEmbeddedDocumentComparesNonFiniteNumbers(t *testing.T) { + nan := bsonRawValue(t, bson.D{{"seq", math.NaN()}}) + negInf := bsonRawValue(t, bson.D{{"seq", math.Inf(-1)}}) + one := bsonRawValue(t, bson.D{{"seq", int32(1)}}) + posInf := bsonRawValue(t, bson.D{{"seq", math.Inf(1)}}) + + assert.Negative(t, compareBSONRawValues(nan, negInf)) + assert.Negative(t, compareBSONRawValues(negInf, one)) + assert.Negative(t, compareBSONRawValues(one, posInf)) +} + +func TestCompareBSONRawValuesBinaryUsesMongoOrdering(t *testing.T) { + shortZ := bsonRawValue(t, bson.Binary{Subtype: 0x00, Data: []byte("z")}) + longAA := bsonRawValue(t, bson.Binary{Subtype: 0x00, Data: []byte("aa")}) + generic := bsonRawValue(t, bson.Binary{Subtype: 0x00, Data: []byte("a")}) + uuid := bsonRawValue(t, bson.Binary{Subtype: 0x04, Data: []byte("a")}) + + assert.Negative(t, compareBSONRawValues(shortZ, longAA)) + assert.Negative(t, compareBSONRawValues(generic, uuid)) +} + +func TestCompareBSONRawValuesEmbeddedDocumentComparesFieldTypeBeforeKey(t *testing.T) { + nullWithEarlierKey := bsonRawValue(t, bson.D{{"a", nil}}) + minKeyWithLaterKey := bsonRawValue(t, bson.D{{"b", bson.MinKey{}}}) + aKey := bsonRawValue(t, bson.D{{"a", int32(1)}}) + bKey := bsonRawValue(t, bson.D{{"b", int32(1)}}) + + assert.Negative(t, compareBSONRawValues(minKeyWithLaterKey, nullWithEarlierKey)) + assert.Negative(t, compareBSONRawValues(aKey, bKey)) +}