feat(firestore): add DML stages, literals source, and atomic execution option to Android SDK pipelines

firebase-firestore/api.txt

@@ -433,13 +433,15 @@ package com.google.firebase.firestore {
     method public com.google.firebase.firestore.Pipeline aggregate(com.google.firebase.firestore.pipeline.AggregateStage aggregateStage, com.google.firebase.firestore.pipeline.AggregateOptions options);
     method public com.google.firebase.firestore.Pipeline aggregate(com.google.firebase.firestore.pipeline.AliasedAggregate accumulator, com.google.firebase.firestore.pipeline.AliasedAggregate... additionalAccumulators);
     method public com.google.firebase.firestore.Pipeline define(com.google.firebase.firestore.pipeline.AliasedExpression aliasedExpression, com.google.firebase.firestore.pipeline.AliasedExpression... additionalExpressions);
+    method public com.google.firebase.firestore.Pipeline delete();
     method public com.google.firebase.firestore.Pipeline distinct(com.google.firebase.firestore.pipeline.Selectable group, java.lang.Object... additionalGroups);
     method public com.google.firebase.firestore.Pipeline distinct(String groupField, java.lang.Object... additionalGroups);
     method public com.google.android.gms.tasks.Task<com.google.firebase.firestore.Pipeline.Snapshot> execute();
     method public com.google.android.gms.tasks.Task<com.google.firebase.firestore.Pipeline.Snapshot> execute(com.google.firebase.firestore.Pipeline.ExecuteOptions options);
     method public com.google.firebase.firestore.Pipeline findNearest(com.google.firebase.firestore.pipeline.Field vectorField, double[] vectorValue, com.google.firebase.firestore.pipeline.FindNearestStage.DistanceMeasure distanceMeasure);
     method public com.google.firebase.firestore.Pipeline findNearest(String vectorField, com.google.firebase.firestore.pipeline.Expression vectorValue, com.google.firebase.firestore.pipeline.FindNearestStage.DistanceMeasure distanceMeasure, com.google.firebase.firestore.pipeline.FindNearestOptions options);
     method public com.google.firebase.firestore.Pipeline findNearest(String vectorField, double[] vectorValue, com.google.firebase.firestore.pipeline.FindNearestStage.DistanceMeasure distanceMeasure);
+    method public com.google.firebase.firestore.Pipeline insert(String collectionPath, com.google.firebase.firestore.pipeline.Expression? documentIdExpr);
     method public com.google.firebase.firestore.Pipeline limit(int limit);
     method public com.google.firebase.firestore.Pipeline offset(int offset);
     method public com.google.firebase.firestore.Pipeline rawStage(com.google.firebase.firestore.pipeline.RawStage rawStage);
@@ -460,11 +462,14 @@ package com.google.firebase.firestore {
     method public com.google.firebase.firestore.Pipeline unnest(com.google.firebase.firestore.pipeline.Selectable arrayWithAlias, com.google.firebase.firestore.pipeline.UnnestOptions options);
     method public com.google.firebase.firestore.Pipeline unnest(com.google.firebase.firestore.pipeline.UnnestStage unnestStage);
     method public com.google.firebase.firestore.Pipeline unnest(String arrayField, String alias);
+    method public com.google.firebase.firestore.Pipeline update(com.google.firebase.firestore.pipeline.Selectable... fields);
+    method public com.google.firebase.firestore.Pipeline upsert(com.google.firebase.firestore.pipeline.Selectable[] transforms, String? collectionPath, com.google.firebase.firestore.pipeline.Expression? documentIdExpr);
     method public com.google.firebase.firestore.Pipeline where(com.google.firebase.firestore.pipeline.BooleanExpression condition);
   }
 
   public static final class Pipeline.ExecuteOptions extends com.google.firebase.firestore.pipeline.AbstractOptions<com.google.firebase.firestore.Pipeline.ExecuteOptions> {
     ctor public Pipeline.ExecuteOptions();
+    method public com.google.firebase.firestore.Pipeline.ExecuteOptions withAtomic(boolean atomic);
     method public com.google.firebase.firestore.Pipeline.ExecuteOptions withIndexMode(com.google.firebase.firestore.Pipeline.ExecuteOptions.IndexMode indexMode);
   }
 
@@ -508,6 +513,8 @@ package com.google.firebase.firestore {
     method public com.google.firebase.firestore.Pipeline database();
     method public com.google.firebase.firestore.Pipeline documents(com.google.firebase.firestore.DocumentReference... documents);
     method public com.google.firebase.firestore.Pipeline documents(java.lang.String... documents);
+    method public com.google.firebase.firestore.Pipeline literals(java.util.List<? extends java.util.Map<java.lang.String, ?>> data);
+    method public com.google.firebase.firestore.Pipeline literals(java.util.Map<java.lang.String, ?>... data);
     method public static com.google.firebase.firestore.Pipeline subcollection(com.google.firebase.firestore.pipeline.SubcollectionSource source);
     method public static com.google.firebase.firestore.Pipeline subcollection(String path);
     field public static final com.google.firebase.firestore.PipelineSource.Companion Companion;

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/PipelineTest.java

@@ -4290,4 +4290,57 @@ static <T> Map<String, T> mapOfEntries(Map.Entry<String, T>... entries) {
     }
     return Collections.unmodifiableMap(res);
   }
+
+  @Test
+  public void testDeleteStage() {
+    CollectionReference collection = testCollection();
+    Pipeline.Snapshot snapshot =
+        waitFor(collection.toPipeline().where(equal(field("__name__"), constant("book1"))).delete().execute());
+    assertThat(snapshot).isNotNull();
+  }
+
+  @Test
+  public void testUpdateStage() {
+    CollectionReference collection = testCollection();
+    Pipeline.Snapshot snapshot =
+        waitFor(
+            collection
+                .toPipeline()
+                .where(equal(field("__name__"), constant("book1")))
+                .update(constant("Updated").as("status"))
+                .execute());
+    assertThat(snapshot).isNotNull();
+  }
+
+  @Test
+  public void testInsertStage() {
+    CollectionReference collection = testCollection();
+    Map<String, Object> data = new HashMap<>();
+    data.put("title", "New Book");
+    Pipeline.Snapshot snapshot =
+        waitFor(
+            db.pipeline()
+                .literals(data)
+                .insert(collection.getPath(), constant("newBook_insert_1"))
+                .execute());
+    assertThat(snapshot).isNotNull();
+  }
+
+  @Test
+  public void testUpsertStage() {
+    CollectionReference collection = testCollection();
+    Map<String, Object> data = new HashMap<>();
+    data.put("title", "Upsert Book");
+    data.put("count", 1);
+    Pipeline.Snapshot snapshot =
+        waitFor(
+            db.pipeline()
+                .literals(data)
+                .upsert(
+                    add(field("count"), constant(1)).as("count"),
+                    collection.getPath(),
+                    constant("upsertBook_1"))
+                .execute(new Pipeline.ExecuteOptions().withAtomic(true)));
+    assertThat(snapshot).isNotNull();
+  }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/Pipeline.kt

@@ -36,6 +36,7 @@ import com.google.firebase.firestore.pipeline.CollectionSource
 import com.google.firebase.firestore.pipeline.CollectionSourceOptions
 import com.google.firebase.firestore.pipeline.DatabaseSource
 import com.google.firebase.firestore.pipeline.DefineStage
+import com.google.firebase.firestore.pipeline.DeleteStage
 import com.google.firebase.firestore.pipeline.DistinctStage
 import com.google.firebase.firestore.pipeline.DocumentsSource
 import com.google.firebase.firestore.pipeline.Expression
@@ -44,8 +45,10 @@ import com.google.firebase.firestore.pipeline.Field
 import com.google.firebase.firestore.pipeline.FindNearestOptions
 import com.google.firebase.firestore.pipeline.FindNearestStage
 import com.google.firebase.firestore.pipeline.FunctionExpression
+import com.google.firebase.firestore.pipeline.InsertStage
 import com.google.firebase.firestore.pipeline.InternalOptions
 import com.google.firebase.firestore.pipeline.LimitStage
+import com.google.firebase.firestore.pipeline.LiteralsSource
 import com.google.firebase.firestore.pipeline.OffsetStage
 import com.google.firebase.firestore.pipeline.Ordering
 import com.google.firebase.firestore.pipeline.RawStage
@@ -61,13 +64,16 @@ import com.google.firebase.firestore.pipeline.SubcollectionSource
 import com.google.firebase.firestore.pipeline.UnionStage
 import com.google.firebase.firestore.pipeline.UnnestOptions
 import com.google.firebase.firestore.pipeline.UnnestStage
+import com.google.firebase.firestore.pipeline.UpdateStage
+import com.google.firebase.firestore.pipeline.UpsertStage
 import com.google.firebase.firestore.pipeline.WhereStage
 import com.google.firebase.firestore.pipeline.evaluation.notImplemented
 import com.google.firebase.firestore.remote.RemoteSerializer
 import com.google.firebase.firestore.util.Logger
 import com.google.firestore.v1.ExecutePipelineRequest
 import com.google.firestore.v1.Pipeline as ProtoPipeline
 import com.google.firestore.v1.StructuredPipeline
+import com.google.firestore.v1.TransactionOptions
 import com.google.firestore.v1.Value
 
 /**
@@ -112,6 +118,8 @@ internal constructor(
     }
 
     fun withIndexMode(indexMode: IndexMode): ExecuteOptions = with("index_mode", indexMode.value)
+
+    fun withAtomic(atomic: Boolean): ExecuteOptions = with("atomic", atomic)
   }
 
   /**
@@ -176,6 +184,13 @@ internal constructor(
     val builder = ExecutePipelineRequest.newBuilder()
     builder.database = "projects/${database.projectId}/databases/${database.databaseId}"
     builder.structuredPipeline = toStructuredPipelineProto(options, firestore.userDataReader)
+    if (options != null && options.hasAtomic()) {
+      builder.newTransaction =
+        TransactionOptions.newBuilder()
+          .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance())
+          .build()
+      builder.autoCommitTransaction = true
+    }
     return builder.build()
   }
 
@@ -1107,11 +1122,32 @@ internal constructor(
    * @return A new `Pipeline` object with this stage appended to the stage list.
    */
   @Beta fun search(searchStage: SearchStage): Pipeline = append(searchStage)
+
+  fun delete(): Pipeline = append(DeleteStage())
+
+  fun update(vararg fields: Selectable): Pipeline = append(UpdateStage(fields))
+
+  fun insert(collectionPath: String, documentIdExpr: Expression? = null): Pipeline =
+    append(InsertStage(collectionPath, documentIdExpr))
+
+  fun upsert(
+    vararg transforms: Selectable,
+    collectionPath: String? = null,
+    documentIdExpr: Expression? = null
+  ): Pipeline = append(UpsertStage(transforms, collectionPath, documentIdExpr))
 }
 
 /** Start of a Firestore Pipeline */
 class PipelineSource internal constructor(private val firestore: FirebaseFirestore) {
 
+  /**
+   * Set the pipeline's source to literal document maps.
+   */
+  fun literals(vararg data: Map<String, Any?>): Pipeline = literals(data.toList())
+
+  fun literals(data: List<Map<String, Any?>>): Pipeline =
+    Pipeline(firestore, firestore.userDataReader, listOf(LiteralsSource(data)))
+
   /**
    * Convert the given Query into an equivalent Pipeline.
    *

firebase-firestore/src/main/java/com/google/firebase/firestore/pipeline/options.kt

@@ -64,6 +64,10 @@ internal constructor(private val options: ImmutableMap<String, Value>) {
     }
   }
 
+  internal fun hasAtomic(): Boolean {
+    return options.containsKey("atomic") && options["atomic"]?.booleanValue == true
+  }
+
   override fun equals(other: Any?): Boolean {
     if (this === other) return true
     if (other !is InternalOptions) return false

firebase-firestore/src/main/java/com/google/firebase/firestore/pipeline/stage.kt

@@ -1736,3 +1736,190 @@ internal constructor(
     return result
   }
 }
+
+internal class DeleteStage
+internal constructor(
+  options: InternalOptions = InternalOptions.EMPTY
+) : Stage<DeleteStage>("delete", options) {
+  override fun self(options: InternalOptions) = DeleteStage(options)
+  override fun canonicalId(): String = "delete()"
+  override fun args(userDataReader: UserDataReader): Sequence<Value> = emptySequence()
+
+  override fun equals(other: Any?): Boolean {
+    if (this === other) return true
+    if (other !is DeleteStage) return false
+    return options == other.options
+  }
+
+  override fun hashCode(): Int = options.hashCode()
+}
+
+internal class UpdateStage
+internal constructor(
+  private val fields: Array<out Selectable>,
+  options: InternalOptions = InternalOptions.EMPTY
+) : Stage<UpdateStage>("update", options) {
+  override fun self(options: InternalOptions) = UpdateStage(fields, options)
+  override fun canonicalId(): String = "update()"
+
+  override fun args(userDataReader: UserDataReader): Sequence<Value> {
+    return if (fields.isNotEmpty()) {
+      sequenceOf(encodeValue(associateWithoutDuplications(fields, userDataReader)))
+    } else {
+      sequenceOf(encodeValue(emptyMap<String, Value>()))
+    }
+  }
+
+  override fun equals(other: Any?): Boolean {
+    if (this === other) return true
+    if (other !is UpdateStage) return false
+    if (!fields.contentEquals(other.fields)) return false
+    return options == other.options
+  }
+
+  override fun hashCode(): Int {
+    var result = fields.contentHashCode()
+    result = 31 * result + options.hashCode()
+    return result
+  }
+}
+
+internal class InsertStage
+internal constructor(
+  internal val collectionPath: String?,
+  internal val documentIdExpr: Expression?,
+  options: InternalOptions = InternalOptions.EMPTY
+) : Stage<InsertStage>("insert", buildOptions(collectionPath, documentIdExpr, options)) {
+
+  override fun self(options: InternalOptions) = InsertStage(collectionPath, documentIdExpr, options)
+  override fun canonicalId(): String = "insert($collectionPath)"
+  override fun args(userDataReader: UserDataReader): Sequence<Value> = emptySequence()
+
+  override fun equals(other: Any?): Boolean {
+    if (this === other) return true
+    if (other !is InsertStage) return false
+    if (collectionPath != other.collectionPath) return false
+    if (documentIdExpr != other.documentIdExpr) return false
+    return options == other.options
+  }
+
+  override fun hashCode(): Int {
+    var result = collectionPath?.hashCode() ?: 0
+    result = 31 * result + (documentIdExpr?.hashCode() ?: 0)
+    result = 31 * result + options.hashCode()
+    return result
+  }
+
+  companion object {
+    private fun buildOptions(
+      collectionPath: String?,
+      documentIdExpr: Expression?,
+      baseOptions: InternalOptions
+    ): InternalOptions {
+      var opts = baseOptions
+      if (collectionPath != null) {
+        val path = if (collectionPath.startsWith("/")) collectionPath else "/$collectionPath"
+        opts = opts.with("collection", Value.newBuilder().setReferenceValue(path).build())
+      }
+      if (documentIdExpr != null) {
+        opts = opts.with("document_id", documentIdExpr.toProto())
+      }
+      return opts
+    }
+  }
+}
+
+internal class UpsertStage
+internal constructor(
+  private val fields: Array<out Selectable>,
+  internal val collectionPath: String?,
+  internal val documentIdExpr: Expression?,
+  options: InternalOptions = InternalOptions.EMPTY
+) : Stage<UpsertStage>("upsert", buildOptions(collectionPath, documentIdExpr, options)) {
+
+  override fun self(options: InternalOptions) =
+    UpsertStage(fields, collectionPath, documentIdExpr, options)
+  override fun canonicalId(): String = "upsert($collectionPath)"
+
+  override fun args(userDataReader: UserDataReader): Sequence<Value> {
+    return if (fields.isNotEmpty()) {
+      sequenceOf(encodeValue(associateWithoutDuplications(fields, userDataReader)))
+    } else {
+      emptySequence()
+    }
+  }
+
+  override fun equals(other: Any?): Boolean {
+    if (this === other) return true
+    if (other !is UpsertStage) return false
+    if (!fields.contentEquals(other.fields)) return false
+    if (collectionPath != other.collectionPath) return false
+    if (documentIdExpr != other.documentIdExpr) return false
+    return options == other.options
+  }
+
+  override fun hashCode(): Int {
+    var result = fields.contentHashCode()
+    result = 31 * result + (collectionPath?.hashCode() ?: 0)
+    result = 31 * result + (documentIdExpr?.hashCode() ?: 0)
+    result = 31 * result + options.hashCode()
+    return result
+  }
+
+  companion object {
+    private fun buildOptions(
+      collectionPath: String?,
+      documentIdExpr: Expression?,
+      baseOptions: InternalOptions
+    ): InternalOptions {
+      var opts = baseOptions
+      if (collectionPath != null) {
+        val path = if (collectionPath.startsWith("/")) collectionPath else "/$collectionPath"
+        opts = opts.with("collection", Value.newBuilder().setReferenceValue(path).build())
+      }
+      if (documentIdExpr != null) {
+        opts = opts.with("document_id", documentIdExpr.toProto())
+      }
+      return opts
+    }
+  }
+}
+
+class LiteralsSource
+internal constructor(
+  internal val data: List<Map<String, Any?>>,
+  options: InternalOptions = InternalOptions.EMPTY
+) : Stage<LiteralsSource>("literals", options) {
+
+  override fun self(options: InternalOptions) = LiteralsSource(data, options)
+  override fun canonicalId(): String = "literals()"
+
+  override fun args(userDataReader: UserDataReader): Sequence<Value> {
+    return data.asSequence().map { encodeLiteralMap(it, userDataReader) }
+  }
+
+  private fun encodeLiteralMap(map: Map<String, Any?>, userDataReader: UserDataReader): Value {
+    val mapValue = com.google.firestore.v1.MapValue.newBuilder()
+    for ((key, value) in map) {
+      when (value) {
+        is Expression -> mapValue.putFields(key, value.toProto())
+        is Map<*, *> -> @Suppress("UNCHECKED_CAST") mapValue.putFields(key, encodeLiteralMap(value as Map<String, Any?>, userDataReader))
+        else -> mapValue.putFields(key, Values.encodeValue(value))
+      }
+    }
+    return Value.newBuilder().setMapValue(mapValue).build()
+  }
+
+  override fun equals(other: Any?): Boolean {
+    if (this === other) return true
+    if (other !is LiteralsSource) return false
+    if (data != other.data) return false
+    return options == other.options
+  }
+
+  override fun hashCode(): Int {
+    var result = data.hashCode()
+    result = 31 * result + options.hashCode()
+    return result
+  }
+}