feat(kits): add firestore-incremental-capture

kits/firestore-incremental-capture/package.json

@@ -1,8 +1,7 @@
 {
   "name": "@firebase/firestore-incremental-capture",
   "version": "0.1.0",
-  "private": true,
-  "description": "Incremental point-in-time capture of Firestore changes",
+  "description": "Incremental point-in-time capture of Firestore changes as a deployable Firebase Function",
   "license": "Apache-2.0",
   "main": "lib/index.js",
   "types": "lib/index.d.ts",
@@ -21,6 +20,18 @@
   },
   "scripts": {
     "build": "tsc -b",
-    "clean": "tsc -b --clean"
+    "clean": "tsc -b --clean",
+    "test": "vitest run",
+    "deploy": "pnpm build && firebase deploy --only functions",
+    "serve": "firebase emulators:start --only functions"
+  },
+  "dependencies": {
+    "@google-cloud/bigquery": "^7.6.0",
+    "@google-cloud/dataflow": "^3.2.0",
+    "firebase-admin": "^14.1.0",
+    "firebase-functions": "7.3.2"
+  },
+  "devDependencies": {
+    "vitest": "^3.2.4"
   }
 }

kits/firestore-incremental-capture/pipeline/.gitignore

@@ -0,0 +1 @@
+target/

kits/firestore-incremental-capture/pipeline/README.md

@@ -0,0 +1,23 @@
+## Debug the pipeline locally
+
+To debug this pipeline locally, use the `DirectRunner`:
+
+Note: If your Cloud Storage bucket was provisioned after September 30, 2024
+the default bucket name will be suffixed with `.firebasestorage.app` instead of `.appspot.com`
+
+```bash
+mvn compile exec:java \
+    -Dexec.mainClass=com.pipeline.RestorationPipeline \
+    -Dexec.args='--timestamp=1697740800 --firestoreCollectionId="test" --firestoreDb="test" --tempLocation="gs://PROJECT_ID.appspot.com" --project="PROJECT_ID"'
+```
+
+### Arguments
+
+- `timestamp`: The timestamp to restore the data to from a PITR, if it's further than 7 days in the past, it will be set to 7 days in the past. The timestamp is in UNIX seconds.
+- `firestoreCollectionId`: The collection to restore, use `*` if you want the full database.
+
+## Compile JAR to run on Dataflow
+
+```bash
+mvn clean package -DskipTests -Dexec.mainClass=com.pipeline.RestorationPipeline
+```

kits/firestore-incremental-capture/pipeline/pom.xml

