diff --git a/docs/docs-cn/source/5.application-development/3.connector/12.mongodb.md b/docs/docs-cn/source/5.application-development/3.connector/12.mongodb.md new file mode 100644 index 000000000..4ff162086 --- /dev/null +++ b/docs/docs-cn/source/5.application-development/3.connector/12.mongodb.md @@ -0,0 +1,94 @@ +# MongoDB Connector 介绍 + +MongoDB Connector 支持读取有界 collection，并将数据插入 MongoDB collection。 + +## 语法 + +```sql +CREATE TABLE mongo_source ( + id BIGINT, + name VARCHAR, + active BOOLEAN +) WITH ( + type = 'mongodb', + `geaflow.dsl.mongodb.uri` = 'mongodb://localhost:27017', + `geaflow.dsl.mongodb.database` = 'geaflow', + `geaflow.dsl.mongodb.collection` = 'source_records', + `geaflow.dsl.mongodb.partition.num` = '4', + `geaflow.dsl.mongodb.partition.field` = 'id', + `geaflow.dsl.mongodb.partition.lowerbound` = '0', + `geaflow.dsl.mongodb.partition.upperbound` = '100' +); +``` + +## 参数 + +| 参数名 | 是否必须 | 默认值 | 描述 | +| --- | --- | --- | --- | +| `geaflow.dsl.mongodb.uri` | 是 | - | MongoDB connection string。 | +| `geaflow.dsl.mongodb.database` | 是 | - | database 名称。 | +| `geaflow.dsl.mongodb.collection` | 是 | - | collection 名称。 | +| `geaflow.dsl.mongodb.batch.size` | 否 | `1000` | sink 每批插入的行数。 | +| `geaflow.dsl.mongodb.partition.num` | 否 | `1` | source 范围分区数上限。 | +| `geaflow.dsl.mongodb.partition.field` | 分区数大于 1 时 | - | 使用整数边界拆分的数值型字段。 | +| `geaflow.dsl.mongodb.partition.lowerbound` | 分区数大于 1 时 | - | 范围下界，包含该值。 | +| `geaflow.dsl.mongodb.partition.upperbound` | 分区数大于 1 时 | - | 范围上界，不包含该值。 | + +分区数为 1 时，source 不使用范围条件。分区数大于 1 时，配置范围会拆分为互不重叠的 +`[lowerbound, upperbound)` 查询条件，范围外的文档不会被读取。实际分区数取 +`partition.num` 和 `upperbound - lowerbound` 中的较小值，避免产生空的整数范围。 + +对于大小窗口，source 使用 bookmark offset 和固定排序。单分区按 `_id` 排序；范围分区按 +分区字段和 `_id` 排序。范围分区读取建议创建对应的升序复合索引，例如 +`{id: 1, _id: 1}`。全量窗口直接流式读取 MongoDB cursor。source 仅支持静态、有界 +collection，读取期间修改数据仍可能造成遗漏或重复。 + +sink 使用 ordered insert，在达到批量阈值或窗口结束时写入。不支持 upsert 和 +Exactly-once，也不在 MongoDB driver 的重试行为之外增加重试。 +批量失败时，MongoDB 可能已经写入部分文档。 + +## 类型映射 + +| GeaFlow 类型 | BSON 类型 | +| --- | --- | +| `VARCHAR` | String | +| `BOOLEAN` | Boolean | +| `TINYINT`、`SMALLINT`、`INTEGER` | Int32 | +| `BIGINT` | Int64 | +| `FLOAT`、`DOUBLE` | Double | +| `DECIMAL` | Decimal128 | +| `DATE`、`TIMESTAMP` | Date | + +字段缺失和 BSON null 均读取为 null。表结构中的 `_id` 为 `VARCHAR` 时，source 会将 +BSON ObjectId 转为十六进制字符串；sink 写入的字符串 `_id` 保持 BSON string。 +首版不支持数组和嵌套文档。 + +## 示例 + +```sql +CREATE TABLE mongo_source ( + id BIGINT, + name VARCHAR, + active BOOLEAN +) WITH ( + type = 'mongodb', + `geaflow.dsl.mongodb.uri` = 'mongodb://localhost:27017', + `geaflow.dsl.mongodb.database` = 'geaflow', + `geaflow.dsl.mongodb.collection` = 'source_records' +); + +CREATE TABLE mongo_sink ( + id BIGINT, + name VARCHAR, + active BOOLEAN +) WITH ( + type = 'mongodb', + `geaflow.dsl.mongodb.uri` = 'mongodb://localhost:27017', + `geaflow.dsl.mongodb.database` = 'geaflow', + `geaflow.dsl.mongodb.collection` = 'sink_records', + `geaflow.dsl.mongodb.batch.size` = '500' +); + +INSERT INTO mongo_sink +SELECT id, name, active FROM mongo_source; +``` diff --git a/docs/docs-cn/source/5.application-development/3.connector/index.rst b/docs/docs-cn/source/5.application-development/3.connector/index.rst index 2511eee35..5c0a9faa0 100644 --- a/docs/docs-cn/source/5.application-development/3.connector/index.rst +++ b/docs/docs-cn/source/5.application-development/3.connector/index.rst @@ -17,4 +17,4 @@ 9.pulsar.md 10.udc.md 11.doris.md - + 12.mongodb.md diff --git a/docs/docs-en/source/5.application-development/3.connector/12.mongodb.md b/docs/docs-en/source/5.application-development/3.connector/12.mongodb.md new file mode 100644 index 000000000..71c0af71d --- /dev/null +++ b/docs/docs-en/source/5.application-development/3.connector/12.mongodb.md @@ -0,0 +1,96 @@ +# MongoDB Connector + +The MongoDB connector reads a bounded collection and inserts rows into a collection. + +## Syntax + +```sql +CREATE TABLE mongo_source ( + id BIGINT, + name VARCHAR, + active BOOLEAN +) WITH ( + type = 'mongodb', + `geaflow.dsl.mongodb.uri` = 'mongodb://localhost:27017', + `geaflow.dsl.mongodb.database` = 'geaflow', + `geaflow.dsl.mongodb.collection` = 'source_records', + `geaflow.dsl.mongodb.partition.num` = '4', + `geaflow.dsl.mongodb.partition.field` = 'id', + `geaflow.dsl.mongodb.partition.lowerbound` = '0', + `geaflow.dsl.mongodb.partition.upperbound` = '100' +); +``` + +## Options + +| Key | Required | Default | Description | +| --- | --- | --- | --- | +| `geaflow.dsl.mongodb.uri` | Yes | - | MongoDB connection string. | +| `geaflow.dsl.mongodb.database` | Yes | - | Database name. | +| `geaflow.dsl.mongodb.collection` | Yes | - | Collection name. | +| `geaflow.dsl.mongodb.batch.size` | No | `1000` | Number of rows in each sink insert batch. | +| `geaflow.dsl.mongodb.partition.num` | No | `1` | Maximum number of source range partitions. | +| `geaflow.dsl.mongodb.partition.field` | When partition number is greater than 1 | - | Numeric field split with integer range bounds. | +| `geaflow.dsl.mongodb.partition.lowerbound` | For multiple partitions | - | Inclusive range lower bound. | +| `geaflow.dsl.mongodb.partition.upperbound` | For multiple partitions | - | Exclusive range upper bound. | + +With one partition, the source scans the collection without a range filter. With multiple +partitions, the configured range is split into non-overlapping `[lowerbound, upperbound)` +filters. Documents outside that range are not read. The actual partition count is the smaller +of `partition.num` and `upperbound - lowerbound`, which avoids empty integer ranges. + +For size windows, the source uses bookmark offsets with a stable sort. A single partition sorts +by `_id`; range partitions sort by the partition field and `_id`. For efficient range-partitioned +reads, create an ascending compound index on both fields, for example `{id: 1, _id: 1}`. +All-window reads stream directly from the MongoDB cursor. The source supports static, bounded +collections; changes made during a read can cause missing or repeated documents. + +The sink uses ordered inserts. It flushes when the batch size is reached and when a window +finishes. It does not provide upsert, exactly-once delivery, or retries beyond the MongoDB driver +behavior configured by the connection string. A failed batch can be partially written by MongoDB. + +## Type mapping + +| GeaFlow type | BSON type | +| --- | --- | +| `VARCHAR` | String | +| `BOOLEAN` | Boolean | +| `TINYINT`, `SMALLINT`, `INTEGER` | Int32 | +| `BIGINT` | Int64 | +| `FLOAT`, `DOUBLE` | Double | +| `DECIMAL` | Decimal128 | +| `DATE`, `TIMESTAMP` | Date | + +Missing fields and BSON null values are read as null. If `_id` is declared as `VARCHAR`, a BSON +ObjectId is returned as its hexadecimal string. A string `_id` written by the sink remains a +BSON string. Arrays and nested documents are not supported. + +## Example + +```sql +CREATE TABLE mongo_source ( + id BIGINT, + name VARCHAR, + active BOOLEAN +) WITH ( + type = 'mongodb', + `geaflow.dsl.mongodb.uri` = 'mongodb://localhost:27017', + `geaflow.dsl.mongodb.database` = 'geaflow', + `geaflow.dsl.mongodb.collection` = 'source_records' +); + +CREATE TABLE mongo_sink ( + id BIGINT, + name VARCHAR, + active BOOLEAN +) WITH ( + type = 'mongodb', + `geaflow.dsl.mongodb.uri` = 'mongodb://localhost:27017', + `geaflow.dsl.mongodb.database` = 'geaflow', + `geaflow.dsl.mongodb.collection` = 'sink_records', + `geaflow.dsl.mongodb.batch.size` = '500' +); + +INSERT INTO mongo_sink +SELECT id, name, active FROM mongo_source; +``` diff --git a/docs/docs-en/source/5.application-development/3.connector/index.rst b/docs/docs-en/source/5.application-development/3.connector/index.rst index ddfa8998c..e77b75458 100644 --- a/docs/docs-en/source/5.application-development/3.connector/index.rst +++ b/docs/docs-en/source/5.application-development/3.connector/index.rst @@ -16,4 +16,5 @@ Connector 8.hudi.md 9.pulsar.md 10.udc.md - 11.doris.md \ No newline at end of file + 11.doris.md + 12.mongodb.md diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/pom.xml b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/pom.xml new file mode 100644 index 000000000..7ea2b12a2 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/pom.xml @@ -0,0 +1,85 @@ + + + + + + org.apache.geaflow + geaflow-dsl-connector + 0.8.0-SNAPSHOT + + + 4.0.0 + + geaflow-dsl-connector-mongodb + + + 4.11.1 + 1.19.8 + + + + + org.apache.geaflow + geaflow-dsl-common + + + org.apache.geaflow + geaflow-dsl-connector-api + + + org.mongodb + mongodb-driver-sync + ${mongodb-driver.version} + + + + org.testng + testng + ${testng.version} + test + + + org.testcontainers + mongodb + ${testcontainers.version} + test + + + + + + + org.apache.maven.plugins + maven-failsafe-plugin + ${maven-surefire.version} + + + + integration-test + verify + + + + + + + diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoConfigKeys.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoConfigKeys.java new file mode 100644 index 000000000..de8d08ea9 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoConfigKeys.java @@ -0,0 +1,69 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.mongodb; + +import org.apache.geaflow.common.config.ConfigKey; +import org.apache.geaflow.common.config.ConfigKeys; + +public class MongoConfigKeys { + + public static final ConfigKey GEAFLOW_DSL_MONGODB_URI = ConfigKeys + .key("geaflow.dsl.mongodb.uri") + .noDefaultValue() + .description("MongoDB connection string."); + + public static final ConfigKey GEAFLOW_DSL_MONGODB_DATABASE = ConfigKeys + .key("geaflow.dsl.mongodb.database") + .noDefaultValue() + .description("MongoDB database name."); + + public static final ConfigKey GEAFLOW_DSL_MONGODB_COLLECTION = ConfigKeys + .key("geaflow.dsl.mongodb.collection") + .noDefaultValue() + .description("MongoDB collection name."); + + public static final ConfigKey GEAFLOW_DSL_MONGODB_BATCH_SIZE = ConfigKeys + .key("geaflow.dsl.mongodb.batch.size") + .defaultValue(1000) + .description("MongoDB sink batch size."); + + public static final ConfigKey GEAFLOW_DSL_MONGODB_PARTITION_NUM = ConfigKeys + .key("geaflow.dsl.mongodb.partition.num") + .defaultValue(1) + .description("MongoDB source partition number."); + + public static final ConfigKey GEAFLOW_DSL_MONGODB_PARTITION_FIELD = ConfigKeys + .key("geaflow.dsl.mongodb.partition.field") + .noDefaultValue() + .description("MongoDB source range partition field."); + + public static final ConfigKey GEAFLOW_DSL_MONGODB_PARTITION_LOWERBOUND = ConfigKeys + .key("geaflow.dsl.mongodb.partition.lowerbound") + .noDefaultValue() + .description("Inclusive lower bound for MongoDB source partitions."); + + public static final ConfigKey GEAFLOW_DSL_MONGODB_PARTITION_UPPERBOUND = ConfigKeys + .key("geaflow.dsl.mongodb.partition.upperbound") + .noDefaultValue() + .description("Exclusive upper bound for MongoDB source partitions."); + + private MongoConfigKeys() { + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoOffset.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoOffset.java new file mode 100644 index 000000000..64e25743f --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoOffset.java @@ -0,0 +1,73 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.mongodb; + +import org.apache.geaflow.dsl.connector.api.Offset; +import org.bson.BsonDocument; +import org.bson.json.JsonMode; +import org.bson.json.JsonWriterSettings; + +public class MongoOffset implements Offset { + + static final String ID_FIELD = "_id"; + static final String PARTITION_VALUE_FIELD = "partitionValue"; + + private static final JsonWriterSettings JSON_SETTINGS = JsonWriterSettings.builder() + .outputMode(JsonMode.EXTENDED) + .build(); + + private final long offset; + private final String bookmark; + + public MongoOffset(long offset) { + this(offset, null); + } + + MongoOffset(long offset, BsonDocument bookmark) { + if (offset < 0) { + throw new IllegalArgumentException("MongoDB offset must not be negative"); + } + this.offset = offset; + this.bookmark = bookmark == null ? null : bookmark.toJson(JSON_SETTINGS); + } + + boolean hasBookmark() { + return bookmark != null; + } + + BsonDocument getBookmark() { + return bookmark == null ? null : BsonDocument.parse(bookmark); + } + + @Override + public String humanReadable() { + return String.valueOf(offset); + } + + @Override + public long getOffset() { + return offset; + } + + @Override + public boolean isTimestamp() { + return false; + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoPartition.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoPartition.java new file mode 100644 index 000000000..cf4331acc --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoPartition.java @@ -0,0 +1,99 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.mongodb; + +import com.mongodb.client.model.Filters; +import java.util.Objects; +import org.apache.geaflow.dsl.connector.api.Partition; +import org.bson.BsonDocument; +import org.bson.conversions.Bson; + +public class MongoPartition implements Partition { + + private final String collection; + private final int index; + private final String field; + private final Long lowerBound; + private final Long upperBound; + + public MongoPartition(String collection, int index) { + this(collection, index, null, null, null); + } + + public MongoPartition(String collection, int index, String field, Long lowerBound, + Long upperBound) { + this.collection = Objects.requireNonNull(collection); + this.index = index; + this.field = field; + this.lowerBound = lowerBound; + this.upperBound = upperBound; + } + + public String getCollection() { + return collection; + } + + public String getField() { + return field; + } + + public Long getLowerBound() { + return lowerBound; + } + + public Long getUpperBound() { + return upperBound; + } + + public boolean hasRange() { + return field != null; + } + + public Bson toFilter() { + if (!hasRange()) { + return new BsonDocument(); + } + return Filters.and(Filters.gte(field, lowerBound), Filters.lt(field, upperBound)); + } + + @Override + public String getName() { + return collection + "-" + index; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof MongoPartition)) { + return false; + } + MongoPartition that = (MongoPartition) o; + return index == that.index && Objects.equals(collection, that.collection) + && Objects.equals(field, that.field) && Objects.equals(lowerBound, that.lowerBound) + && Objects.equals(upperBound, that.upperBound); + } + + @Override + public int hashCode() { + return Objects.hash(collection, index, field, lowerBound, upperBound); + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoRowConverter.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoRowConverter.java new file mode 100644 index 000000000..4347f8f14 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoRowConverter.java @@ -0,0 +1,236 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.mongodb; + +import java.math.BigDecimal; +import java.sql.Timestamp; +import java.util.Collections; +import java.util.List; +import org.apache.geaflow.common.binary.BinaryString; +import org.apache.geaflow.common.config.Configuration; +import org.apache.geaflow.common.type.IType; +import org.apache.geaflow.common.type.Types; +import org.apache.geaflow.dsl.common.data.Row; +import org.apache.geaflow.dsl.common.data.impl.ObjectRow; +import org.apache.geaflow.dsl.common.exception.GeaFlowDSLException; +import org.apache.geaflow.dsl.common.types.StructType; +import org.apache.geaflow.dsl.common.types.TableField; +import org.apache.geaflow.dsl.connector.api.serde.TableDeserializer; +import org.bson.Document; +import org.bson.types.Decimal128; +import org.bson.types.ObjectId; + +class MongoRowConverter implements TableDeserializer { + + private StructType schema; + + MongoRowConverter(StructType schema) { + this.schema = schema; + } + + @Override + public void init(Configuration conf, StructType structType) { + this.schema = structType; + } + + @Override + public List deserialize(Document document) { + return Collections.singletonList(toRow(document)); + } + + Row toRow(Document document) { + Object[] values = new Object[schema.size()]; + for (int i = 0; i < schema.size(); i++) { + TableField field = schema.getField(i); + values[i] = fromBson(field.getName(), field.getType(), document.get(field.getName())); + } + return ObjectRow.create(values); + } + + Document toDocument(Row row) { + Document document = new Document(); + for (int i = 0; i < schema.size(); i++) { + TableField field = schema.getField(i); + Object value = row.getField(i, field.getType()); + document.put(field.getName(), toBson(field.getName(), field.getType(), value)); + } + return document; + } + + private Object fromBson(String field, IType type, Object value) { + if (value == null) { + return null; + } + switch (type.getName()) { + case Types.TYPE_NAME_STRING: + return asString(field, type, value); + case Types.TYPE_NAME_BINARY_STRING: + return BinaryString.fromString(asString(field, type, value)); + case Types.TYPE_NAME_BOOLEAN: + return requireType(field, type, value, Boolean.class); + case Types.TYPE_NAME_BYTE: + return (byte) integralValue(field, type, value, Byte.MIN_VALUE, Byte.MAX_VALUE); + case Types.TYPE_NAME_SHORT: + return (short) integralValue(field, type, value, Short.MIN_VALUE, Short.MAX_VALUE); + case Types.TYPE_NAME_INTEGER: + return (int) integralValue(field, type, value, Integer.MIN_VALUE, Integer.MAX_VALUE); + case Types.TYPE_NAME_LONG: + return integralValue(field, type, value, Long.MIN_VALUE, Long.MAX_VALUE); + case Types.TYPE_NAME_FLOAT: + return floatValue(field, type, value); + case Types.TYPE_NAME_DOUBLE: + return decimalValue(field, type, value); + case Types.TYPE_NAME_DECIMAL: + return decimal128Value(field, type, value).bigDecimalValue(); + case Types.TYPE_NAME_TIMESTAMP: + return new Timestamp(dateValue(field, type, value).getTime()); + case Types.TYPE_NAME_DATE: + return new java.sql.Date(dateValue(field, type, value).getTime()); + default: + throw conversionError(field, type, value); + } + } + + private Object toBson(String field, IType type, Object value) { + if (value == null) { + return null; + } + switch (type.getName()) { + case Types.TYPE_NAME_STRING: + return requireType(field, type, value, String.class); + case Types.TYPE_NAME_BINARY_STRING: + return requireType(field, type, value, BinaryString.class).toString(); + case Types.TYPE_NAME_BOOLEAN: + return requireType(field, type, value, Boolean.class); + case Types.TYPE_NAME_BYTE: + return (int) requireType(field, type, value, Byte.class); + case Types.TYPE_NAME_SHORT: + return (int) requireType(field, type, value, Short.class); + case Types.TYPE_NAME_INTEGER: + return requireType(field, type, value, Integer.class); + case Types.TYPE_NAME_LONG: + return requireType(field, type, value, Long.class); + case Types.TYPE_NAME_FLOAT: + return (double) requireType(field, type, value, Float.class); + case Types.TYPE_NAME_DOUBLE: + return requireType(field, type, value, Double.class); + case Types.TYPE_NAME_DECIMAL: + return new Decimal128(requireType(field, type, value, BigDecimal.class)); + case Types.TYPE_NAME_TIMESTAMP: + Timestamp timestamp = requireType(field, type, value, Timestamp.class); + return new java.util.Date(timestamp.getTime()); + case Types.TYPE_NAME_DATE: + java.sql.Date date = requireType(field, type, value, java.sql.Date.class); + return new java.util.Date(date.getTime()); + default: + throw conversionError(field, type, value); + } + } + + private String asString(String field, IType type, Object value) { + if (value instanceof String) { + return (String) value; + } + if ("_id".equals(field) && value instanceof ObjectId) { + return ((ObjectId) value).toHexString(); + } + throw conversionError(field, type, value); + } + + private long integralValue(String field, IType type, Object value, long min, long max) { + if (!(value instanceof Number)) { + throw conversionError(field, type, value); + } + Number number = (Number) value; + try { + BigDecimal decimal; + if (number instanceof Decimal128) { + decimal = ((Decimal128) number).bigDecimalValue(); + } else if (number instanceof BigDecimal) { + decimal = (BigDecimal) number; + } else if (number instanceof Byte || number instanceof Short + || number instanceof Integer || number instanceof Long) { + decimal = BigDecimal.valueOf(number.longValue()); + } else { + double doubleValue = number.doubleValue(); + if (!Double.isFinite(doubleValue)) { + throw conversionError(field, type, value); + } + decimal = BigDecimal.valueOf(doubleValue); + } + long longValue = decimal.longValueExact(); + if (longValue < min || longValue > max) { + throw conversionError(field, type, value); + } + return longValue; + } catch (ArithmeticException e) { + throw conversionError(field, type, value); + } + } + + private float floatValue(String field, IType type, Object value) { + double result = decimalValue(field, type, value); + float floatResult = (float) result; + if (!Float.isFinite(floatResult)) { + throw conversionError(field, type, value); + } + return floatResult; + } + + private double decimalValue(String field, IType type, Object value) { + if (!(value instanceof Number)) { + throw conversionError(field, type, value); + } + double result = ((Number) value).doubleValue(); + if (!Double.isFinite(result)) { + throw conversionError(field, type, value); + } + return result; + } + + private Decimal128 decimal128Value(String field, IType type, Object value) { + if (value instanceof Decimal128) { + return (Decimal128) value; + } + if (value instanceof BigDecimal) { + return new Decimal128((BigDecimal) value); + } + throw conversionError(field, type, value); + } + + private java.util.Date dateValue(String field, IType type, Object value) { + if (!(value instanceof java.util.Date)) { + throw conversionError(field, type, value); + } + return (java.util.Date) value; + } + + private T requireType(String field, IType type, Object value, Class valueClass) { + if (!valueClass.isInstance(value)) { + throw conversionError(field, type, value); + } + return valueClass.cast(value); + } + + private GeaFlowDSLException conversionError(String field, IType type, Object value) { + return new GeaFlowDSLException("Cannot map MongoDB field '{}' from {} to {}", field, + value.getClass().getSimpleName(), type.getName()); + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableConnector.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableConnector.java new file mode 100644 index 000000000..c60d27ee5 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableConnector.java @@ -0,0 +1,46 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.mongodb; + +import org.apache.geaflow.common.config.Configuration; +import org.apache.geaflow.dsl.connector.api.TableReadableConnector; +import org.apache.geaflow.dsl.connector.api.TableSink; +import org.apache.geaflow.dsl.connector.api.TableSource; +import org.apache.geaflow.dsl.connector.api.TableWritableConnector; + +public class MongoTableConnector implements TableReadableConnector, TableWritableConnector { + + public static final String TYPE = "MONGODB"; + + @Override + public String getType() { + return TYPE; + } + + @Override + public TableSource createSource(Configuration conf) { + return new MongoTableSource(); + } + + @Override + public TableSink createSink(Configuration conf) { + return new MongoTableSink(); + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableSink.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableSink.java new file mode 100644 index 000000000..694670f6c --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableSink.java @@ -0,0 +1,131 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.mongodb; + +import com.mongodb.MongoBulkWriteException; +import com.mongodb.MongoException; +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoClients; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.model.InsertManyOptions; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import org.apache.geaflow.api.context.RuntimeContext; +import org.apache.geaflow.common.config.ConfigKey; +import org.apache.geaflow.common.config.Configuration; +import org.apache.geaflow.dsl.common.data.Row; +import org.apache.geaflow.dsl.common.exception.GeaFlowDSLException; +import org.apache.geaflow.dsl.common.types.StructType; +import org.apache.geaflow.dsl.connector.api.TableSink; +import org.bson.Document; + +public class MongoTableSink implements TableSink { + + private String uri; + private String databaseName; + private String collectionName; + private int batchSize; + private MongoRowConverter converter; + private List batch; + + private transient MongoClient client; + private transient MongoCollection collection; + + @Override + public void init(Configuration tableConf, StructType schema) { + this.uri = required(tableConf, MongoConfigKeys.GEAFLOW_DSL_MONGODB_URI); + this.databaseName = required(tableConf, MongoConfigKeys.GEAFLOW_DSL_MONGODB_DATABASE); + this.collectionName = required(tableConf, MongoConfigKeys.GEAFLOW_DSL_MONGODB_COLLECTION); + this.batchSize = tableConf.getInteger(MongoConfigKeys.GEAFLOW_DSL_MONGODB_BATCH_SIZE); + if (batchSize <= 0) { + throw new GeaFlowDSLException("MongoDB batch size must be greater than zero"); + } + this.converter = new MongoRowConverter(schema); + this.batch = new ArrayList<>(batchSize); + } + + @Override + public void open(RuntimeContext context) { + try { + client = MongoClients.create(uri); + collection = client.getDatabase(databaseName).getCollection(collectionName); + } catch (MongoException | IllegalArgumentException e) { + close(); + throw new GeaFlowDSLException("Failed to create MongoDB sink client", e); + } + } + + @Override + public void write(Row row) throws IOException { + if (collection == null) { + throw new GeaFlowDSLException("MongoDB sink is not open"); + } + batch.add(converter.toDocument(row)); + if (batch.size() >= batchSize) { + flush(); + } + } + + @Override + public void finish() throws IOException { + flush(); + } + + @Override + public void close() { + if (client != null) { + client.close(); + client = null; + collection = null; + } + } + + private void flush() throws IOException { + if (batch.isEmpty()) { + return; + } + try { + collection.insertMany(batch, new InsertManyOptions().ordered(true)); + batch.clear(); + } catch (MongoBulkWriteException e) { + String writeConcernError = e.getWriteConcernError() == null ? "" + : "; write concern error: " + e.getWriteConcernError(); + throw new IOException("MongoDB bulk write failed for collection " + collectionName + + "; successful inserts in this batch: " + + e.getWriteResult().getInsertedCount() + "; write errors: " + + e.getWriteErrors() + writeConcernError, e); + } catch (MongoException e) { + throw new IOException("Failed to write MongoDB collection " + collectionName, e); + } + } + + private static String required(Configuration conf, ConfigKey key) { + if (!conf.contains(key)) { + throw new GeaFlowDSLException("Missing MongoDB configuration '{}'", key.getKey()); + } + String value = conf.getString(key); + if (value == null || value.trim().isEmpty()) { + throw new GeaFlowDSLException("MongoDB configuration '{}' must not be blank", + key.getKey()); + } + return value; + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableSource.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableSource.java new file mode 100644 index 000000000..de77d702f --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableSource.java @@ -0,0 +1,378 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.mongodb; + +import com.mongodb.MongoException; +import com.mongodb.client.FindIterable; +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoClients; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoCursor; +import com.mongodb.client.model.Filters; +import com.mongodb.client.model.Sorts; +import java.io.IOException; +import java.math.BigInteger; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.Optional; +import org.apache.geaflow.api.context.RuntimeContext; +import org.apache.geaflow.api.window.WindowType; +import org.apache.geaflow.common.config.ConfigKey; +import org.apache.geaflow.common.config.Configuration; +import org.apache.geaflow.dsl.common.exception.GeaFlowDSLException; +import org.apache.geaflow.dsl.common.types.TableSchema; +import org.apache.geaflow.dsl.connector.api.FetchData; +import org.apache.geaflow.dsl.connector.api.Offset; +import org.apache.geaflow.dsl.connector.api.Partition; +import org.apache.geaflow.dsl.connector.api.TableSource; +import org.apache.geaflow.dsl.connector.api.serde.TableDeserializer; +import org.apache.geaflow.dsl.connector.api.window.FetchWindow; +import org.bson.BsonDocument; +import org.bson.BsonValue; +import org.bson.Document; +import org.bson.conversions.Bson; + +public class MongoTableSource implements TableSource { + + private String uri; + private String databaseName; + private String collectionName; + private int partitionNum; + private String partitionField; + private long lowerBound; + private long upperBound; + private TableSchema schema; + + private transient MongoClient client; + private transient MongoCollection collection; + + @Override + public void init(Configuration tableConf, TableSchema tableSchema) { + this.uri = required(tableConf, MongoConfigKeys.GEAFLOW_DSL_MONGODB_URI); + this.databaseName = required(tableConf, MongoConfigKeys.GEAFLOW_DSL_MONGODB_DATABASE); + this.collectionName = required(tableConf, MongoConfigKeys.GEAFLOW_DSL_MONGODB_COLLECTION); + this.partitionNum = tableConf.getInteger(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_NUM); + this.schema = tableSchema; + + if (partitionNum <= 0) { + throw new GeaFlowDSLException("MongoDB partition number must be greater than zero"); + } + if (partitionNum > 1) { + this.partitionField = required(tableConf, + MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_FIELD); + requireConfig(tableConf, MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_LOWERBOUND); + requireConfig(tableConf, MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_UPPERBOUND); + this.lowerBound = tableConf.getLong( + MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_LOWERBOUND); + this.upperBound = tableConf.getLong( + MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_UPPERBOUND); + if (lowerBound >= upperBound) { + throw new GeaFlowDSLException( + "MongoDB partition upper bound must be greater than lower bound"); + } + } + } + + @Override + public void open(RuntimeContext context) { + try { + client = MongoClients.create(uri); + collection = client.getDatabase(databaseName).getCollection(collectionName); + } catch (MongoException | IllegalArgumentException e) { + close(); + throw new GeaFlowDSLException("Failed to create MongoDB source client", e); + } + } + + @Override + public List listPartitions() { + if (partitionNum == 1) { + return Collections.singletonList(new MongoPartition(collectionName, 0)); + } + + BigInteger lower = BigInteger.valueOf(lowerBound); + BigInteger range = BigInteger.valueOf(upperBound).subtract(lower); + int count = range.min(BigInteger.valueOf(partitionNum)).intValueExact(); + BigInteger[] strideAndRemainder = range.divideAndRemainder(BigInteger.valueOf(count)); + BigInteger current = lower; + List partitions = new ArrayList<>(count); + for (int i = 0; i < count; i++) { + BigInteger increment = strideAndRemainder[0]; + if (i < strideAndRemainder[1].intValue()) { + increment = increment.add(BigInteger.ONE); + } + BigInteger next = current.add(increment); + partitions.add(new MongoPartition(collectionName, i, partitionField, + current.longValueExact(), next.longValueExact())); + current = next; + } + return partitions; + } + + @Override + public List listPartitions(int parallelism) { + return listPartitions(); + } + + @Override + @SuppressWarnings("unchecked") + public TableDeserializer getDeserializer(Configuration conf) { + return (TableDeserializer) new MongoRowConverter(schema); + } + + @Override + @SuppressWarnings("unchecked") + public FetchData fetch(Partition partition, Optional startOffset, + FetchWindow windowInfo) throws IOException { + if (collection == null) { + throw new GeaFlowDSLException("MongoDB source is not open"); + } + if (!(partition instanceof MongoPartition)) { + throw new GeaFlowDSLException("Invalid MongoDB partition"); + } + MongoPartition mongoPartition = (MongoPartition) partition; + if (!collectionName.equals(mongoPartition.getCollection())) { + throw new GeaFlowDSLException("MongoDB partition belongs to another collection"); + } + if (windowInfo.getType() != WindowType.SIZE_TUMBLING_WINDOW + && windowInfo.getType() != WindowType.ALL_WINDOW) { + throw new GeaFlowDSLException("Unsupported MongoDB fetch window: {}", + windowInfo.getType()); + } + + MongoOffset offset = getMongoOffset(startOffset); + + boolean allWindow = windowInfo.getType() == WindowType.ALL_WINDOW; + int limit = 0; + if (!allWindow) { + long windowSize = windowInfo.windowSize(); + if (windowSize <= 0 || windowSize > Integer.MAX_VALUE) { + throw new GeaFlowDSLException("MongoDB fetch window size is out of range: {}", + windowSize); + } + limit = (int) windowSize; + } + + Bson filter = createFilter(mongoPartition, offset); + if (allWindow) { + try { + MongoCursor cursor = collection.find(filter).iterator(); + return (FetchData) FetchData.createBatchFetch( + new MongoCursorIterator(cursor, collectionName), offset); + } catch (MongoException e) { + throw new IOException("Failed to read MongoDB collection " + collectionName, e); + } + } + + List documents = new ArrayList<>(); + try { + FindIterable query = collection.find(filter) + .sort(mongoPartition.hasRange() + ? Sorts.ascending(mongoPartition.getField(), "_id") : Sorts.ascending("_id")) + .limit(limit); + try (MongoCursor cursor = query.iterator()) { + while (cursor.hasNext()) { + documents.add(cursor.next()); + } + } + } catch (MongoException e) { + throw new IOException("Failed to read MongoDB collection " + collectionName, e); + } + + MongoOffset nextOffset = createNextOffset(offset, mongoPartition, documents); + boolean finished = documents.size() < limit; + return (FetchData) FetchData.createStreamFetch(documents, nextOffset, finished); + } + + static Bson createFilter(MongoPartition partition, MongoOffset offset) { + if (!offset.hasBookmark()) { + return partition.toFilter(); + } + + Bson bookmarkFilter = createBookmarkFilter(partition, offset.getBookmark()); + return partition.hasRange() + ? Filters.and(partition.toFilter(), bookmarkFilter) : bookmarkFilter; + } + + static Bson createBookmarkFilter(MongoPartition partition, BsonDocument bookmark) { + BsonValue lastId = requiredBookmarkValue(bookmark, MongoOffset.ID_FIELD); + if (!partition.hasRange()) { + return Filters.gt(MongoOffset.ID_FIELD, lastId); + } + + BsonValue lastPartitionValue = requiredBookmarkValue(bookmark, + MongoOffset.PARTITION_VALUE_FIELD); + return Filters.or( + Filters.gt(partition.getField(), lastPartitionValue), + Filters.and( + Filters.eq(partition.getField(), lastPartitionValue), + Filters.gt(MongoOffset.ID_FIELD, lastId))); + } + + private static BsonValue requiredBookmarkValue(BsonDocument bookmark, String field) { + BsonValue value = bookmark.get(field); + if (value == null) { + throw new GeaFlowDSLException("MongoDB offset is missing bookmark field '{}'", + field); + } + return value; + } + + private MongoOffset createNextOffset(MongoOffset offset, MongoPartition partition, + List documents) { + if (documents.isEmpty()) { + return offset; + } + + Document lastDocument = documents.get(documents.size() - 1); + BsonDocument bsonDocument = lastDocument.toBsonDocument(Document.class, + collection.getCodecRegistry()); + BsonValue lastId = getDocumentValue(bsonDocument, MongoOffset.ID_FIELD); + if (lastId == null) { + throw new GeaFlowDSLException("MongoDB document is missing field '{}'", + MongoOffset.ID_FIELD); + } + + BsonDocument bookmark = new BsonDocument(MongoOffset.ID_FIELD, lastId); + if (partition.hasRange()) { + BsonValue partitionValue = getDocumentValue(bsonDocument, partition.getField()); + if (partitionValue == null) { + throw new GeaFlowDSLException("MongoDB document is missing partition field '{}'", + partition.getField()); + } + bookmark.append(MongoOffset.PARTITION_VALUE_FIELD, partitionValue); + } + + long nextOffset; + try { + nextOffset = Math.addExact(offset.getOffset(), documents.size()); + } catch (ArithmeticException e) { + throw new GeaFlowDSLException("MongoDB offset exceeds the supported range", e); + } + return new MongoOffset(nextOffset, bookmark); + } + + private static BsonValue getDocumentValue(BsonDocument document, String field) { + BsonValue value = document; + for (String name : field.split("\\.")) { + if (!value.isDocument()) { + return null; + } + value = value.asDocument().get(name); + if (value == null) { + return null; + } + } + return value; + } + + private static MongoOffset getMongoOffset(Optional startOffset) { + if (!startOffset.isPresent()) { + return new MongoOffset(0L); + } + Offset offset = startOffset.get(); + if (!(offset instanceof MongoOffset)) { + throw new GeaFlowDSLException("Invalid MongoDB offset"); + } + MongoOffset mongoOffset = (MongoOffset) offset; + if (mongoOffset.getOffset() > 0 && !mongoOffset.hasBookmark()) { + throw new GeaFlowDSLException("MongoDB offset is missing a bookmark"); + } + return mongoOffset; + } + + @Override + public void close() { + if (client != null) { + client.close(); + client = null; + collection = null; + } + } + + private static String required(Configuration conf, ConfigKey key) { + requireConfig(conf, key); + String value = conf.getString(key); + if (value == null || value.trim().isEmpty()) { + throw new GeaFlowDSLException("MongoDB configuration '{}' must not be blank", + key.getKey()); + } + return value; + } + + private static void requireConfig(Configuration conf, ConfigKey key) { + if (!conf.contains(key)) { + throw new GeaFlowDSLException("Missing MongoDB configuration '{}'", key.getKey()); + } + } + + private static class MongoCursorIterator implements Iterator { + + private final MongoCursor cursor; + private final String collectionName; + private boolean closed; + + private MongoCursorIterator(MongoCursor cursor, String collectionName) { + this.cursor = cursor; + this.collectionName = collectionName; + } + + @Override + public boolean hasNext() { + if (closed) { + return false; + } + try { + boolean hasNext = cursor.hasNext(); + if (!hasNext) { + close(); + } + return hasNext; + } catch (MongoException e) { + close(); + throw readException(e); + } + } + + @Override + public Document next() { + try { + return cursor.next(); + } catch (MongoException e) { + close(); + throw readException(e); + } + } + + private void close() { + if (!closed) { + closed = true; + cursor.close(); + } + } + + private GeaFlowDSLException readException(MongoException cause) { + return new GeaFlowDSLException( + "Failed to read MongoDB collection " + collectionName, cause); + } + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/resources/META-INF/services/org.apache.geaflow.dsl.connector.api.TableConnector b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/resources/META-INF/services/org.apache.geaflow.dsl.connector.api.TableConnector new file mode 100644 index 000000000..6dc59cf73 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/main/resources/META-INF/services/org.apache.geaflow.dsl.connector.api.TableConnector @@ -0,0 +1,19 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +org.apache.geaflow.dsl.connector.mongodb.MongoTableConnector diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/test/java/org/apache/geaflow/dsl/connector/mongodb/MongoRowConverterTest.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/test/java/org/apache/geaflow/dsl/connector/mongodb/MongoRowConverterTest.java new file mode 100644 index 000000000..c0e127bd6 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/test/java/org/apache/geaflow/dsl/connector/mongodb/MongoRowConverterTest.java @@ -0,0 +1,146 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.mongodb; + +import java.math.BigDecimal; +import java.sql.Date; +import java.sql.Timestamp; +import org.apache.geaflow.common.binary.BinaryString; +import org.apache.geaflow.common.type.Types; +import org.apache.geaflow.dsl.common.data.Row; +import org.apache.geaflow.dsl.common.data.impl.ObjectRow; +import org.apache.geaflow.dsl.common.exception.GeaFlowDSLException; +import org.apache.geaflow.dsl.common.types.StructType; +import org.apache.geaflow.dsl.common.types.TableField; +import org.bson.Document; +import org.bson.types.Decimal128; +import org.bson.types.ObjectId; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class MongoRowConverterTest { + + private static final StructType SCHEMA = new StructType( + new TableField("_id", Types.BINARY_STRING), + new TableField("name", Types.BINARY_STRING), + new TableField("count", Types.INTEGER), + new TableField("total", Types.LONG), + new TableField("score", Types.DOUBLE), + new TableField("active", Types.BOOLEAN), + new TableField("price", Types.DECIMAL), + new TableField("created", Types.TIMESTAMP), + new TableField("day", Types.DATE)); + + @Test + public void testDocumentToRow() { + ObjectId id = new ObjectId(); + java.util.Date time = new java.util.Date(1710000000000L); + Document document = new Document("_id", id) + .append("name", "alice") + .append("count", 3) + .append("total", 8L) + .append("score", 1.5D) + .append("active", true) + .append("price", new Decimal128(new BigDecimal("12.30"))) + .append("created", time) + .append("day", time); + + Row row = new MongoRowConverter(SCHEMA).toRow(document); + + Assert.assertEquals(row.getField(0, Types.BINARY_STRING), + BinaryString.fromString(id.toHexString())); + Assert.assertEquals(row.getField(1, Types.BINARY_STRING), + BinaryString.fromString("alice")); + Assert.assertEquals(row.getField(2, Types.INTEGER), 3); + Assert.assertEquals(row.getField(3, Types.LONG), 8L); + Assert.assertEquals(row.getField(4, Types.DOUBLE), 1.5D); + Assert.assertEquals(row.getField(5, Types.BOOLEAN), true); + Assert.assertEquals(row.getField(6, Types.DECIMAL), new BigDecimal("12.30")); + Assert.assertEquals(((Timestamp) row.getField(7, Types.TIMESTAMP)).getTime(), time.getTime()); + Assert.assertEquals(((Date) row.getField(8, Types.DATE)).getTime(), time.getTime()); + } + + @Test + public void testRowToDocument() { + Timestamp timestamp = new Timestamp(1710000000000L); + Date date = new Date(1710028800000L); + Row row = ObjectRow.create( + BinaryString.fromString("record-1"), + BinaryString.fromString("alice"), + 3, + 8L, + 1.5D, + true, + new BigDecimal("12.30"), + timestamp, + date); + + Document document = new MongoRowConverter(SCHEMA).toDocument(row); + + Assert.assertEquals(document.getString("_id"), "record-1"); + Assert.assertEquals(document.getString("name"), "alice"); + Assert.assertEquals(document.getInteger("count"), Integer.valueOf(3)); + Assert.assertEquals(document.getLong("total"), Long.valueOf(8)); + Assert.assertEquals(document.getDouble("score"), Double.valueOf(1.5D)); + Assert.assertEquals(document.getBoolean("active"), Boolean.TRUE); + Assert.assertEquals(document.get("price"), new Decimal128(new BigDecimal("12.30"))); + Assert.assertEquals(document.getDate("created").getTime(), timestamp.getTime()); + Assert.assertEquals(document.getDate("day").getTime(), date.getTime()); + } + + @Test + public void testExactDecimalToLong() { + StructType schema = new StructType(new TableField("value", Types.LONG)); + Document document = new Document("value", + new Decimal128(new BigDecimal("9007199254740992"))); + + Row row = new MongoRowConverter(schema).toRow(document); + + Assert.assertEquals(row.getField(0, Types.LONG), 9007199254740992L); + } + + @Test(expectedExceptions = GeaFlowDSLException.class, + expectedExceptionsMessageRegExp = ".*value.*Double.*FLOAT.*") + public void testRejectFloatOverflow() { + StructType schema = new StructType(new TableField("value", Types.FLOAT)); + Document document = new Document("value", Double.MAX_VALUE); + + new MongoRowConverter(schema).toRow(document); + } + + @Test(expectedExceptions = GeaFlowDSLException.class, + expectedExceptionsMessageRegExp = ".*value.*Decimal128.*LONG.*") + public void testRejectFractionalDecimalAsLong() { + StructType schema = new StructType(new TableField("value", Types.LONG)); + Document document = new Document("value", + new Decimal128(new BigDecimal("9007199254740992.1"))); + + new MongoRowConverter(schema).toRow(document); + } + + @Test(expectedExceptions = GeaFlowDSLException.class, + expectedExceptionsMessageRegExp = ".*name.*Document.*BINARY_STRING.*") + public void testRejectNestedDocument() { + StructType schema = new StructType(new TableField("name", Types.BINARY_STRING)); + Document document = new Document("name", new Document("first", "alice")); + + new MongoRowConverter(schema).toRow(document); + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/test/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableConnectorIT.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/test/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableConnectorIT.java new file mode 100644 index 000000000..30f295486 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/test/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableConnectorIT.java @@ -0,0 +1,298 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.mongodb; + +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoClients; +import java.io.IOException; +import java.util.Arrays; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import org.apache.geaflow.common.binary.BinaryString; +import org.apache.geaflow.common.config.Configuration; +import org.apache.geaflow.common.type.Types; +import org.apache.geaflow.dsl.common.data.Row; +import org.apache.geaflow.dsl.common.data.impl.ObjectRow; +import org.apache.geaflow.dsl.common.types.TableField; +import org.apache.geaflow.dsl.common.types.TableSchema; +import org.apache.geaflow.dsl.connector.api.FetchData; +import org.apache.geaflow.dsl.connector.api.Offset; +import org.apache.geaflow.dsl.connector.api.Partition; +import org.apache.geaflow.dsl.connector.api.serde.TableDeserializer; +import org.apache.geaflow.dsl.connector.api.window.AllFetchWindow; +import org.apache.geaflow.dsl.connector.api.window.SizeFetchWindow; +import org.bson.Document; +import org.testcontainers.DockerClientFactory; +import org.testcontainers.containers.MongoDBContainer; +import org.testcontainers.utility.DockerImageName; +import org.testng.Assert; +import org.testng.SkipException; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +public class MongoTableConnectorIT { + + private static final String DATABASE = "geaflow"; + private static final String COLLECTION = "records"; + private static final TableSchema SCHEMA = new TableSchema( + new TableField("id", Types.INTEGER, false), + new TableField("name", Types.BINARY_STRING), + new TableField("active", Types.BOOLEAN)); + private static final TableSchema BOOKMARK_SCHEMA = new TableSchema( + new TableField("_id", Types.BINARY_STRING, false), + new TableField("group_id", Types.INTEGER, false)); + + private MongoDBContainer container; + + @BeforeClass + public void startMongoDB() { + boolean dockerAvailable; + try { + dockerAvailable = DockerClientFactory.instance().isDockerAvailable(); + } catch (RuntimeException e) { + throw new SkipException("Docker is not available", e); + } + if (!dockerAvailable) { + throw new SkipException("Docker is not available"); + } + container = new MongoDBContainer(DockerImageName.parse("mongo:6.0.14")); + container.start(); + } + + @AfterClass(alwaysRun = true) + public void stopMongoDB() { + if (container != null) { + container.stop(); + } + } + + @BeforeMethod + public void clearCollection() { + try (MongoClient client = MongoClients.create(container.getReplicaSetUrl())) { + client.getDatabase(DATABASE).getCollection(COLLECTION).deleteMany(new Document()); + } + } + + @Test + public void testBatchWriteAndPartitionedRead() throws IOException { + Configuration sinkConf = baseConfig(); + sinkConf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_BATCH_SIZE, "2"); + MongoTableSink sink = new MongoTableSink(); + sink.init(sinkConf, SCHEMA); + sink.open(null); + + sink.write(row(1, "alice")); + sink.write(row(2, "bob")); + Assert.assertEquals(countDocuments(), 2L); + sink.write(row(3, "carol")); + Assert.assertEquals(countDocuments(), 2L); + sink.finish(); + sink.close(); + Assert.assertEquals(countDocuments(), 3L); + + Configuration sourceConf = baseConfig(); + sourceConf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_NUM, "2"); + sourceConf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_FIELD, "id"); + sourceConf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_LOWERBOUND, "1"); + sourceConf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_UPPERBOUND, "4"); + MongoTableSource source = new MongoTableSource(); + source.init(sourceConf, SCHEMA); + source.open(null); + + Set ids = new HashSet<>(); + List partitions = source.listPartitions(); + for (Partition partition : partitions) { + readPartition(source, sourceConf, partition, ids); + } + source.close(); + + Assert.assertEquals(ids, new HashSet<>(Arrays.asList(1, 2, 3))); + } + + @Test + public void testCompositeBookmarkPagination() throws IOException { + insertDocuments(Arrays.asList( + new Document("_id", "z").append("group_id", 1), + new Document("_id", "a").append("group_id", 2), + new Document("_id", "y").append("group_id", 6), + new Document("_id", "b").append("group_id", 7))); + + Configuration conf = baseConfig(); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_NUM, "2"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_FIELD, "group_id"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_LOWERBOUND, "0"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_UPPERBOUND, "10"); + MongoTableSource source = new MongoTableSource(); + source.init(conf, BOOKMARK_SCHEMA); + source.open(null); + + Set ids = new HashSet<>(); + try { + for (Partition partition : source.listPartitions()) { + readStringPartition(source, conf, partition, ids); + } + } finally { + source.close(); + } + + Assert.assertEquals(ids, new HashSet<>(Arrays.asList("z", "a", "y", "b"))); + } + + @Test + public void testAllWindowRead() throws IOException { + insertDocuments(Arrays.asList( + new Document("id", 1), + new Document("id", 2), + new Document("id", 3))); + + Configuration conf = baseConfig(); + MongoTableSource source = new MongoTableSource(); + source.init(conf, SCHEMA); + source.open(null); + + try { + Partition partition = source.listPartitions().get(0); + FetchData fetchData = source.fetch(partition, Optional.empty(), + new AllFetchWindow(0)); + Assert.assertEquals(fetchData.getDataSize(), -1); + Assert.assertTrue(fetchData.isFinish()); + + int count = 0; + Iterator iterator = fetchData.getDataIterator(); + while (iterator.hasNext()) { + iterator.next(); + count++; + } + Assert.assertEquals(count, 3); + } finally { + source.close(); + } + } + + @Test + public void testOrderedBulkWriteFailure() throws IOException { + insertDocuments(Arrays.asList( + new Document("_id", "duplicate").append("group_id", 0))); + + Configuration conf = baseConfig(); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_BATCH_SIZE, "3"); + MongoTableSink sink = new MongoTableSink(); + sink.init(conf, BOOKMARK_SCHEMA); + sink.open(null); + + IOException failure = null; + try { + sink.write(bookmarkRow("first", 1)); + sink.write(bookmarkRow("duplicate", 2)); + try { + sink.write(bookmarkRow("last", 3)); + } catch (IOException e) { + failure = e; + } + } finally { + sink.close(); + } + + Assert.assertNotNull(failure); + Assert.assertTrue(failure.getMessage().contains( + "successful inserts in this batch: 1")); + Assert.assertTrue(failure.getMessage().contains("write errors:")); + Assert.assertEquals(countDocuments(), 2L); + Assert.assertEquals(countDocuments(new Document("_id", "last")), 0L); + } + + private void readPartition(MongoTableSource source, Configuration conf, Partition partition, + Set ids) throws IOException { + Optional offset = Optional.empty(); + boolean finished; + TableDeserializer deserializer = source.getDeserializer(conf); + deserializer.init(conf, SCHEMA); + do { + FetchData fetchData = source.fetch(partition, offset, + new SizeFetchWindow(0, 1)); + Iterator iterator = fetchData.getDataIterator(); + while (iterator.hasNext()) { + Row row = deserializer.deserialize(iterator.next()).get(0); + Integer id = (Integer) row.getField(0, Types.INTEGER); + Assert.assertTrue(ids.add(id), "Duplicate id: " + id); + } + offset = Optional.of(fetchData.getNextOffset()); + finished = fetchData.isFinish(); + } while (!finished); + } + + private void readStringPartition(MongoTableSource source, Configuration conf, + Partition partition, Set ids) throws IOException { + Optional offset = Optional.empty(); + boolean finished; + TableDeserializer deserializer = source.getDeserializer(conf); + deserializer.init(conf, BOOKMARK_SCHEMA); + do { + FetchData fetchData = source.fetch(partition, offset, + new SizeFetchWindow(0, 1)); + Iterator iterator = fetchData.getDataIterator(); + while (iterator.hasNext()) { + Row row = deserializer.deserialize(iterator.next()).get(0); + BinaryString id = (BinaryString) row.getField(0, Types.BINARY_STRING); + Assert.assertTrue(ids.add(id.toString()), "Duplicate id: " + id); + } + offset = Optional.of(fetchData.getNextOffset()); + finished = fetchData.isFinish(); + } while (!finished); + } + + private Row row(int id, String name) { + return ObjectRow.create(id, BinaryString.fromString(name), true); + } + + private Row bookmarkRow(String id, int groupId) { + return ObjectRow.create(BinaryString.fromString(id), groupId); + } + + private void insertDocuments(List documents) { + try (MongoClient client = MongoClients.create(container.getReplicaSetUrl())) { + client.getDatabase(DATABASE).getCollection(COLLECTION).insertMany(documents); + } + } + + private long countDocuments() { + return countDocuments(new Document()); + } + + private long countDocuments(Document filter) { + try (MongoClient client = MongoClients.create(container.getReplicaSetUrl())) { + return client.getDatabase(DATABASE).getCollection(COLLECTION) + .countDocuments(filter); + } + } + + private Configuration baseConfig() { + Configuration conf = new Configuration(); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_URI, container.getReplicaSetUrl()); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_DATABASE, DATABASE); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_COLLECTION, COLLECTION); + return conf; + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/test/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableConnectorTest.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/test/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableConnectorTest.java new file mode 100644 index 000000000..0a4e9ed08 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-mongodb/src/test/java/org/apache/geaflow/dsl/connector/mongodb/MongoTableConnectorTest.java @@ -0,0 +1,206 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.mongodb; + +import com.mongodb.MongoClientSettings; +import java.util.List; +import org.apache.geaflow.common.config.Configuration; +import org.apache.geaflow.common.serialize.SerializerFactory; +import org.apache.geaflow.common.type.Types; +import org.apache.geaflow.dsl.common.exception.GeaFlowDSLException; +import org.apache.geaflow.dsl.common.types.TableField; +import org.apache.geaflow.dsl.common.types.TableSchema; +import org.apache.geaflow.dsl.connector.api.Partition; +import org.apache.geaflow.dsl.connector.api.TableConnector; +import org.apache.geaflow.dsl.connector.api.util.ConnectorFactory; +import org.bson.BsonDocument; +import org.bson.BsonInt32; +import org.bson.BsonString; +import org.bson.conversions.Bson; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class MongoTableConnectorTest { + + private static final TableSchema SCHEMA = new TableSchema( + new TableField("id", Types.INTEGER, false)); + + @Test + public void testLoadConnector() { + TableConnector connector = ConnectorFactory.loadConnector("mongodb"); + Assert.assertEquals(connector.getType(), MongoTableConnector.TYPE); + } + + @Test + public void testSinglePartition() { + MongoTableSource source = new MongoTableSource(); + source.init(baseConfig(), SCHEMA); + + List partitions = source.listPartitions(); + + Assert.assertEquals(partitions.size(), 1); + MongoPartition partition = (MongoPartition) partitions.get(0); + Assert.assertFalse(partition.hasRange()); + } + + @Test + public void testRangePartitions() { + Configuration conf = baseConfig(); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_NUM, "3"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_FIELD, "id"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_LOWERBOUND, "0"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_UPPERBOUND, "10"); + MongoTableSource source = new MongoTableSource(); + source.init(conf, SCHEMA); + + List partitions = source.listPartitions(); + + Assert.assertEquals(partitions.size(), 3); + assertBounds((MongoPartition) partitions.get(0), 0L, 4L); + assertBounds((MongoPartition) partitions.get(1), 4L, 7L); + assertBounds((MongoPartition) partitions.get(2), 7L, 10L); + } + + @Test(expectedExceptions = GeaFlowDSLException.class, + expectedExceptionsMessageRegExp = ".*partition.upperbound.*") + public void testMissingPartitionBound() { + Configuration conf = baseConfig(); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_NUM, "2"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_FIELD, "id"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_LOWERBOUND, "0"); + + new MongoTableSource().init(conf, SCHEMA); + } + + @Test(expectedExceptions = GeaFlowDSLException.class, + expectedExceptionsMessageRegExp = ".*batch size.*") + public void testInvalidBatchSize() { + Configuration conf = baseConfig(); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_BATCH_SIZE, "0"); + + new MongoTableSink().init(conf, SCHEMA); + } + + @Test(expectedExceptions = GeaFlowDSLException.class, + expectedExceptionsMessageRegExp = ".*partition number.*") + public void testInvalidPartitionNumber() { + Configuration conf = baseConfig(); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_NUM, "0"); + + new MongoTableSource().init(conf, SCHEMA); + } + + @Test(expectedExceptions = GeaFlowDSLException.class, + expectedExceptionsMessageRegExp = ".*upper bound.*") + public void testInvalidPartitionRange() { + Configuration conf = baseConfig(); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_NUM, "2"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_FIELD, "id"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_LOWERBOUND, "10"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_UPPERBOUND, "10"); + + new MongoTableSource().init(conf, SCHEMA); + } + + @Test(expectedExceptions = GeaFlowDSLException.class, + expectedExceptionsMessageRegExp = ".*mongodb.uri.*") + public void testMissingUri() { + Configuration conf = new Configuration(); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_DATABASE, "geaflow"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_COLLECTION, "records"); + + new MongoTableSource().init(conf, SCHEMA); + } + + @Test + public void testPartitionCountLimitedByRange() { + Configuration conf = baseConfig(); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_NUM, "4"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_FIELD, "id"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_LOWERBOUND, "0"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_PARTITION_UPPERBOUND, "2"); + MongoTableSource source = new MongoTableSource(); + source.init(conf, SCHEMA); + + List partitions = source.listPartitions(); + + Assert.assertEquals(partitions.size(), 2); + assertBounds((MongoPartition) partitions.get(0), 0L, 1L); + assertBounds((MongoPartition) partitions.get(1), 1L, 2L); + } + + @Test + public void testSinglePartitionBookmarkFilter() { + BsonDocument bookmark = new BsonDocument(MongoOffset.ID_FIELD, + new BsonString("record-1")); + + Bson filter = MongoTableSource.createBookmarkFilter( + new MongoPartition("records", 0), bookmark); + + Assert.assertEquals(toBsonDocument(filter), + BsonDocument.parse("{'_id': {'$gt': 'record-1'}}")); + } + + @Test + public void testRangePartitionBookmarkFilter() { + BsonDocument bookmark = new BsonDocument(MongoOffset.ID_FIELD, new BsonString("z")) + .append(MongoOffset.PARTITION_VALUE_FIELD, new BsonInt32(3)); + MongoPartition partition = new MongoPartition("records", 0, "id", 0L, 10L); + + Bson filter = MongoTableSource.createBookmarkFilter(partition, bookmark); + + Assert.assertEquals(toBsonDocument(filter), BsonDocument.parse( + "{'$or': [{'id': {'$gt': 3}}, {'$and': [{'id': 3}, " + + "{'_id': {'$gt': 'z'}}]}]}")); + } + + @Test + public void testBookmarkOffsetSerialization() { + BsonDocument bookmark = new BsonDocument(MongoOffset.ID_FIELD, + new BsonString("record-1")) + .append(MongoOffset.PARTITION_VALUE_FIELD, new BsonInt32(3)); + MongoOffset offset = new MongoOffset(5L, bookmark); + + byte[] bytes = SerializerFactory.getKryoSerializer().serialize(offset); + MongoOffset restored = (MongoOffset) SerializerFactory.getKryoSerializer() + .deserialize(bytes); + + Assert.assertEquals(restored.getOffset(), 5L); + Assert.assertEquals(restored.getBookmark(), bookmark); + } + + private static Configuration baseConfig() { + Configuration conf = new Configuration(); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_URI, "mongodb://localhost:27017"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_DATABASE, "geaflow"); + conf.put(MongoConfigKeys.GEAFLOW_DSL_MONGODB_COLLECTION, "records"); + return conf; + } + + private static void assertBounds(MongoPartition partition, long lower, long upper) { + Assert.assertEquals(partition.getLowerBound(), Long.valueOf(lower)); + Assert.assertEquals(partition.getUpperBound(), Long.valueOf(upper)); + } + + private static BsonDocument toBsonDocument(Bson filter) { + return filter.toBsonDocument(BsonDocument.class, + MongoClientSettings.getDefaultCodecRegistry()); + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/pom.xml b/geaflow/geaflow-dsl/geaflow-dsl-connector/pom.xml index 4fe03c6de..9782d73ad 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-connector/pom.xml +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/pom.xml @@ -50,6 +50,7 @@ geaflow-dsl-connector-neo4j geaflow-dsl-connector-elasticsearch geaflow-dsl-connector-doris + geaflow-dsl-connector-mongodb diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/pom.xml b/geaflow/geaflow-dsl/geaflow-dsl-runtime/pom.xml index 87c9f040a..f1e31d8bf 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/pom.xml +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/pom.xml @@ -86,6 +86,11 @@ geaflow-dsl-connector-doris + + org.apache.geaflow + geaflow-dsl-connector-mongodb + + org.apache.geaflow geaflow-dsl-connector-odps diff --git a/geaflow/geaflow-dsl/pom.xml b/geaflow/geaflow-dsl/pom.xml index a7a37063f..b644025ac 100644 --- a/geaflow/geaflow-dsl/pom.xml +++ b/geaflow/geaflow-dsl/pom.xml @@ -148,6 +148,12 @@ ${project.version} + + org.apache.geaflow + geaflow-dsl-connector-mongodb + ${project.version} + + org.apache.geaflow geaflow-dsl-connector-odps diff --git a/geaflow/geaflow-examples/gql/mongodb_connector.sql b/geaflow/geaflow-examples/gql/mongodb_connector.sql new file mode 100644 index 000000000..41b3f7cbe --- /dev/null +++ b/geaflow/geaflow-examples/gql/mongodb_connector.sql @@ -0,0 +1,48 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +CREATE TABLE mongo_source ( + id BIGINT, + name VARCHAR, + active BOOLEAN +) WITH ( + type = 'mongodb', + `geaflow.dsl.mongodb.uri` = 'mongodb://localhost:27017', + `geaflow.dsl.mongodb.database` = 'geaflow', + `geaflow.dsl.mongodb.collection` = 'source_records', + `geaflow.dsl.mongodb.partition.num` = '4', + `geaflow.dsl.mongodb.partition.field` = 'id', + `geaflow.dsl.mongodb.partition.lowerbound` = '0', + `geaflow.dsl.mongodb.partition.upperbound` = '100' +); + +CREATE TABLE mongo_sink ( + id BIGINT, + name VARCHAR, + active BOOLEAN +) WITH ( + type = 'mongodb', + `geaflow.dsl.mongodb.uri` = 'mongodb://localhost:27017', + `geaflow.dsl.mongodb.database` = 'geaflow', + `geaflow.dsl.mongodb.collection` = 'sink_records', + `geaflow.dsl.mongodb.batch.size` = '500' +); + +INSERT INTO mongo_sink +SELECT id, name, active FROM mongo_source;