Add missed fields to bulk's UpdateOperation: source (#947)

Signed-off-by: Andriy Redko <andriy.redko@aiven.io>
(cherry picked from commit a6c7dc9)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

CHANGELOG.md

@@ -4,6 +4,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ## [Unreleased 2.x]
 ### Added
 - Add missed fields to PhraseSuggestOption: collapseMatch ([#940](https://github.com/opensearch-project/opensearch-java/pull/940))
+- Add missed fields to bulk's UpdateOperation: source ([#947](https://github.com/opensearch-project/opensearch-java/pull/947))
 
 ### Dependencies
 - Bumps `com.github.jk1.dependency-license-report` from 2.6 to 2.7

java-client/src/main/java/org/opensearch/client/opensearch/core/bulk/UpdateOperation.java

@@ -41,6 +41,7 @@
 import org.opensearch.client.json.JsonpSerializer;
 import org.opensearch.client.json.NdJsonpSerializable;
 import org.opensearch.client.opensearch._types.Script;
+import org.opensearch.client.opensearch.core.search.SourceConfig;
 import org.opensearch.client.util.ApiTypeHelper;
 import org.opensearch.client.util.ObjectBuilder;
 
@@ -156,6 +157,9 @@ public static class Builder<TDocument> extends BulkOperationBase.AbstractBuilder
         @Nullable
         private Script script;
 
+        @Nullable
+        private SourceConfig source;
+
         /**
          * API name: {@code document}
          */
@@ -220,6 +224,21 @@ public final Builder<TDocument> retryOnConflict(@Nullable Integer value) {
             return this;
         }
 
+        /**
+         * API name: {@code _source}
+         */
+        public final Builder<TDocument> source(@Nullable SourceConfig value) {
+            this.source = value;
+            return this;
+        }
+
+        /**
+         * API name: {@code _source}
+         */
+        public final Builder<TDocument> source(Function<SourceConfig.Builder, ObjectBuilder<SourceConfig>> fn) {
+            return this.source(fn.apply(new SourceConfig.Builder()).build());
+        }
+
         /**
          * Serializer for TDocument. If not set, an attempt will be made to find a
          * serializer from the JSON context.
@@ -249,6 +268,7 @@ public UpdateOperation<TDocument> build() {
                 .detectNoop(detectNoop)
                 .script(script)
                 .upsert(upsert)
+                .source(source)
                 .tDocumentSerializer(tDocumentSerializer)
                 .build();
 

java-client/src/main/java/org/opensearch/client/opensearch/core/bulk/UpdateOperationData.java

@@ -9,12 +9,14 @@
 package org.opensearch.client.opensearch.core.bulk;
 
 import jakarta.json.stream.JsonGenerator;
+import java.util.function.Function;
 import javax.annotation.Nullable;
 import org.opensearch.client.json.JsonpMapper;
 import org.opensearch.client.json.JsonpSerializable;
 import org.opensearch.client.json.JsonpSerializer;
 import org.opensearch.client.json.JsonpUtils;
 import org.opensearch.client.opensearch._types.Script;
+import org.opensearch.client.opensearch.core.search.SourceConfig;
 import org.opensearch.client.util.ObjectBuilder;
 
 public class UpdateOperationData<TDocument> implements JsonpSerializable {
@@ -39,6 +41,9 @@ public class UpdateOperationData<TDocument> implements JsonpSerializable {
     @Nullable
     private final JsonpSerializer<TDocument> tDocumentSerializer;
 
+    @Nullable
+    private final SourceConfig source;
+
     private UpdateOperationData(Builder<TDocument> builder) {
         this.document = builder.document;
         this.docAsUpsert = builder.docAsUpsert;
@@ -47,7 +52,7 @@ private UpdateOperationData(Builder<TDocument> builder) {
         this.script = builder.script;
         this.upsert = builder.upsert;
         this.tDocumentSerializer = builder.tDocumentSerializer;
-
+        this.source = builder.source;
     }
 
     @Override
@@ -87,6 +92,11 @@ protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) {
             generator.writeKey("script");
             this.script.serialize(generator, mapper);
         }
+
+        if (this.source != null) {
+            generator.writeKey("_source");
+            this.source.serialize(generator, mapper);
+        }
     }
 
     /**
@@ -117,6 +127,9 @@ public static class Builder<TDocument> extends BulkOperationBase.AbstractBuilder
         @Nullable
         private Script script;
 
+        @Nullable
+        private SourceConfig source;
+
         /**
          * API name: {@code document}
          */
@@ -165,6 +178,21 @@ public final Builder<TDocument> script(@Nullable Script value) {
             return this;
         }
 
+        /**
+         * API name: {@code _source}
+         */
+        public final Builder<TDocument> source(@Nullable SourceConfig value) {
+            this.source = value;
+            return this;
+        }
+
+        /**
+         * API name: {@code _source}
+         */
+        public final Builder<TDocument> source(Function<SourceConfig.Builder, ObjectBuilder<SourceConfig>> fn) {
+            return this.source(fn.apply(new SourceConfig.Builder()).build());
+        }
+
         /**
          * Serializer for TDocument. If not set, an attempt will be made to find a
          * serializer from the JSON context.

java-client/src/test/java11/org/opensearch/client/opensearch/integTest/AbstractCrudIT.java

@@ -204,6 +204,7 @@ public void testUpdate() throws Exception {
             UpdateRequest<AppData, AppData> updateRequest = new UpdateRequest.Builder<AppData, AppData>().index("index")
                 .id("does_not_exist")
                 .doc(appData)
+                .source(s -> s.fetch(randomBoolean()))
                 .build();
             try {
                 javaClient().update(updateRequest, AppData.class);