@@ -0,0 +1,122 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+    xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+
+    <groupId>com.pipeline</groupId>
+    <artifactId>pipeline</artifactId>
+    <version>1.0</version>
+
+    <properties>
+        <maven.compiler.source>11</maven.compiler.source>
+        <maven.compiler.target>11</maven.compiler.target>
+        <beam.version>2.75.0</beam.version>
+    </properties>
+    <build>
+        <finalName>restore-firestore</finalName>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-shade-plugin</artifactId>
+                <version>3.2.4</version>
+                <executions>
+                    <execution>
+                        <phase>package</phase>
+                        <goals>
+                            <goal>shade</goal>
+                        </goals>
+                        <configuration>
+                            <createDependencyReducedPom>false</createDependencyReducedPom>
+                            <filters>
+                                <filter>
+                                    <artifact>*:*</artifact>
+                                    <excludes>
+                                        <exclude>META-INF/*.SF</exclude>
+                                        <exclude>META-INF/*.DSA</exclude>
+                                        <exclude>META-INF/*.RSA</exclude>
+                                    </excludes>
+                                </filter>
+                            </filters>
+                            <transformers>
+                                <transformer
+                                    implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer" />
+                            </transformers>
+                        </configuration>
+                    </execution>
+                </executions>
+            </plugin>
+            <plugin>
+                <groupId>org.codehaus.mojo</groupId>
+                <artifactId>exec-maven-plugin</artifactId>
+                <version>3.1.0</version>
+                <executions>
+                    <execution>
+                        <goals>
+                            <goal>java</goal>
+                        </goals>
+                    </execution>
+                </executions>
+                <configuration>
+                    <mainClass>com.pipeline.RestorationPipeline</mainClass>
+                </configuration>
+            </plugin>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-jar-plugin</artifactId>
+                <version>3.3.0</version>
+                <configuration>
+                    <archive>
+                        <manifest>
+                            <addClasspath>true</addClasspath>
+                            <classpathPrefix>lib/</classpathPrefix>
+                            <mainClass>com.pipeline.RestorationPipeline</mainClass>
+                        </manifest>
+                    </archive>
+                </configuration>
+            </plugin>
+        </plugins>
+    </build>
+
+    <dependencies>
+
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>slf4j-simple</artifactId>
+            <version>2.0.9</version>
+        </dependency>
+
+
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>slf4j-api</artifactId>
+            <version>2.0.9</version>
+        </dependency>
+
+        <!--
+        https://mvnrepository.com/artifact/org.apache.beam/beam-sdks-java-google-cloud-platform-bom -->
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-sdks-java-core</artifactId>
+            <version>${beam.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-runners-google-cloud-dataflow-java</artifactId>
+            <version>${beam.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-sdks-java-io-google-cloud-platform</artifactId>
+            <version>${beam.version}</version>
+        </dependency>
+        <!-- https://mvnrepository.com/artifact/org.apache.beam/beam-runners-direct-java -->
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-runners-direct-java</artifactId>
+            <version>${beam.version}</version>
+        </dependency>
+    </dependencies>
+
+
+</project>

kits/firestore-incremental-capture/pipeline/src/main/java/com/pipeline/FirestoreHelpers.java

@@ -0,0 +1,156 @@
+/**
+ * Copyright 2023 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
+ *
+ *    https://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.pipeline;
+
+import org.apache.beam.sdk.io.gcp.firestore.FirestoreIO;
+import org.apache.beam.sdk.io.gcp.firestore.FirestoreV1.BatchWriteWithSummary;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.google.firestore.v1.Document;
+import com.google.firestore.v1.RunQueryRequest;
+import com.google.firestore.v1.RunQueryResponse;
+import com.google.firestore.v1.StructuredQuery;
+import com.google.firestore.v1.StructuredQuery.CollectionSelector;
+import com.google.firestore.v1.Write;
+
+public class FirestoreHelpers {
+  public static final class RunQuery extends BasePTransform<String, RunQueryRequest> {
+    private static final Logger LOG = LoggerFactory.getLogger(Utils.class);
+
+    final String projectId;
+
+    public RunQuery(String projectId, String database) {
+      super("projects/" + projectId + "/databases/" + database + "/documents");
+      this.projectId = projectId;
+    }
+
+    @Override
+    public PCollection<RunQueryRequest> expand(PCollection<String> input) {
+      LOG.info(baseDocumentPath);
+      return input.apply(
+          ParDo.of(
+              new DoFn<String, RunQueryRequest>() {
+                @ProcessElement
+                public void processElement(ProcessContext c) {
+                  final String collectionId = c.element();
+
+                  if (collectionId.equals("*")) {
+                    LOG.info("Querying all collections");
+                    RunQueryRequest runQueryRequest = RunQueryRequest.newBuilder()
+                        .setParent(baseDocumentPath)
+                        .setStructuredQuery(StructuredQuery.newBuilder().build())
+                        .build();
+
+                    c.output(runQueryRequest);
+                    return;
+                  }
+
+                  CollectionSelector collection = CollectionSelector
+                      .newBuilder()
+                      .setCollectionId(collectionId)
+                      .build();
+
+                  RunQueryRequest runQueryRequest = RunQueryRequest.newBuilder()
+                      .setParent(baseDocumentPath)
+                      .setStructuredQuery(
+                          com.google.firestore.v1.StructuredQuery.newBuilder()
+                              .addFrom(collection)
+                              .build())
+                      .build();
+
+                  c.output(runQueryRequest);
+                }
+              }));
+    }
+  }
+
+  public static final class RunQueryResponseToDocument extends BasePTransform<RunQueryResponse, Document> {
+
+    public RunQueryResponseToDocument() {
+      super("");
+    }
+
+    @Override
+    public PCollection<Document> expand(PCollection<RunQueryResponse> input) {
+      return input.apply(
+          ParDo.of(
+              new DoFn<RunQueryResponse, Document>() {
+                @ProcessElement
+                public void processElement(ProcessContext c) {
+                  RunQueryResponse response = c.element();
+                  c.output(response.getDocument());
+                }
+              }));
+    }
+  }
+
+  public static final class DocumentToWrite extends BasePTransform<KV<String, Document>, Write> {
+
+    final String projectId;
+
+    public DocumentToWrite(String projectId, String database) {
+      super("projects/" + projectId + "/databases/" + database + "/documents");
+      this.projectId = projectId;
+    }
+
+    @Override
+    public PCollection<Write> expand(PCollection<KV<String, Document>> input) {
+      return input.apply(
+          ParDo.of(
+              new DoFn<KV<String, Document>, Write>() {
+                @ProcessElement
+                public void processElement(ProcessContext c) {
+                  String changeType = c.element().getKey();
+                  Document document = c.element().getValue();
+
+                  // LOG.info("STEP ONE >>>>>>> changeType: {}, documentName: {}", changeType,
+                  // document.getName());
+
+                  switch (changeType) {
+                    case "DELETE":
+                      c.output(Write.newBuilder()
+                          .setDelete(document.getName())
+                          .build());
+
+                      break;
+
+                    default:
+                      c.output(Write.newBuilder()
+                          .setUpdate(document)
+                          .build());
+                  }
+                }
+              }));
+    }
+  }
+
+  private abstract static class BasePTransform<InT, OutT>
+      extends PTransform<PCollection<InT>, PCollection<OutT>> {
+
+    protected final String baseDocumentPath;
+
+    private BasePTransform(String baseDocumentPath) {
+      this.baseDocumentPath = baseDocumentPath;
+    }
+  }
+}