From 7a3a1ddeb97ee82d5955bfc3e2311806208a0b81 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pablo=20Francisco=20P=C3=A9rez=20Hidalgo?= Date: Fri, 18 Sep 2026 20:47:29 +0200 Subject: [PATCH] feat(mongodb-structured-ingestion) DBI-1098: Expose BSON to JSON conversion in flow/pkg This enables ClickPipes' Discovery MongoDB schema inference. Part of: https://linear.app/clickhouse/issue/DBI-1098 --- flow/connectors/mongo/qvalue_convert.go | 227 +------------------ flow/connectors/mongo/qvalue_convert_test.go | 20 -- flow/pkg/go.mod | 3 + flow/pkg/go.sum | 8 + flow/pkg/mongo/schema.go | 224 ++++++++++++++++++ flow/pkg/mongo/schema_test.go | 37 +++ 6 files changed, 280 insertions(+), 239 deletions(-) create mode 100644 flow/pkg/mongo/schema_test.go diff --git a/flow/connectors/mongo/qvalue_convert.go b/flow/connectors/mongo/qvalue_convert.go index b351e9e098..a7c431518c 100644 --- a/flow/connectors/mongo/qvalue_convert.go +++ b/flow/connectors/mongo/qvalue_convert.go @@ -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" ) @@ -52,7 +49,7 @@ 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 @@ -60,7 +57,7 @@ func (c *DirectBsonConverter) QValueJSONFromDocument(raw bson.Raw) (types.QValue 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 @@ -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 @@ -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: @@ -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: @@ -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: @@ -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 @@ -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": , "Data": ""}. -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": "", "Options": ""}. -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": , "I": }. -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)) - } -} diff --git a/flow/connectors/mongo/qvalue_convert_test.go b/flow/connectors/mongo/qvalue_convert_test.go index e7e95d3b8e..f9f10b3d82 100644 --- a/flow/connectors/mongo/qvalue_convert_test.go +++ b/flow/connectors/mongo/qvalue_convert_test.go @@ -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) { diff --git a/flow/pkg/go.mod b/flow/pkg/go.mod index 6af3bab8dd..e0a69531b0 100644 --- a/flow/pkg/go.mod +++ b/flow/pkg/go.mod @@ -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 @@ -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 diff --git a/flow/pkg/go.sum b/flow/pkg/go.sum index 42b8aed36a..0ee03f3747 100644 --- a/flow/pkg/go.sum +++ b/flow/pkg/go.sum @@ -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= @@ -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= diff --git a/flow/pkg/mongo/schema.go b/flow/pkg/mongo/schema.go index ac6dd261cb..6b685bfb73 100644 --- a/flow/pkg/mongo/schema.go +++ b/flow/pkg/mongo/schema.go @@ -2,11 +2,19 @@ package mongo import ( "context" + "encoding/base64" + "encoding/hex" + "fmt" + "math" "slices" + "strconv" "strings" + "time" + jsoniter "github.com/json-iterator/go" "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/mongo" + "go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore" ) // Filter out system databases after listing all databases. @@ -65,3 +73,219 @@ func GetCollectionNames(ctx context.Context, client *mongo.Client, databaseName slices.Sort(filteredCollNames) return filteredCollNames, nil } + +func RawDocumentToJSON(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 +} + +// RawValueToJSON writes a BSON value to stream as the JSON PeerDB lands MongoDB documents with, following the +// MongoDB ClickPipes type mapping (https://clickhouse.com/docs/integrations/clickpipes/mongodb/datatypes). +// +// The deprecated Undefined, MinKey and MaxKey render as {}, DBPointer as {"DB": "", "Pointer": ""} +// and CodeWithScope as {"Code": "", "Scope": {...}}. +// +// RawDocumentToJSON and RawArrayToJSON are its entry points for a whole document or array; WriteBinaryJSON, +// WriteRegexJSON and WriteTimestampJSON expose the encodings of the BSON types rendered as JSON objects. +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 RawDocumentToJSON(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 := RawDocumentToJSON(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": , "Data": ""}. +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": "", "Options": ""}. +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": , "I": }. +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)) + } +} diff --git a/flow/pkg/mongo/schema_test.go b/flow/pkg/mongo/schema_test.go new file mode 100644 index 0000000000..a5f9daffdf --- /dev/null +++ b/flow/pkg/mongo/schema_test.go @@ -0,0 +1,37 @@ +package mongo + +import ( + "fmt" + "math" + "strings" + "testing" + + jsoniter "github.com/json-iterator/go" + "github.com/stretchr/testify/require" + "go.mongodb.org/mongo-driver/v2/bson" + "go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore" +) + +func rawValueOf(t *testing.T, value any) bsoncore.Value { + t.Helper() + raw, err := bson.Marshal(bson.D{{Key: "a", Value: value}}) + require.NoError(t, err) + rv := bson.Raw(raw).Lookup("a") + return bsoncore.Value{Type: bsoncore.Type(rv.Type), Data: rv.Value} +} + +func TestRawValueToJSONFloatLimitBoundary(t *testing.T) { + stream := jsoniter.NewStream(jsoniter.ConfigCompatibleWithStandardLibrary, nil, 64) + 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)), + } { + t.Run(fmt.Sprint(value), func(t *testing.T) { + stream.Reset(nil) + require.NoError(t, RawValueToJSON(rawValueOf(t, value), stream)) + result := string(stream.Buffer()) + require.Less(t, len(result), 31) + require.True(t, strings.Contains(result, ".") || strings.Contains(result, "e"), result) + }) + } +}