如何在Spring Batch的MongoItemWriter中指定shard key？

问题重现 #

使用Spring Batch的MongoItemWriter向分片MongoDB集群写入数据时，触发如下报错：

Write errors: [BulkWriteError{index=0, code=61, message='Failed to target upsert by query :: could not extract exact shard key', details={}}]

当前MongoItemWriter配置代码（Kotlin）：

@Bean
@StepScope
fun mongoItemWriter(mongoOperations: MongoOperations): MongoItemWriter<MongoWord> {
    return MongoItemWriterBuilder<MongoWord>()
        .collection("word")
        .template(mongoOperations)
        .build()
}

实体类MongoWord已配置分片注解：

@Document("word")
@Sharded(shardKey = ["userId"], shardingStrategy = ShardingStrategy.HASH)
data class MongoWord(
    @Id
    var id: String? = null,
    @Field(name = "user_id")
    var userId = 0L
)

完整报错堆栈：

exception is com.mongodb.MongoBulkWriteException: Bulk write operation error on server host:port. Write errors: [BulkWriteError{index=0, code=61, message='Failed to target upsert by query :: could not extract exact shard key', details={}}]. 
    at org.springframework.data.mongodb.core.DefaultBulkOperations.bulkWriteTo(DefaultBulkOperations.java:324) ~[spring-data-mongodb-3.3.3.jar!/:3.3.3]
    at org.springframework.data.mongodb.core.MongoTemplate.execute(MongoTemplate.java:560) ~[spring-data-mongodb-3.3.3.jar!/:3.3.3]
    at org.springframework.data.mongodb.core.DefaultBulkOperations.execute(DefaultBulkOperations.java:290) ~[spring-data-mongodb-3.3.3.jar!/:3.3.3]
    at org.springframework.batch.item.data.MongoItemWriter.saveOrUpdate(MongoItemWriter.java:171) ~[spring-batch-infrastructure-4.3.5.jar!/:4.3.5]
    at org.springframework.batch.item.data.MongoItemWriter.doWrite(MongoItemWriter.java:140) ~[spring-batch-infrastructure-4.3.5.jar!/:4.3.5]
    at org.springframework.batch.item.data.MongoItemWriter$1.beforeCommit(MongoItemWriter.java:201) ~[spring-batch-infrastructure-4.3.5.jar!/:4.3.5]
    at org.springframework.transaction.support.TransactionSynchronizationUtils.triggerBeforeCommit(TransactionSynchronizationUtils.java:97) ~[spring-tx-5.3.17.jar!/:5.3.17]
    at org.springframework.transaction.support.AbstractPlatformTransactionManager.triggerBeforeCommit(AbstractPlatformTransactionManager.java:916) ~[spring-tx-5.3.17.jar!/:5.3.17]
    at org.springframework.transaction.support.AbstractPlatformTransactionManager.processCommit(AbstractPlatformTransactionManager.java:727) ~[spring-tx-5.3.17.jar!/:5.3.17]
    at org.springframework.transaction.support.AbstractPlatformTransactionManager.commit(AbstractPlatformTransactionManager.java:711) ~[spring-tx-5.3.17.jar!/:5.3.17]
    at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke0(Native Method) ~[na:na]
    at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:77) ~[na:na]
    at java.base/jdk.internal.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43) ~[na:na]
    at java.base/java.lang.reflect.Method.invoke(Method.java:568) ~[na:na]

报错原因 #

MongoItemWriter默认仅使用@Id字段作为查询条件执行upsert操作。在分片集群中，当执行upsert且目标文档不存在时，MongoDB需要从查询条件中提取分片键定位分片，但默认查询条件只有_id，无法获取分片键userId，因此触发报错。

解决方案 #

通过MongoItemWriterBuilder的idMapping方法自定义查询条件，将分片键userId与id一起作为匹配条件，确保MongoDB能从查询条件中提取分片键。

修改后的配置代码：

@Bean
@StepScope
fun mongoItemWriter(mongoOperations: MongoOperations): MongoItemWriter<MongoWord> {
    return MongoItemWriterBuilder<MongoWord>()
        .collection("word")
        .template(mongoOperations)
        .idMapping { mongoWord ->
            // 构建包含分片键userId和id的查询条件
            Query.query(
                Criteria.where("_id").`is`(mongoWord.id)
                    .and("user_id").`is`(mongoWord.userId)
            )
        }
        .build()
}

补充说明 #

• 该配置确保upsert操作的查询条件同时包含_id和分片键user_id，MongoDB可通过分片键定位目标分片，即使文档不存在也能正常执行插入。
• 如果业务场景允许仅用分片键作为查询条件（如userId是唯一标识），可调整查询条件为仅匹配user_id，但需确保业务逻辑符合唯一性要求。

内容的提问来源于stack exchange，提问作者mozzi