diff --git a/CHANGES.md b/CHANGES.md index 151cbe14aaa2..05e4a7ea42e0 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -89,6 +89,7 @@ * (Python) Fixed incorrect profiler options handling on portable runners ([#39613](https://github.com/apache/beam/issues/39613)). * (Java) KafkaIO dynamic reads no longer require the obsolete `beam_fn_api` experiment ([#29998](https://github.com/apache/beam/issues/29998)). * (Prism) Self-checkpointing splittable DoFns now resume after their requested delay instead of immediately, so polling SDFs no longer busy-spin ([#39848](https://github.com/apache/beam/issues/39848)). +* (Java) MongoDbIO read splitting now preserves non-ObjectId `_id` types (e.g. string ids) instead of failing to parse the generated range filters ([#39900](https://github.com/apache/beam/issues/39900)). ## Security Fixes diff --git a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java index 46c3f8fcd58a..501eaa4867c4 100644 --- a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java +++ b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java @@ -67,11 +67,9 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.bson.BsonDocument; import org.bson.BsonInt32; -import org.bson.BsonObjectId; import org.bson.BsonString; import org.bson.Document; import org.bson.conversions.Bson; -import org.bson.types.ObjectId; import org.checkerframework.checker.nullness.qual.Nullable; import org.checkerframework.dataflow.qual.Pure; import org.slf4j.Logger; @@ -606,38 +604,26 @@ public List> split( @VisibleForTesting static List splitKeysToFilters(List splitKeys) { ArrayList filters = new ArrayList<>(); - String lowestBound = null; // lower boundary (previous split in the iteration) + Object lowestBound = null; // lower boundary (previous split in the iteration) for (int i = 0; i < splitKeys.size(); i++) { - String splitKey = splitKeys.get(i).get("_id").toString(); - String rangeFilter; + Object splitKey = splitKeys.get(i).get("_id"); if (i == 0) { // this is the first split in the list, the filter defines // the range from the beginning up to this split - rangeFilter = String.format("{ $and: [ {\"_id\":{$lte:ObjectId(\"%s\")}}", splitKey); - filters.add(String.format("%s ]}", rangeFilter)); + filters.add(rangeFilter(null, splitKey)); // If there is only one split, also generate a range from the split to the end if (splitKeys.size() == 1) { - rangeFilter = String.format("{ $and: [ {\"_id\":{$gt:ObjectId(\"%s\")}}", splitKey); - filters.add(String.format("%s ]}", rangeFilter)); + filters.add(rangeFilter(splitKey, null)); } } else if (i == splitKeys.size() - 1) { // this is the last split in the list, the filters define // the range from the previous split to the current split and also // the current split to the end - rangeFilter = - String.format( - "{ $and: [ {\"_id\":{$gt:ObjectId(\"%s\")," + "$lte:ObjectId(\"%s\")}}", - lowestBound, splitKey); - filters.add(String.format("%s ]}", rangeFilter)); - rangeFilter = String.format("{ $and: [ {\"_id\":{$gt:ObjectId(\"%s\")}}", splitKey); - filters.add(String.format("%s ]}", rangeFilter)); + filters.add(rangeFilter(lowestBound, splitKey)); + filters.add(rangeFilter(splitKey, null)); } else { // we are between two splits - rangeFilter = - String.format( - "{ $and: [ {\"_id\":{$gt:ObjectId(\"%s\")," + "$lte:ObjectId(\"%s\")}}", - lowestBound, splitKey); - filters.add(String.format("%s ]}", rangeFilter)); + filters.add(rangeFilter(lowestBound, splitKey)); } lowestBound = splitKey; @@ -646,6 +632,23 @@ static List splitKeysToFilters(List splitKeys) { return filters; } + /** + * Builds a JSON range filter on {@code _id} with the given bounds. Bounds are serialized with + * their actual BSON types (as extended JSON) so that ids that are not ObjectIds, such as + * application-defined string ids, are preserved. + */ + private static String rangeFilter( + @Nullable Object greaterThan, @Nullable Object lessThanOrEqualTo) { + Document range = new Document(); + if (greaterThan != null) { + range.append("$gt", greaterThan); + } + if (lessThanOrEqualTo != null) { + range.append("$lte", lessThanOrEqualTo); + } + return new Document("$and", Collections.singletonList(new Document("_id", range))).toJson(); + } + /** * Transform a list of split keys as a list of filters containing corresponding range. * @@ -674,9 +677,11 @@ static List splitKeysToFilters(List splitKeys) { @VisibleForTesting static List splitKeysToMatch(List splitKeys) { List aggregates = new ArrayList<>(); - ObjectId lowestBound = null; // lower boundary (previous split in the iteration) + Object lowestBound = null; // lower boundary (previous split in the iteration) for (int i = 0; i < splitKeys.size(); i++) { - ObjectId splitKey = splitKeys.get(i).getObjectId("_id"); + // Keep the raw value so that ids that are not ObjectIds, such as application-defined + // string ids, are preserved. + Object splitKey = splitKeys.get(i).get("_id"); if (i == 0) { aggregates.add(Aggregates.match(Filters.lte("_id", splitKey))); if (splitKeys.size() == 1) { @@ -687,22 +692,20 @@ static List splitKeysToMatch(List splitKeys) { // the range from the previous split to the current split and also // the current split to the end // Create a custom BSON document with multiple conditions on the same field - BsonDocument rangeFilter = - new BsonDocument( + Document rangeFilter = + new Document( "_id", - new BsonDocument( - "$gt", new BsonObjectId(Preconditions.checkStateNotNull(lowestBound))) - .append("$lte", new BsonObjectId(splitKey))); + new Document("$gt", Preconditions.checkStateNotNull(lowestBound)) + .append("$lte", splitKey)); aggregates.add(Aggregates.match(rangeFilter)); aggregates.add(Aggregates.match(Filters.gt("_id", splitKey))); } else { // Create a custom BSON document with multiple conditions on the same field - BsonDocument rangeFilter = - new BsonDocument( + Document rangeFilter = + new Document( "_id", - new BsonDocument( - "$gt", new BsonObjectId(Preconditions.checkStateNotNull(lowestBound))) - .append("$lte", new BsonObjectId(splitKey))); + new Document("$gt", Preconditions.checkStateNotNull(lowestBound)) + .append("$lte", splitKey)); aggregates.add(Aggregates.match(rangeFilter)); } diff --git a/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbIOTest.java b/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbIOTest.java index 94b9df527d2b..e0a0eabd5be3 100644 --- a/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbIOTest.java +++ b/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbIOTest.java @@ -103,20 +103,40 @@ public void testSplitIntoFilters() { documents.add(new Document("_id", 56)); List filters = MongoDbIO.BoundedMongoDbSource.splitKeysToFilters(documents); assertEquals(2, filters.size()); - assertEquals("{ $and: [ {\"_id\":{$lte:ObjectId(\"56\")}} ]}", filters.get(0)); - assertEquals("{ $and: [ {\"_id\":{$gt:ObjectId(\"56\")}} ]}", filters.get(1)); + assertEquals(56, idRange(filters.get(0)).getInt32("$lte").getValue()); + assertEquals(56, idRange(filters.get(1)).getInt32("$gt").getValue()); // Add two more splits; now we should have 4 filters documents.add(new Document("_id", 109)); documents.add(new Document("_id", 256)); filters = MongoDbIO.BoundedMongoDbSource.splitKeysToFilters(documents); assertEquals(4, filters.size()); - assertEquals("{ $and: [ {\"_id\":{$lte:ObjectId(\"56\")}} ]}", filters.get(0)); - assertEquals( - "{ $and: [ {\"_id\":{$gt:ObjectId(\"56\"),$lte:ObjectId(\"109\")}} ]}", filters.get(1)); - assertEquals( - "{ $and: [ {\"_id\":{$gt:ObjectId(\"109\"),$lte:ObjectId(\"256\")}} ]}", filters.get(2)); - assertEquals("{ $and: [ {\"_id\":{$gt:ObjectId(\"256\")}} ]}", filters.get(3)); + assertEquals(56, idRange(filters.get(0)).getInt32("$lte").getValue()); + assertEquals(56, idRange(filters.get(1)).getInt32("$gt").getValue()); + assertEquals(109, idRange(filters.get(1)).getInt32("$lte").getValue()); + assertEquals(109, idRange(filters.get(2)).getInt32("$gt").getValue()); + assertEquals(256, idRange(filters.get(2)).getInt32("$lte").getValue()); + assertEquals(256, idRange(filters.get(3)).getInt32("$gt").getValue()); + } + + @Test + public void testSplitIntoFiltersWithStringId() { + // Ids that are not ObjectIds, such as application-defined string ids, must keep their type: + // they used to be formatted as ObjectId("...") and could never be parsed back (#39900). + ArrayList documents = new ArrayList<>(); + documents.add(new Document("_id", "id-aaa")); + documents.add(new Document("_id", "id-mmm")); + List filters = MongoDbIO.BoundedMongoDbSource.splitKeysToFilters(documents); + assertEquals(3, filters.size()); + assertEquals("id-aaa", idRange(filters.get(0)).getString("$lte").getValue()); + assertEquals("id-aaa", idRange(filters.get(1)).getString("$gt").getValue()); + assertEquals("id-mmm", idRange(filters.get(1)).getString("$lte").getValue()); + assertEquals("id-mmm", idRange(filters.get(2)).getString("$gt").getValue()); + } + + /** Parses a filter generated by splitKeysToFilters and returns the range on {@code _id}. */ + private static BsonDocument idRange(String filter) { + return BsonDocument.parse(filter).getArray("$and").get(0).asDocument().getDocument("_id"); } @Test @@ -152,6 +172,22 @@ public void testSplitIntoBucket() { buckets.get(3).toString()); } + @Test + public void testSplitIntoBucketWithStringId() { + // Ids that are not ObjectIds, such as application-defined string ids, must keep their type + // instead of being read as ObjectIds (#39900). + ArrayList documents = new ArrayList<>(); + documents.add(new Document("_id", "id-aaa")); + documents.add(new Document("_id", "id-mmm")); + List buckets = MongoDbIO.BoundedMongoDbSource.splitKeysToMatch(documents); + assertEquals(3, buckets.size()); + assertEquals("{\"$match\": {\"_id\": {\"$lte\": \"id-aaa\"}}}", buckets.get(0).toString()); + assertEquals( + "{\"$match\": {\"_id\": {\"$gt\": \"id-aaa\", \"$lte\": \"id-mmm\"}}}", + buckets.get(1).toString()); + assertEquals("{\"$match\": {\"_id\": {\"$gt\": \"id-mmm\"}}}", buckets.get(2).toString()); + } + @Test public void testBuildAutoBuckets() { List aggregates = new ArrayList();