Merge pull request #7537 from googleapis/dl/dml

feat: [firestore] Add support for DML

Author: wu-hui

handwritten/firestore/dev/src/pipelines/pipelines.ts

@@ -52,6 +52,7 @@ import {isOptionalEqual, isPlainObject} from '../util';
 import {
   AggregateFunction,
   AliasedAggregate,
+  AliasedExpression,
   Expression,
   Field,
   BooleanExpression,
@@ -81,6 +82,7 @@ import {
   Sample,
   Union,
   Unnest,
+  DeleteStage,
   InternalWhereStageOptions,
   InternalOffsetStageOptions,
   InternalLimitStageOptions,
@@ -95,6 +97,7 @@ import {
   InternalDocumentsStageOptions,
   InternalCollectionGroupStageOptions,
   InternalCollectionStageOptions,
+  UpdateStage,
 } from './stage';
 import {StructuredPipeline} from './structured-pipeline';
 import Selectable = FirebaseFirestore.Pipelines.Selectable;
@@ -1486,6 +1489,42 @@ export class Pipeline implements firestore.Pipelines.Pipeline {
     return this._addStage(new Sort(internalOptions));
   }
 
+  /**
+   * @beta
+   * Performs a delete operation on documents from previous stages.
+   *
+   * @example
+   * ```typescript
+   * // Deletes all documents in the "books" collection.
+   * firestore.pipeline().collection("books")
+   *    .delete();
+   * ```
+   *
+   * @return A new {@code Pipeline} object with this stage appended to the stage list.
+   */
+  delete(): Pipeline {
+    return this._addStage(new DeleteStage());
+  }
+
+  /**
+   * @beta
+   * Performs an update operation using documents from previous stages.
+   *
+   * @return A new {@code Pipeline} object with this stage appended to the stage list.
+   */
+  update(): Pipeline;
+  /**
+   * @beta
+   * Performs an update operation using documents from previous stages.
+   *
+   * @param transformedFields - The list of transformations to apply.
+   * @return A new {@code Pipeline} object with this stage appended to the stage list.
+   */
+  update(transformedFields: AliasedExpression[]): Pipeline;
+  update(transformedFields?: AliasedExpression[]): Pipeline {
+    return this._addStage(new UpdateStage(transformedFields));
+  }
+
   /**
    * @beta
    * Adds a raw stage to the pipeline.

handwritten/firestore/dev/src/pipelines/stage.ts

@@ -21,6 +21,7 @@ import {ProtoSerializable, Serializer} from '../serializer';
 
 import {
   AggregateFunction,
+  AliasedExpression,
   BooleanExpression,
   Expression,
   Field,
@@ -29,7 +30,7 @@ import {
 } from './expression';
 import {OptionsUtil} from './options-util';
 import {CollectionReference} from '../reference/collection-reference';
-import {validateUserDataHelper} from './pipeline-util';
+import {validateUserDataHelper, selectablesToMap} from './pipeline-util';
 import {Pipeline} from './pipelines';
 
 /**
@@ -759,3 +760,58 @@ export class RawStage implements Stage {
     validateUserDataHelper(this.params, ignoreUndefinedProperties);
   }
 }
+
+/**
+ * Delete stage.
+ */
+export class DeleteStage implements Stage {
+  name = 'delete';
+  readonly optionsUtil = new OptionsUtil({});
+
+  constructor() {}
+
+  _toProto(serializer: Serializer): api.Pipeline.IStage {
+    const args: api.IValue[] = [];
+
+    return {
+      name: this.name,
+      args,
+      options: this.optionsUtil.getOptionsProto(serializer, {}, {}),
+    };
+  }
+
+  _validateUserData(ignoreUndefinedProperties: boolean): void {}
+}
+
+/**
+ * Update stage.
+ */
+export class UpdateStage implements Stage {
+  name = 'update';
+  readonly optionsUtil = new OptionsUtil({});
+
+  constructor(private transformedFields?: AliasedExpression[]) {}
+
+  _toProto(serializer: Serializer): api.Pipeline.IStage {
+    const args: api.IValue[] = [];
+
+    if (this.transformedFields && this.transformedFields.length > 0) {
+      const mapped = selectablesToMap(this.transformedFields);
+      args.push(serializer.encodeValue(mapped)!);
+    } else {
+      args.push(serializer.encodeValue(new Map())!);
+    }
+
+    return {
+      name: this.name,
+      args,
+      options: this.optionsUtil.getOptionsProto(serializer, {}, {}),
+    };
+  }
+
+  _validateUserData(ignoreUndefinedProperties: boolean): void {
+    if (this.transformedFields) {
+      validateUserDataHelper(this.transformedFields, ignoreUndefinedProperties);
+    }
+  }
+}

handwritten/firestore/dev/system-test/pipeline.ts

@@ -189,16 +189,19 @@ describe.skipClassic('Pipeline class', () => {
   let beginDocCreation = 0;
   let endDocCreation = 0;
 
-  async function testCollectionWithDocs(docs: {
-    [id: string]: DocumentData;
-  }): Promise<CollectionReference<DocumentData>> {
+  async function testCollectionWithDocs(
+    targetCol: CollectionReference,
+    docs: {
+      [id: string]: DocumentData;
+    },
+  ): Promise<CollectionReference<DocumentData>> {
     beginDocCreation = new Date().valueOf();
     for (const id in docs) {
-      const ref = randomCol.doc(id);
+      const ref = targetCol.doc(id);
       await ref.set(docs[id]);
     }
     endDocCreation = new Date().valueOf();
-    return randomCol;
+    return targetCol;
   }
 
   function expectResults(result: PipelineSnapshot, ...docs: string[]): void;
@@ -224,7 +227,9 @@ describe.skipClassic('Pipeline class', () => {
     }
   }
 
-  async function setupBookDocs(): Promise<CollectionReference<DocumentData>> {
+  async function setupBookDocs(
+    targetCol: CollectionReference,
+  ): Promise<CollectionReference<DocumentData>> {
     const bookDocs: {[id: string]: DocumentData} = {
       book1: {
         title: "The Hitchhiker's Guide to the Galaxy",
@@ -334,18 +339,153 @@ describe.skipClassic('Pipeline class', () => {
         embedding: FieldValue.vector([1, 1, 1, 1, 1, 1, 1, 1, 1, 10]),
       },
     };
-    return testCollectionWithDocs(bookDocs);
+    return testCollectionWithDocs(targetCol, bookDocs);
   }
 
   before(async () => {
     randomCol = getTestRoot();
-    await setupBookDocs();
+    await setupBookDocs(randomCol);
     firestore = randomCol.firestore;
   });
 
   afterEach(() => verifyInstance(firestore as unknown as InternalFirestore));
 
   describe('pipeline results', () => {
+    describe('DML stages', () => {
+      let dmlCol: CollectionReference;
+
+      beforeEach(async () => {
+        dmlCol = getTestRoot();
+        await setupBookDocs(dmlCol);
+      });
+
+      it('can execute delete stage multiple documents', async () => {
+        const deletePpl = firestore
+          .pipeline()
+          .collection(dmlCol.path)
+          .where(equal(field('genre'), 'Science Fiction'))
+          .delete();
+
+        const promise = deletePpl.execute();
+
+        if (process.env.FIRESTORE_TARGET_BACKEND?.toUpperCase() === 'NIGHTLY') {
+          const deleteRes = await promise;
+          expectResults(deleteRes, {documents_modified: 2});
+
+          const docSnap1 = await dmlCol.doc('book1').get();
+          expect(docSnap1.exists).to.be.false;
+
+          const docSnap10 = await dmlCol.doc('book10').get();
+          expect(docSnap10.exists).to.be.false;
+        } else {
+          await expect(promise).to.be.rejected;
+        }
+      });
+
+      it('can execute delete stage within a transaction', async () => {
+        const promise = firestore.runTransaction(async transaction => {
+          const deletePpl = firestore
+            .pipeline()
+            .collection(dmlCol.path)
+            .where(equal(field('__name__').documentId(), 'book2'))
+            .delete();
+
+          const deleteRes = await transaction.execute(deletePpl);
+          expectResults(deleteRes, {documents_modified: 1});
+        });
+
+        if (process.env.FIRESTORE_TARGET_BACKEND?.toUpperCase() === 'NIGHTLY') {
+          await promise;
+          const docSnap = await dmlCol.doc('book2').get();
+          expect(docSnap.exists).to.be.false;
+        } else {
+          await expect(promise).to.be.rejected;
+        }
+      });
+
+      it('can execute update stage with addFields', async () => {
+        const ppl = firestore
+          .pipeline()
+          .collection(dmlCol.path)
+          .where(equal(field('__name__').documentId(), 'book3'))
+          .addFields(field('__name__').documentId().as('id'))
+          .update([constant('baz').as('foo')]);
+
+        const promise = ppl.execute();
+
+        if (process.env.FIRESTORE_TARGET_BACKEND?.toUpperCase() === 'NIGHTLY') {
+          const res = await promise;
+          expectResults(res, {documents_modified: 1});
+
+          const docSnap = await dmlCol.doc('book3').get();
+          expect(docSnap.get('foo')).to.equal('baz');
+          expect(docSnap.get('id')).to.equal('book3');
+        } else {
+          await expect(promise).to.be.rejected;
+        }
+      });
+
+      it('can update multiple documents and remove fields', async () => {
+        const promise = firestore
+          .pipeline()
+          .collection(dmlCol.path)
+          .where(equal(field('genre'), 'Science Fiction'))
+          .removeFields('awards')
+          .update([constant('Updated').as('status')])
+          .execute();
+
+        if (process.env.FIRESTORE_TARGET_BACKEND?.toUpperCase() === 'NIGHTLY') {
+          const res = await promise;
+          expectResults(res, {documents_modified: 2});
+
+          const docSnap1 = await dmlCol.doc('book1').get();
+          expect(docSnap1.get('status')).to.equal('Updated');
+          expect(docSnap1.get('awards')).to.be.undefined;
+
+          const docSnap10 = await dmlCol.doc('book10').get();
+          expect(docSnap10.get('status')).to.equal('Updated');
+          expect(docSnap10.get('awards')).to.be.undefined;
+        } else {
+          await expect(promise).to.be.rejected;
+        }
+      });
+
+      it('can update with expressions', async () => {
+        const promise = firestore
+          .pipeline()
+          .collection(dmlCol.path)
+          .where(equal(field('__name__').documentId(), 'book1'))
+          .update([add(field('rating'), constant(1.0)).as('rating')])
+          .execute();
+
+        if (process.env.FIRESTORE_TARGET_BACKEND?.toUpperCase() === 'NIGHTLY') {
+          const res = await promise;
+          expectResults(res, {documents_modified: 1});
+
+          const docSnap = await dmlCol.doc('book1').get();
+          expect(docSnap.get('rating')).to.equal(5.2);
+        } else {
+          await expect(promise).to.be.rejected;
+        }
+      });
+
+      it('can update non existing document modifies zero documents', async () => {
+        const nonExistingId = 'nonExistingId_123';
+        const promise = firestore
+          .pipeline()
+          .documents([dmlCol.doc(nonExistingId)])
+          .update([constant('Updated').as('status')])
+          .execute();
+
+        if (process.env.FIRESTORE_TARGET_BACKEND?.toUpperCase() === 'NIGHTLY') {
+          const res = await promise;
+          expectResults(res, {documents_modified: 0});
+        } else {
+          await expect(promise).to.be.rejected;
+        }
+      });
+    });
+
     it('empty snapshot as expected', async () => {
       const snapshot = await firestore
         .pipeline()

handwritten/firestore/types/firestore.d.ts

@@ -12314,6 +12314,28 @@ declare namespace FirebaseFirestore {
        * @returns A new `Pipeline` object with this select stage appended to its list of stages.
        */
       select(options: SelectStageOptions): Pipeline;
+      /**
+       * @beta
+       * Performs a delete operation on documents from previous stages.
+       *
+       * @return A new {@link Pipeline} object with this stage appended to the stage list.
+       */
+      delete(): Pipeline;
+      /**
+       * @beta
+       * Performs an update operation using documents from previous stages.
+       *
+       * @return A new {@link Pipeline} object with this stage appended to the stage list.
+       */
+      update(): Pipeline;
+      /**
+       * @beta
+       * Performs an update operation using documents from previous stages.
+       *
+       * @param transformedFields - The list of transformations to apply.
+       * @return A new {@link Pipeline} object with this stage appended to the stage list.
+       */
+      update(transformedFields: AliasedExpression[]): Pipeline;
       /**
        * @beta
        * Filters the documents from previous stages to only include those matching the specified {@link
@@ -13178,6 +13200,7 @@ declare namespace FirebaseFirestore {
        */
       docs: Array<string | DocumentReference>;
     };
+
     /**
      * @beta
      * Options defining how an AddFieldsStage is evaluated. See {@link Pipeline.addFields}.