diff --git a/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Pipeline.java b/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Pipeline.java index c49ee0ad2f9a..bedf37b64f39 100644 --- a/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Pipeline.java +++ b/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Pipeline.java @@ -47,6 +47,7 @@ import com.google.cloud.firestore.pipeline.stages.Distinct; import com.google.cloud.firestore.pipeline.stages.FindNearest; import com.google.cloud.firestore.pipeline.stages.FindNearestOptions; +import com.google.cloud.firestore.pipeline.stages.Insert; import com.google.cloud.firestore.pipeline.stages.Limit; import com.google.cloud.firestore.pipeline.stages.Offset; import com.google.cloud.firestore.pipeline.stages.PipelineExecuteOptions; @@ -63,6 +64,7 @@ import com.google.cloud.firestore.pipeline.stages.Unnest; import com.google.cloud.firestore.pipeline.stages.UnnestOptions; import com.google.cloud.firestore.pipeline.stages.Update; +import com.google.cloud.firestore.pipeline.stages.Upsert; import com.google.cloud.firestore.pipeline.stages.Where; import com.google.cloud.firestore.telemetry.MetricsUtil.MetricsContext; import com.google.cloud.firestore.telemetry.TelemetryConstants; @@ -76,6 +78,7 @@ import com.google.firestore.v1.ExecutePipelineRequest; import com.google.firestore.v1.ExecutePipelineResponse; import com.google.firestore.v1.StructuredPipeline; +import com.google.firestore.v1.TransactionOptions; import com.google.firestore.v1.Value; import com.google.protobuf.ByteString; import java.util.ArrayList; @@ -1306,6 +1309,21 @@ public Pipeline update(Update update) { return append(update); } + @BetaApi + public Pipeline insert(Insert insert) { + return append(insert); + } + + @BetaApi + public Pipeline upsert(Upsert upsert) { + return append(upsert); + } + + @BetaApi + public Pipeline upsert(Selectable... transformedFields) { + return append(new Upsert(transformedFields)); + } + /** * Performs an insert operation using documents from previous stages. Adds a generic stage to the * pipeline. @@ -1526,6 +1544,12 @@ void executeInternal( if (transactionId != null) { request.setTransaction(transactionId); + } else if (options.isAtomic()) { + request.setNewTransaction( + TransactionOptions.newBuilder() + .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance()) + .build()); + request.setAutoCommitTransaction(true); } if (readTime != null) { diff --git a/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/Insert.java b/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/Insert.java new file mode 100644 index 000000000000..e8ad81e8ba08 --- /dev/null +++ b/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/Insert.java @@ -0,0 +1,78 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed 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 com.google.cloud.firestore.pipeline.stages; + +import com.google.api.core.BetaApi; +import com.google.api.core.InternalApi; +import com.google.cloud.firestore.PipelineUtils; +import com.google.cloud.firestore.pipeline.expressions.Expression; +import com.google.firestore.v1.Value; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Map; +import javax.annotation.Nullable; + +@InternalApi +public final class Insert extends Stage { + + @Nullable private final String collectionPath; + @Nullable private final Expression documentIdExpr; + + private Insert( + @Nullable String collectionPath, + @Nullable Expression documentIdExpr, + InternalOptions options) { + super("insert", buildOptions(collectionPath, documentIdExpr, options)); + this.collectionPath = collectionPath; + this.documentIdExpr = documentIdExpr; + } + + @BetaApi + public Insert() { + this(null, null, InternalOptions.EMPTY); + } + + @BetaApi + public Insert withCollection(String collectionPath) { + return new Insert(collectionPath, this.documentIdExpr, this.options); + } + + @BetaApi + public Insert withDocumentId(Expression documentIdExpr) { + return new Insert(this.collectionPath, documentIdExpr, this.options); + } + + private static InternalOptions buildOptions( + @Nullable String collectionPath, + @Nullable Expression documentIdExpr, + InternalOptions baseOptions) { + Map optsMap = new HashMap<>(baseOptions.options); + if (collectionPath != null) { + String path = collectionPath.startsWith("/") ? collectionPath : "/" + collectionPath; + optsMap.put("collection", Value.newBuilder().setReferenceValue(path).build()); + } + if (documentIdExpr != null) { + optsMap.put("document_id", PipelineUtils.encodeValue(documentIdExpr)); + } + return new InternalOptions(optsMap); + } + + @Override + Iterable toStageArgs() { + return new ArrayList<>(); + } +} diff --git a/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/PipelineExecuteOptions.java b/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/PipelineExecuteOptions.java index 754215e675f2..5dc947e80a45 100644 --- a/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/PipelineExecuteOptions.java +++ b/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/PipelineExecuteOptions.java @@ -38,4 +38,13 @@ public PipelineExecuteOptions withExplainOptions(ExplainOptions options) { public PipelineExecuteOptions withIndexMode(String indexMode) { return with("index_mode", indexMode); } + + public PipelineExecuteOptions withAtomic(boolean atomic) { + return with("atomic", atomic); + } + + boolean isAtomic() { + return options.options.containsKey("atomic") + && options.options.get("atomic").getBooleanValue(); + } } diff --git a/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/Upsert.java b/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/Upsert.java new file mode 100644 index 000000000000..928c758fbd0d --- /dev/null +++ b/java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/Upsert.java @@ -0,0 +1,97 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed 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 com.google.cloud.firestore.pipeline.stages; + +import com.google.api.core.BetaApi; +import com.google.api.core.InternalApi; +import com.google.cloud.firestore.PipelineUtils; +import com.google.cloud.firestore.pipeline.expressions.Expression; +import com.google.cloud.firestore.pipeline.expressions.Selectable; +import com.google.firestore.v1.Value; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; + +@InternalApi +public final class Upsert extends Stage { + + @Nullable private final Selectable[] transformedFields; + @Nullable private final String collectionPath; + @Nullable private final Expression documentIdExpr; + + private Upsert( + @Nullable Selectable[] transformedFields, + @Nullable String collectionPath, + @Nullable Expression documentIdExpr, + InternalOptions options) { + super("upsert", buildOptions(collectionPath, documentIdExpr, options)); + this.transformedFields = transformedFields; + this.collectionPath = collectionPath; + this.documentIdExpr = documentIdExpr; + } + + @BetaApi + public Upsert() { + this(null, null, null, InternalOptions.EMPTY); + } + + @BetaApi + public Upsert(Selectable... transformedFields) { + this(transformedFields, null, null, InternalOptions.EMPTY); + } + + @BetaApi + public Upsert withCollection(String collectionPath) { + return new Upsert(this.transformedFields, collectionPath, this.documentIdExpr, this.options); + } + + @BetaApi + public Upsert withDocumentId(Expression documentIdExpr) { + return new Upsert(this.transformedFields, this.collectionPath, documentIdExpr, this.options); + } + + private static InternalOptions buildOptions( + @Nullable String collectionPath, + @Nullable Expression documentIdExpr, + InternalOptions baseOptions) { + Map optsMap = new HashMap<>(baseOptions.options); + if (collectionPath != null) { + String path = collectionPath.startsWith("/") ? collectionPath : "/" + collectionPath; + optsMap.put("collection", Value.newBuilder().setReferenceValue(path).build()); + } + if (documentIdExpr != null) { + optsMap.put("document_id", PipelineUtils.encodeValue(documentIdExpr)); + } + return new InternalOptions(optsMap); + } + + @Override + Iterable toStageArgs() { + List args = new ArrayList<>(); + if (transformedFields != null && transformedFields.length > 0) { + Map map = PipelineUtils.selectablesToMap(transformedFields); + Map encodedMap = new HashMap<>(); + for (Map.Entry entry : map.entrySet()) { + encodedMap.put(entry.getKey(), PipelineUtils.encodeValue(entry.getValue())); + } + args.add(PipelineUtils.encodeValue(encodedMap)); + } + return args; + } +} diff --git a/java-firestore/google-cloud-firestore/src/test/java/com/google/cloud/firestore/PipelineProtoTest.java b/java-firestore/google-cloud-firestore/src/test/java/com/google/cloud/firestore/PipelineProtoTest.java index f62e4ef83f63..26ef6c367375 100644 --- a/java-firestore/google-cloud-firestore/src/test/java/com/google/cloud/firestore/PipelineProtoTest.java +++ b/java-firestore/google-cloud-firestore/src/test/java/com/google/cloud/firestore/PipelineProtoTest.java @@ -103,4 +103,75 @@ public void testSearchStageProtoEncoding() { Value addFields = optionsMap.get("add_fields"); assertThat(addFields.getMapValue().getFieldsMap().get("bar").getBooleanValue()).isTrue(); } + + @Test + public void testInsertStageProtoEncoding() { + FirestoreOptions options = + FirestoreOptions.newBuilder() + .setProjectId("new-project") + .setDatabaseId("(default)") + .build(); + Firestore firestore = options.getService(); + + java.util.Map data = new java.util.HashMap<>(); + data.put("title", "Test Book"); + + Pipeline pipeline = + firestore + .pipeline() + .literals(data) + .insert( + new com.google.cloud.firestore.pipeline.stages.Insert() + .withCollection("books") + .withDocumentId(constant("book1"))); + + com.google.firestore.v1.Pipeline protoPipeline = pipeline.toProto(); + assertThat(protoPipeline.getStagesCount()).isEqualTo(2); + + Stage insertStage = protoPipeline.getStages(1); + assertThat(insertStage.getName()).isEqualTo("insert"); + assertThat(insertStage.getArgsCount()).isEqualTo(0); + + java.util.Map optionsMap = insertStage.getOptionsMap(); + assertThat(optionsMap.get("collection").getReferenceValue()).isEqualTo("/books"); + assertThat(optionsMap.get("document_id").getStringValue()).isEqualTo("book1"); + } + + @Test + public void testUpsertStageProtoEncoding() { + FirestoreOptions options = + FirestoreOptions.newBuilder() + .setProjectId("new-project") + .setDatabaseId("(default)") + .build(); + Firestore firestore = options.getService(); + + java.util.Map data = new java.util.HashMap<>(); + data.put("title", "Upsert Book"); + data.put("count", 1); + + Pipeline pipeline = + firestore + .pipeline() + .literals(data) + .upsert( + new com.google.cloud.firestore.pipeline.stages.Upsert( + com.google.cloud.firestore.pipeline.expressions.Expression.add( + field("count"), constant(1)) + .as("count")) + .withCollection("books") + .withDocumentId(constant("book1"))); + + com.google.firestore.v1.Pipeline protoPipeline = pipeline.toProto(); + assertThat(protoPipeline.getStagesCount()).isEqualTo(2); + + Stage upsertStage = protoPipeline.getStages(1); + assertThat(upsertStage.getName()).isEqualTo("upsert"); + assertThat(upsertStage.getArgsCount()).isEqualTo(1); + assertThat(upsertStage.getArgs(0).getMapValue().getFieldsMap()).containsKey("count"); + + java.util.Map optionsMap = upsertStage.getOptionsMap(); + assertThat(optionsMap.get("collection").getReferenceValue()).isEqualTo("/books"); + assertThat(optionsMap.get("document_id").getStringValue()).isEqualTo("book1"); + } } diff --git a/java-firestore/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITPipelineTest.java b/java-firestore/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITPipelineTest.java index 170859f5962e..484589676e84 100644 --- a/java-firestore/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITPipelineTest.java +++ b/java-firestore/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITPipelineTest.java @@ -136,11 +136,13 @@ import com.google.cloud.firestore.pipeline.stages.ExplainOptions; import com.google.cloud.firestore.pipeline.stages.FindNearest; import com.google.cloud.firestore.pipeline.stages.FindNearestOptions; +import com.google.cloud.firestore.pipeline.stages.Insert; import com.google.cloud.firestore.pipeline.stages.PipelineExecuteOptions; import com.google.cloud.firestore.pipeline.stages.RawOptions; import com.google.cloud.firestore.pipeline.stages.RawStage; import com.google.cloud.firestore.pipeline.stages.Sample; import com.google.cloud.firestore.pipeline.stages.UnnestOptions; +import com.google.cloud.firestore.pipeline.stages.Upsert; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Lists; @@ -4737,4 +4739,102 @@ public void testLiteralsWithExpressions() throws Exception { assertThat(results.get(0).getData().get("base")).isEqualTo(10L); assertThat(results.get(0).getData().get("doubled")).isEqualTo(20L); } + + @Test + public void testInsertStage() throws Exception { + CollectionReference dmlCol = firestore.collection(LocalFirestoreHelper.autoId()); + java.util.Map data = new java.util.HashMap<>(); + data.put("title", "New Book"); + data.put("author", "Author 1"); + + List results = + firestore + .pipeline() + .literals(data) + .insert(new Insert().withCollection(dmlCol.getPath()).withDocumentId(constant("newBook_insert_1"))) + .execute() + .get() + .getResults(); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getData().get("documents_modified")).isEqualTo(1L); + + DocumentSnapshot snap = dmlCol.document("newBook_insert_1").get().get(); + assertThat(snap.exists()).isTrue(); + assertThat(snap.get("title")).isEqualTo("New Book"); + } + + @Test + public void testUpsertStageWithTransforms() throws Exception { + CollectionReference dmlCol = firestore.collection(LocalFirestoreHelper.autoId()); + java.util.Map data = new java.util.HashMap<>(); + data.put("title", "Upserted Book"); + data.put("count", 1); + + List results = + firestore + .pipeline() + .literals(data) + .upsert( + new Upsert(add(field("count"), constant(1)).as("count")) + .withCollection(dmlCol.getPath()) + .withDocumentId(constant("upsertBook_1"))) + .execute(new PipelineExecuteOptions().withAtomic(true)) + .get() + .getResults(); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getData().get("documents_modified")).isEqualTo(1L); + + DocumentSnapshot snap = dmlCol.document("upsertBook_1").get().get(); + assertThat(snap.exists()).isTrue(); + assertThat(snap.get("title")).isEqualTo("Upserted Book"); + } + + @Test + public void testDMLWithAtomicOption() throws Exception { + CollectionReference dmlCol = firestore.collection(LocalFirestoreHelper.autoId()); + java.util.Map data = new java.util.HashMap<>(); + data.put("title", "Atomic Book"); + + List results = + firestore + .pipeline() + .literals(data) + .insert(new Insert().withCollection(dmlCol.getPath()).withDocumentId(constant("atomicBook_1"))) + .execute(new PipelineExecuteOptions().withAtomic(true)) + .get() + .getResults(); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getData().get("documents_modified")).isEqualTo(1L); + + DocumentSnapshot snap = dmlCol.document("atomicBook_1").get().get(); + assertThat(snap.exists()).isTrue(); + } + + @Test + public void testInsertUpsertStagesInsideTransaction() throws Exception { + CollectionReference dmlCol = firestore.collection(LocalFirestoreHelper.autoId()); + firestore + .runTransaction( + transaction -> { + java.util.Map data = new java.util.HashMap<>(); + data.put("title", "Tx Book"); + Pipeline insertPpl = + firestore + .pipeline() + .literals(data) + .insert(new Insert().withCollection(dmlCol.getPath()).withDocumentId(constant("txBook_1"))); + List res = transaction.execute(insertPpl).get().getResults(); + assertThat(res).hasSize(1); + assertThat(res.get(0).getData().get("documents_modified")).isEqualTo(1L); + return null; + }) + .get(); + + DocumentSnapshot snap = dmlCol.document("txBook_1").get().get(); + assertThat(snap.exists()).isTrue(); + assertThat(snap.get("title")).isEqualTo("Tx Book"); + } }