Skip to content

feat(mongodb-structured-ingestion) DBI-1098: Move BSON to JSON conversion to flow/pkg #4825

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

Open
pfcoperez wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(mongodb-structured-ingestion) DBI-1098: Move BSON to JSON conversion to flow/pkg #4825

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
227 changes: 8 additions & 219 deletions flow/connectors/mongo/qvalue_convert.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,17 +1,14 @@
package connmongo

import (
"encoding/base64"
"encoding/hex"
"fmt"
"math"
"strconv"
"time"

jsoniter "github.com/json-iterator/go"
"go.mongodb.org/mongo-driver/v2/bson"
"go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore"

shared_mongo "github.com/PeerDB-io/peerdb/flow/pkg/mongo"
"github.com/PeerDB-io/peerdb/flow/shared"
"github.com/PeerDB-io/peerdb/flow/shared/types"
)
Expand Down Expand Up @@ -52,15 +49,15 @@ func NewDirectBsonConverter() *DirectBsonConverter {

func (c *DirectBsonConverter) QValueJSONFromDocument(raw bson.Raw) (types.QValueJSON, error) {
c.stream.Reset(nil)
if err := rawDocToJSON(bsoncore.Document(raw), c.stream); err != nil {
if err := shared_mongo.RawDocumentToJSON(bsoncore.Document(raw), c.stream); err != nil {
return types.QValueJSON{}, fmt.Errorf("failed to convert document: %w", err)
}
return types.QValueJSON{Val: string(c.stream.Buffer())}, nil
}

func (c *DirectBsonConverter) QValueJSONFromArray(arr bson.RawArray) (types.QValueJSON, error) {
c.stream.Reset(nil)
if err := rawArrayToJSON(bsoncore.Array(arr), c.stream); err != nil {
if err := shared_mongo.RawArrayToJSON(bsoncore.Array(arr), c.stream); err != nil {
return types.QValueJSON{}, fmt.Errorf("failed to convert array: %w", err)
}
return types.QValueJSON{Val: string(c.stream.Buffer()), IsArray: true}, nil
Expand All @@ -76,7 +73,7 @@ func (c *DirectBsonConverter) QValueStringFromId(id bson.RawValue, version uint3
}
}
c.stream.Reset(nil)
if err := rawValueToJSON(bsoncore.Value{Type: bsoncore.Type(id.Type), Data: id.Value}, c.stream); err != nil {
if err := shared_mongo.RawValueToJSON(bsoncore.Value{Type: bsoncore.Type(id.Type), Data: id.Value}, c.stream); err != nil {
return types.QValueString{}, fmt.Errorf("failed to convert %s: %w", DefaultDocumentKeyColumnName, err)
}
return types.QValueString{Val: string(c.stream.Buffer())}, nil
Expand Down Expand Up @@ -105,7 +102,7 @@ func (c *DirectBsonConverter) QValueFromBsonValue(rv bson.RawValue, nullKind typ
case bsoncore.TypeBinary:
subtype, data := v.Binary()
c.stream.Reset(nil)
writeBinaryJSON(c.stream, subtype, data)
shared_mongo.WriteBinaryJSON(c.stream, subtype, data)
return types.QValueJSON{Val: string(c.stream.Buffer())}, nil

case bsoncore.TypeObjectID:
Expand All @@ -123,7 +120,7 @@ func (c *DirectBsonConverter) QValueFromBsonValue(rv bson.RawValue, nullKind typ
case bsoncore.TypeRegex:
pattern, options := v.Regex()
c.stream.Reset(nil)
writeRegexJSON(c.stream, pattern, options)
shared_mongo.WriteRegexJSON(c.stream, pattern, options)
return types.QValueJSON{Val: string(c.stream.Buffer())}, nil

case bsoncore.TypeJavaScript:
Expand All @@ -139,7 +136,7 @@ func (c *DirectBsonConverter) QValueFromBsonValue(rv bson.RawValue, nullKind typ
case bsoncore.TypeTimestamp:
t, i := v.Timestamp()
c.stream.Reset(nil)
writeTimestampJSON(c.stream, t, i)
shared_mongo.WriteTimestampJSON(c.stream, t, i)
return types.QValueJSON{Val: string(c.stream.Buffer())}, nil

case bsoncore.TypeInt64:
Expand All @@ -153,7 +150,7 @@ func (c *DirectBsonConverter) QValueFromBsonValue(rv bson.RawValue, nullKind typ
// Undefined, MinKey, MaxKey, DBPointer and CodeWithScope are deprecated and not part of the documented
// mapping; they are rendered as JSON exactly as they are inside a full document.
c.stream.Reset(nil)
if err := rawValueToJSON(v, c.stream); err != nil {
if err := shared_mongo.RawValueToJSON(v, c.stream); err != nil {
return nil, fmt.Errorf("failed to convert %s value: %w", v.Type.String(), err)
}
return types.QValueJSON{Val: string(c.stream.Buffer())}, nil
Expand All @@ -167,211 +164,3 @@ func (c *DirectBsonConverter) QValueStringFromObjectID(oid bson.ObjectID) types.
func (c *DirectBsonConverter) QValueStringFromString(s string) types.QValueString {
return types.QValueString{Val: s}
}

func rawDocToJSON(doc bsoncore.Document, stream *jsoniter.Stream) error {
length, rem, ok := bsoncore.ReadLength(doc)
if !ok {
return fmt.Errorf("failed to read document length")
}
length -= 4

stream.WriteRaw("{")
first := true
for length > 1 {
elem, next, ok := bsoncore.ReadElement(rem)
if !ok {
return fmt.Errorf("failed to read document element")
}
length -= int32(len(elem))
rem = next

if !first {
stream.WriteRaw(",")
}
first = false

stream.WriteStringWithHTMLEscaped(elem.Key())
stream.WriteRaw(":")
if err := rawValueToJSON(elem.Value(), stream); err != nil {
return err
}
}
stream.WriteRaw("}")
return nil
}

func rawArrayToJSON(arr bsoncore.Array, stream *jsoniter.Stream) error {
length, rem, ok := bsoncore.ReadLength(arr)
if !ok {
return fmt.Errorf("failed to read array length")
}
length -= 4

stream.WriteRaw("[")
first := true
for length > 1 {
elem, next, ok := bsoncore.ReadElement(rem)
if !ok {
return fmt.Errorf("failed to read array element")
}
length -= int32(len(elem))
rem = next

if !first {
stream.WriteRaw(",")
}
first = false

if err := rawValueToJSON(elem.Value(), stream); err != nil {
return err
}
}
stream.WriteRaw("]")
return nil
}

func rawValueToJSON(v bsoncore.Value, stream *jsoniter.Stream) error {
switch v.Type {
case bsoncore.TypeDouble:
writeFloat64JSON(stream, v.Double())

case bsoncore.TypeString:
stream.WriteStringWithHTMLEscaped(v.StringValue())

case bsoncore.TypeEmbeddedDocument:
return rawDocToJSON(v.Document(), stream)

case bsoncore.TypeArray:
return rawArrayToJSON(v.Array(), stream)

case bsoncore.TypeBinary:
subtype, data := v.Binary()
writeBinaryJSON(stream, subtype, data)

case bsoncore.TypeUndefined:
stream.WriteEmptyObject()

case bsoncore.TypeObjectID:
oid := v.ObjectID()
stream.WriteRaw(`"`)
stream.SetBuffer(hex.AppendEncode(stream.Buffer(), oid[:]))
stream.WriteRaw(`"`)

case bsoncore.TypeBoolean:
stream.WriteBool(v.Boolean())

case bsoncore.TypeDateTime:
stream.WriteRaw(`"`)
stream.SetBuffer(v.Time().UTC().AppendFormat(stream.Buffer(), time.RFC3339Nano))
stream.WriteRaw(`"`)

case bsoncore.TypeNull:
stream.WriteNil()

case bsoncore.TypeRegex:
pattern, options := v.Regex()
writeRegexJSON(stream, pattern, options)

case bsoncore.TypeJavaScript:
stream.WriteStringWithHTMLEscaped(v.JavaScript())

case bsoncore.TypeSymbol:
stream.WriteStringWithHTMLEscaped(v.Symbol())

case bsoncore.TypeInt32:
stream.WriteInt32(v.Int32())

case bsoncore.TypeTimestamp:
t, i := v.Timestamp()
writeTimestampJSON(stream, t, i)

case bsoncore.TypeInt64:
stream.WriteInt64(v.Int64())

case bsoncore.TypeDecimal128:
h, l := v.Decimal128()
stream.WriteString(bson.NewDecimal128(h, l).String())

case bsoncore.TypeMinKey, bsoncore.TypeMaxKey:
stream.WriteEmptyObject()

case bsoncore.TypeDBPointer: // deprecated type, kept for backwards-compatibility
ns, oid := v.DBPointer()
stream.WriteRaw(`{"DB":`)
stream.WriteStringWithHTMLEscaped(ns)
stream.WriteRaw(`,"Pointer":"`)
stream.SetBuffer(hex.AppendEncode(stream.Buffer(), oid[:]))
stream.WriteRaw(`"}`)

case bsoncore.TypeCodeWithScope: // deprecated type, kept for backwards-compatibility
code, scope := v.CodeWithScope()
stream.WriteRaw(`{"Code":`)
stream.WriteStringWithHTMLEscaped(code)
stream.WriteRaw(`,"Scope":`)
if err := rawDocToJSON(scope, stream); err != nil {
return err
}
stream.WriteRaw("}")

default:
return fmt.Errorf("unknown type: %v", v.Type.String())
}
return nil
}

// writeBinaryJSON encodes BSON binary data as {"Subtype": <int>, "Data": "<base64>"}.
func writeBinaryJSON(stream *jsoniter.Stream, subtype byte, data []byte) {
stream.WriteRaw(`{"Subtype":`)
stream.WriteUint8(subtype)
stream.WriteRaw(`,"Data":"`)
stream.SetBuffer(base64.StdEncoding.AppendEncode(stream.Buffer(), data))
stream.WriteRaw(`"}`)
}

// writeRegexJSON encodes a BSON regular expression as {"Pattern": "<pattern>", "Options": "<flags>"}.
func writeRegexJSON(stream *jsoniter.Stream, pattern string, options string) {
stream.WriteRaw(`{"Pattern":`)
stream.WriteStringWithHTMLEscaped(pattern)
stream.WriteRaw(`,"Options":`)
stream.WriteStringWithHTMLEscaped(options)
stream.WriteRaw("}")
}

// writeTimestampJSON encodes a BSON (internal) timestamp as {"T": <seconds>, "I": <increment>}.
func writeTimestampJSON(stream *jsoniter.Stream, t uint32, i uint32) {
stream.WriteRaw(`{"T":`)
stream.WriteUint32(t)
stream.WriteRaw(`,"I":`)
stream.WriteUint32(i)
stream.WriteRaw("}")
}

// Assume (and test) that values outside of these limits will come out in scientific notation
// and will be parsed as floats either way
var (
floatLimit = math.Pow10(21)
floatNegLimit = -floatLimit
)

// writeFloat64JSON encodes NaN/Inf as quoted strings, integer-valued floats with an explicit
// ".0" suffix (to hint ClickHouse to parse as float), and other values in standard notation.
func writeFloat64JSON(stream *jsoniter.Stream, v float64) {
if math.IsNaN(v) {
stream.WriteRaw(`"NaN"`)
} else if math.IsInf(v, 1) {
stream.WriteRaw(`"+Inf"`)
} else if math.IsInf(v, -1) {
stream.WriteRaw(`"-Inf"`)
} else if v < floatLimit && v > floatNegLimit && v == math.Trunc(v) {
// use explicit decimal to hint ClickHouse to parse as float
stream.SetBuffer(strconv.AppendFloat(stream.Buffer(), v, 'f', 1, 64))
} else {
// standard notation, with implementation copied from json-iterator's WriteFloat64
abs := math.Abs(v)
format := byte('f')
if abs != 0 && (abs < 1e-6 || abs >= 1e21) {
format = 'e'
}
stream.SetBuffer(strconv.AppendFloat(stream.Buffer(), v, format, -1, 64))
}
}
20 changes: 0 additions & 20 deletions flow/connectors/mongo/qvalue_convert_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -723,26 +723,6 @@ func TestMarshalFloatLengths(t *testing.T) {
})
}
}

// Test the boundary around the limit itself
for _, value := range []float64{
floatLimit, math.Nextafter(floatLimit, math.Inf(1)), math.Nextafter(floatLimit, math.Inf(-1)),
floatNegLimit, math.Nextafter(floatNegLimit, math.Inf(1)), math.Nextafter(floatNegLimit, math.Inf(-1)),
} {
name := fmt.Sprint(value)
t.Run(name, func(t *testing.T) {
input := bson.D{{Key: "a", Value: value}}
raw, err := bson.Marshal(input)
require.NoError(t, err)
result, err := converter.QValueJSONFromDocument(raw)
require.NoError(t, err)
require.Less(t, len(result.Val), 33)
require.True(t,
strings.Contains(result.Val, ".") ||
strings.Contains(result.Val, "e"),
result.Val)
})
}
}

func TestQValuesFromBsonRawInvalidIds(t *testing.T) {
Expand Down
3 changes: 3 additions & 0 deletions flow/pkg/go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ require (
github.com/go-mysql-org/go-mysql v1.15.0
github.com/jackc/pgx/v5 v5.10.0
github.com/joho/godotenv v1.5.1
github.com/json-iterator/go v1.1.12
github.com/stretchr/testify v1.12.1
go.mongodb.org/mongo-driver/v2 v2.9.0
go.temporal.io/sdk v1.48.0
Expand Down Expand Up @@ -56,6 +57,8 @@ require (
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
github.com/klauspost/compress v1.19.2 // indirect
github.com/klauspost/cpuid/v2 v2.2.5 // indirect
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
github.com/modern-go/reflect2 v1.0.2 // indirect
github.com/nexus-rpc/nexus-proto-annotations v0.1.0 // indirect
github.com/nexus-rpc/sdk-go v0.7.0 // indirect
github.com/paulmach/orb v0.13.0 // indirect
Expand Down
8 changes: 8 additions & 0 deletions flow/pkg/go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@ github.com/google/flatbuffers v23.5.26+incompatible h1:M9dgRyhJemaM4Sw8+66GHBu8i
github.com/google/flatbuffers v23.5.26+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8=
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
github.com/google/martian/v3 v3.3.3 h1:DIhPTQrbPkgs2yJYdXU/eNACCG5DVQjySNRNlflZ9Fc=
github.com/google/martian/v3 v3.3.3/go.mod h1:iEPrYcgCF7jA9OtScMFQyAlZZ4YXTKEtJ1E6RWzmBA0=
github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0=
Expand All @@ -113,12 +114,19 @@ github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
github.com/klauspost/compress v1.19.2 h1:hMRETovs/pu/dVWN7zIT1PGG8t509MwT6bO7XSi26R8=
github.com/klauspost/compress v1.19.2/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ=
github.com/klauspost/cpuid/v2 v2.2.5 h1:0E5MSMDEoAulmXNFquVs//DdoomxaoTY1kUhbc/qbZg=
github.com/klauspost/cpuid/v2 v2.2.5/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws=
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg=
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
github.com/nexus-rpc/nexus-proto-annotations v0.1.0 h1:2fELd+9sqUtNu6Fg//pw8YFsxOvp8vZ8hfP0nHhNI80=
github.com/nexus-rpc/nexus-proto-annotations v0.1.0/go.mod h1:n3UjF1bPCW8llR8tHvbxJ+27yPWrhpo8w/Yg1IOuY0Y=
github.com/nexus-rpc/sdk-go v0.7.0 h1:38NrfY5rLnZAiMMs2ZfCKI/CSDzdfJG+27iAgfA8bUI=
Expand Down
Loading
Loading