diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java index b11f12722a51..77d9a3967d20 100644 --- a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java +++ b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java @@ -84,6 +84,7 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableSet.toImmutableSet; import static io.trino.plugin.mongodb.ObjectIdType.OBJECT_ID; +import static io.trino.plugin.mongodb.ptf.Query.parseFilter; import static io.trino.spi.HostAddress.fromParts; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.BooleanType.BOOLEAN; @@ -474,7 +475,7 @@ static Document buildFilter(MongoTableHandle table) { // Use $and operator because Document.putAll method overwrites existing entries where the key already exists ImmutableList.Builder filter = ImmutableList.builder(); - table.getFilter().ifPresent(filter::add); + table.getFilter().ifPresent(json -> filter.add(parseFilter(json))); filter.add(buildQuery(table.getConstraint())); return andPredicate(filter.build()); } diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoTableHandle.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoTableHandle.java index 26c18b21de9c..c19e9bdaa4df 100644 --- a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoTableHandle.java +++ b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoTableHandle.java @@ -19,7 +19,6 @@ import io.trino.spi.connector.ConnectorTableHandle; import io.trino.spi.connector.SchemaTableName; import io.trino.spi.predicate.TupleDomain; -import org.bson.Document; import java.util.Objects; import java.util.Optional; @@ -34,10 +33,10 @@ public class MongoTableHandle private final SchemaTableName schemaTableName; private final RemoteTableName remoteTableName; private final TupleDomain constraint; - private final Optional filter; + private final Optional filter; private final OptionalInt limit; - public MongoTableHandle(SchemaTableName schemaTableName, RemoteTableName remoteTableName, Optional filter) + public MongoTableHandle(SchemaTableName schemaTableName, RemoteTableName remoteTableName, Optional filter) { this(schemaTableName, remoteTableName, filter, TupleDomain.all(), OptionalInt.empty()); } @@ -46,7 +45,7 @@ public MongoTableHandle(SchemaTableName schemaTableName, RemoteTableName remoteT public MongoTableHandle( @JsonProperty("schemaTableName") SchemaTableName schemaTableName, @JsonProperty("remoteTableName") RemoteTableName remoteTableName, - @JsonProperty("filter") Optional filter, + @JsonProperty("filter") Optional filter, @JsonProperty("constraint") TupleDomain constraint, @JsonProperty("limit") OptionalInt limit) { @@ -70,7 +69,7 @@ public RemoteTableName getRemoteTableName() } @JsonProperty - public Optional getFilter() + public Optional getFilter() { return filter; } diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/ptf/Query.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/ptf/Query.java index df109f20f2be..28d2737ec95c 100644 --- a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/ptf/Query.java +++ b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/ptf/Query.java @@ -120,8 +120,10 @@ public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransact throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "Only lowercase collection name is supported"); } RemoteTableName remoteTableName = mongoSession.toRemoteSchemaTableName(new SchemaTableName(database, collection)); + // Don't store Document object to MongoTableHandle for avoiding serialization issue + parseFilter(filter); - MongoTableHandle tableHandle = new MongoTableHandle(new SchemaTableName(database, collection), remoteTableName, Optional.of(parseFilter(filter))); + MongoTableHandle tableHandle = new MongoTableHandle(new SchemaTableName(database, collection), remoteTableName, Optional.of(filter)); ConnectorTableSchema tableSchema = metadata.getTableSchema(session, tableHandle); Map columnsByName = metadata.getColumnHandles(session, tableHandle); List columns = tableSchema.getColumns().stream() @@ -144,7 +146,7 @@ public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransact } } - private static Document parseFilter(String filter) + public static Document parseFilter(String filter) { try { return Document.parse(filter); diff --git a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java index 6c0c264f555f..a90171c4ae35 100644 --- a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java +++ b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java @@ -742,6 +742,20 @@ public void testNativeQueryNestedRow() assertUpdate("DROP TABLE " + tableName); } + @Test + public void testNativeQueryHelperFunction() + { + String tableName = "test_query_helper_function" + randomNameSuffix(); + MongoCollection collection = client.getDatabase("tpch").getCollection(tableName); + collection.insertOne(new Document(ImmutableMap.of("id", 1, "timestamp", LocalDateTime.of(2023, 3, 20, 1, 2, 3)))); + collection.insertOne(new Document(ImmutableMap.of("id", 2, "timestamp", LocalDateTime.of(2024, 3, 20, 1, 2, 3)))); + + assertQuery( + "SELECT id FROM TABLE(mongodb.system.query(database => 'tpch', collection => '" + tableName + "', filter => '{ timestamp: ISODate(\"2023-03-20T01:02:03.000Z\") }'))", + "VALUES 1"); + assertUpdate("DROP TABLE " + tableName); + } + @Test public void testNativeQueryFilterAndWhere() { diff --git a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoTableHandle.java b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoTableHandle.java index 9792b72989d1..60042f6042c5 100644 --- a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoTableHandle.java +++ b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoTableHandle.java @@ -15,7 +15,6 @@ import io.airlift.json.JsonCodec; import io.trino.spi.connector.SchemaTableName; -import org.bson.Document; import org.testng.annotations.Test; import java.util.Optional; @@ -57,7 +56,20 @@ public void testRoundTripWithQuery() { SchemaTableName schemaTableName = new SchemaTableName("schema", "table"); RemoteTableName remoteTableName = new RemoteTableName("schema", "table"); - MongoTableHandle expected = new MongoTableHandle(schemaTableName, remoteTableName, Optional.of(new Document("key", "value"))); + MongoTableHandle expected = new MongoTableHandle(schemaTableName, remoteTableName, Optional.of("{\"key\": \"value\"}")); + + String json = codec.toJson(expected); + MongoTableHandle actual = codec.fromJson(json); + + assertEquals(actual.getSchemaTableName(), expected.getSchemaTableName()); + } + + @Test + public void testRoundTripWithQueryHavingHelperFunction() + { + SchemaTableName schemaTableName = new SchemaTableName("schema", "table"); + RemoteTableName remoteTableName = new RemoteTableName("schema", "table"); + MongoTableHandle expected = new MongoTableHandle(schemaTableName, remoteTableName, Optional.of("{timestamp: ISODate(\"2023-03-20T01:02:03.000Z\")}")); String json = codec.toJson(expected); MongoTableHandle actual = codec.fromJson(json);