Merge pull request #747 from tapdata/TAP-10570

fix: TAP-10570 Multi-update operations are not allowed when updating …

Author: jarad0628

connectors/mongodb-connector/src/main/java/io/tapdata/mongodb/writer/MongodbWriter.java

@@ -130,10 +130,10 @@ public void writeRecord(List<TapRecordEvent> tapRecordEvents, TapTable table, Co
 
 	private void doubleActiveWrite(List<TapRecordEvent> tapRecordEvents, TapTable table, Consumer<WriteListResult<TapRecordEvent>> writeListResultConsumer, ClientSession session) throws Throwable {
 		Document doubleActiveDoc = new Document("_id", "aaaaaaaa");
-        UpdateOptions options = new UpdateOptions().upsert(true);
-        mongoDatabase.getCollection("_tap_double_active").updateOne(session, doubleActiveDoc, new Document("$set", new Document("ts", System.currentTimeMillis())), options);
-        write(table, tapRecordEvents, writeListResultConsumer, session);
-    }
+		UpdateOptions options = new UpdateOptions().upsert(true);
+		mongoDatabase.getCollection("_tap_double_active").updateOne(session, doubleActiveDoc, new Document("$set", new Document("ts", System.currentTimeMillis())), options);
+		write(table, tapRecordEvents, writeListResultConsumer, session);
+	}
 
 	private void write(TapTable table, List<TapRecordEvent> tapRecordEvents, Consumer<WriteListResult<TapRecordEvent>> writeListResultConsumer, ClientSession session) throws Throwable {
 		AtomicLong inserted = new AtomicLong(0); //insert count
@@ -528,6 +528,18 @@ private static BulkWriteOptions buildBulkWriteOptions(BulkWriteModel bulkWriteMo
 		return bulkWriteOptions;
 	}
 
+	private boolean isShardedCollection(String tableId) {
+		return shardKeyMap != null && shardKeyMap.containsKey(tableId);
+	}
+
+	private WriteModel<Document> createUpdateModel(String tableId, Document filter, Document update, UpdateOptions options) {
+		if (isShardedCollection(tableId)) {
+			return new UpdateOneModel<>(filter, update, options);
+		} else {
+			return new UpdateManyModel<>(filter, update, options);
+		}
+	}
+
 	protected List<WriteModel<Document>> normalWriteMode(AtomicLong inserted, AtomicLong updated, AtomicLong deleted, UpdateOptions options, TapTable tapTable, Collection<String> pks, TapRecordEvent recordEvent) {
 		List<WriteModel<Document>> writeModels = new ArrayList<>();
 		if (recordEvent instanceof TapInsertRecordEvent) {
@@ -552,17 +564,17 @@ protected List<WriteModel<Document>> normalWriteMode(AtomicLong inserted, Atomic
 				}
 				MongodbUtil.removeIdIfNeed(pks, insertRecordEvent.getAfter());
 				Document update = new Document(operation, insertRecordEvent.getAfter());
-				writeModels.add(new UpdateManyModel<>(pkFilter, update, options));
+				writeModels.add(createUpdateModel(tapTable.getId(), pkFilter, update, options));
 				if (MapUtils.isNotEmpty(unsetDoc)) {
-					writeModels.add(new UpdateManyModel<>(pkFilter, new Document("$unset", unsetDoc), options));
+					writeModels.add(createUpdateModel(tapTable.getId(), pkFilter, new Document("$unset", unsetDoc), options));
 				}
 			} else {
 				if (CollectionUtils.isNotEmpty(pks) && MapUtils.isNotEmpty(unsetDoc)) {
 					Document pkFilter = getPkFilter(pks, insertRecordEvent.getAfter());
 					Document update = new Document("$set", insertRecordEvent.getAfter());
-					writeModels.add(new UpdateManyModel<>(pkFilter, update, options));
+					writeModels.add(createUpdateModel(tapTable.getId(), pkFilter, update, options));
 					if (MapUtils.isNotEmpty(unsetDoc)) {
-						writeModels.add(new UpdateManyModel<>(pkFilter, new Document("$unset", unsetDoc), options));
+						writeModels.add(createUpdateModel(tapTable.getId(), pkFilter, new Document("$unset", unsetDoc), options));
 					}
 				} else {
 					writeModels.add(new InsertOneModel<>(new Document(insertRecordEvent.getAfter())));
@@ -590,7 +602,7 @@ protected List<WriteModel<Document>> normalWriteMode(AtomicLong inserted, Atomic
 				u.remove("$v"); // Exists '$v' in update operation of MongoDB(v3.6), remove it because can't apply in write model.
 				boolean isUpdate = u.keySet().stream().anyMatch(k -> k.startsWith("$"));
 				if (isUpdate) {
-					writeModel = new UpdateManyModel<>(pkFilter, u, options);
+					writeModel = createUpdateModel(tapTable.getId(), pkFilter, u, options);
 					options.upsert(false);
 				} else {
 					writeModel = new ReplaceOneModel<>(pkFilter, u, new ReplaceOptions().upsert(false));
@@ -607,10 +619,10 @@ protected List<WriteModel<Document>> normalWriteMode(AtomicLong inserted, Atomic
 					}
 					MongodbUtil.removeIdIfNeed(pks, after);
 					u.append("$set", after);
-					writeModels.add(new UpdateManyModel<>(pkFilter, u, options));
+					writeModels.add(createUpdateModel(tapTable.getId(), pkFilter, u, options));
 					Document unsetDoc = wrapUnset(recordEvent);
 					if (MapUtils.isNotEmpty(unsetDoc)) {
-						writeModels.add(new UpdateManyModel<>(pkFilter, new Document("$unset", unsetDoc), options));
+						writeModels.add(createUpdateModel(tapTable.getId(), pkFilter, new Document("$unset", unsetDoc), options));
 					}
 				}
 			}

connectors/mongodb-connector/src/main/java/io/tapdata/mongodb/writer/error/BulkWriteErrorCodeHandlerEnum.java

@@ -3,6 +3,7 @@
 import io.tapdata.mongodb.writer.error.handler.Code11000Handler;
 import io.tapdata.mongodb.writer.error.handler.Code28Handler;
 import io.tapdata.mongodb.writer.error.handler.Code2Handler;
+import io.tapdata.mongodb.writer.error.handler.Code72Handler;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -14,6 +15,7 @@
  **/
 public enum BulkWriteErrorCodeHandlerEnum {
 	CODE_28(28, new Code28Handler()),
+	CODE_72(72, new Code72Handler()),
 	CODE_11000(11000, new Code11000Handler()),
 	CODE_2(2, new Code2Handler()),
 	;

connectors/mongodb-connector/src/main/java/io/tapdata/mongodb/writer/error/handler/Code72Handler.java

@@ -0,0 +1,46 @@
+package io.tapdata.mongodb.writer.error.handler;
+
+import com.mongodb.MongoBulkWriteException;
+import com.mongodb.bulk.BulkWriteError;
+import com.mongodb.client.MongoCollection;
+import com.mongodb.client.model.*;
+import io.tapdata.mongodb.writer.BulkWriteModel;
+import io.tapdata.mongodb.writer.error.BulkWriteErrorHandler;
+import org.bson.Document;
+
+/**
+ * @author Tapdata
+ * @Description Handle error code 72: Multi-update operations are not allowed when updating the shard key field
+ * This error occurs when using UpdateManyModel on a sharded collection where the update involves shard key fields.
+ * The solution is to convert UpdateManyModel to UpdateOneModel.
+ * @create 2026-03-27
+ **/
+public class Code72Handler implements BulkWriteErrorHandler {
+
+    @Override
+    public WriteModel<Document> handle(
+            BulkWriteModel bulkWriteModel,
+            WriteModel<Document> writeModel,
+            BulkWriteOptions bulkWriteOptions,
+            MongoBulkWriteException mongoBulkWriteException,
+            BulkWriteError writeError,
+            MongoCollection<Document> collection
+    ) {
+        // Convert UpdateManyModel to UpdateOneModel to avoid shard key update restriction
+        try {
+            if (writeModel instanceof UpdateManyModel) {
+                UpdateManyModel<Document> updateManyModel = (UpdateManyModel<Document>) writeModel;
+                Document filter = (Document) updateManyModel.getFilter();
+                Document update = (Document) updateManyModel.getUpdate();
+                UpdateOptions options = updateManyModel.getOptions();
+                
+                // Convert to UpdateOneModel with the same filter, update and options
+                return new UpdateOneModel<>(filter, update, options);
+            }
+        } catch (Throwable ignored) {
+            // fall through to return null (can't handle)
+        }
+        return null;
+    }
+}
+