This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new d79a2487a87 normalize firestoreIO (#39941)
d79a2487a87 is described below
commit d79a2487a8745b882560f029486d347c69a024fe
Author: Abdelrahman Ibrahim <[email protected]>
AuthorDate: Wed Sep 9 03:00:07 2026 +0300
normalize firestoreIO (#39941)
* normalize firestoreIO
* move firestore to gcp_test extra
---
.../beam_PostCommit_Yaml_Xlang_Direct.json | 3 +-
.../beam_PostCommit_Yaml_Xlang_Direct.yml | 2 +-
.../workflows/beam_PreCommit_Yaml_Xlang_Direct.yml | 2 +-
.../FirestoreReadSchemaTransformConfiguration.java | 83 ++++++
.../FirestoreReadSchemaTransformProvider.java | 216 ++++++++++++++
.../beam/sdk/io/gcp/firestore/FirestoreUtils.java | 312 +++++++++++++++++++++
.../beam/sdk/io/gcp/firestore/FirestoreV1.java | 7 +-
.../sdk/io/gcp/firestore/FirestoreV1ReadFn.java | 38 ++-
...FirestoreWriteSchemaTransformConfiguration.java | 82 ++++++
.../FirestoreWriteSchemaTransformProvider.java | 201 +++++++++++++
.../FirestoreSchemaTransformProviderTest.java | 179 ++++++++++++
.../sdk/io/gcp/firestore/FirestoreUtilsTest.java | 65 +++++
.../yaml/extended_tests/databases/firestore.yaml | 84 ++++++
sdks/python/apache_beam/yaml/integration_tests.py | 54 ++++
sdks/python/apache_beam/yaml/standard_io.yaml | 27 ++
sdks/python/build.gradle | 8 +-
sdks/python/setup.py | 4 +
17 files changed, 1337 insertions(+), 30 deletions(-)
diff --git a/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json
b/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json
index 0942404b0eb..86bf1193abd 100644
--- a/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json
+++ b/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json
@@ -1,5 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to
run",
- "revision": 3,
- "https://github.com/apache/beam/pull/39990": "removing dead code from
FnApiDoFnRunner"
+ "revision": 7
}
diff --git a/.github/workflows/beam_PostCommit_Yaml_Xlang_Direct.yml
b/.github/workflows/beam_PostCommit_Yaml_Xlang_Direct.yml
index a2de7f285ed..b390f8faf4e 100644
--- a/.github/workflows/beam_PostCommit_Yaml_Xlang_Direct.yml
+++ b/.github/workflows/beam_PostCommit_Yaml_Xlang_Direct.yml
@@ -82,7 +82,7 @@ jobs:
- name: run PostCommit Yaml Xlang Direct script
uses: ./.github/actions/gradle-command-self-hosted-action
with:
- gradle-command: :sdks:python:postCommitYamlIntegrationTests
-PyamlTestSet=${{ matrix.test_set }} -PbeamPythonExtra=p310_ml_test,yaml
+ gradle-command: :sdks:python:postCommitYamlIntegrationTests
-PyamlTestSet=${{ matrix.test_set }}
-PbeamPythonExtra=p310_ml_test,yaml,gcp_test
- name: Archive Python Test Results
uses: actions/upload-artifact@v7
if: failure()
diff --git a/.github/workflows/beam_PreCommit_Yaml_Xlang_Direct.yml
b/.github/workflows/beam_PreCommit_Yaml_Xlang_Direct.yml
index cf36408537b..d92783c5318 100644
--- a/.github/workflows/beam_PreCommit_Yaml_Xlang_Direct.yml
+++ b/.github/workflows/beam_PreCommit_Yaml_Xlang_Direct.yml
@@ -95,7 +95,7 @@ jobs:
- name: run PreCommit Yaml Xlang Direct script
uses: ./.github/actions/gradle-command-self-hosted-action
with:
- gradle-command: :sdks:python:yamlIntegrationTests
-PbeamPythonExtra=p310_ml_test,yaml
+ gradle-command: :sdks:python:yamlIntegrationTests
-PbeamPythonExtra=p310_ml_test,yaml,gcp_test
- name: Archive Python Test Results
uses: actions/upload-artifact@v7
if: failure()
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreReadSchemaTransformConfiguration.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreReadSchemaTransformConfiguration.java
new file mode 100644
index 00000000000..0a1e0e20e59
--- /dev/null
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreReadSchemaTransformConfiguration.java
@@ -0,0 +1,83 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.beam.sdk.io.gcp.firestore;
+
+import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Configuration class for the Firestore Read transform. */
+@DefaultSchema(AutoValueSchema.class)
+@AutoValue
+public abstract class FirestoreReadSchemaTransformConfiguration implements
Serializable {
+
+ @SchemaFieldDescription("GCP project id. Defaults to GcpOptions project when
unset.")
+ @Nullable
+ public abstract String getProjectId();
+
+ @SchemaFieldDescription(
+ "Firestore database id. Defaults to FirestoreOptions firestoreDb when
unset.")
+ @Nullable
+ public abstract String getDatabaseId();
+
+ @SchemaFieldDescription("Firestore collection id to read from.")
+ public abstract String getCollectionId();
+
+ @SchemaFieldDescription(
+ "The schema in which the data is encoded, defined with JSON-schema
syntax "
+ + "(https://json-schema.org/).")
+ public abstract String getSchema();
+
+ @SchemaFieldDescription(
+ "This option specifies whether and where to output rows that failed to
be read.")
+ @Nullable
+ public abstract ErrorHandling getErrorHandling();
+
+ public void validate() {
+ checkArgument(
+ getCollectionId() != null && !getCollectionId().isEmpty(),
+ "Firestore collection id must be specified.");
+ checkArgument(
+ getSchema() != null && !getSchema().isEmpty(), "Firestore schema must
be specified.");
+ }
+
+ public static Builder builder() {
+ return new AutoValue_FirestoreReadSchemaTransformConfiguration.Builder();
+ }
+
+ @AutoValue.Builder
+ public abstract static class Builder {
+ public abstract Builder setProjectId(String projectId);
+
+ public abstract Builder setDatabaseId(String databaseId);
+
+ public abstract Builder setCollectionId(String collectionId);
+
+ public abstract Builder setSchema(String schema);
+
+ public abstract Builder setErrorHandling(ErrorHandling errorHandling);
+
+ public abstract FirestoreReadSchemaTransformConfiguration build();
+ }
+}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreReadSchemaTransformProvider.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreReadSchemaTransformProvider.java
new file mode 100644
index 00000000000..c035a53acda
--- /dev/null
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreReadSchemaTransformProvider.java
@@ -0,0 +1,216 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.beam.sdk.io.gcp.firestore;
+
+import com.google.auto.service.AutoService;
+import com.google.firestore.v1.Document;
+import com.google.firestore.v1.ListDocumentsRequest;
+import java.nio.charset.StandardCharsets;
+import java.util.Collections;
+import java.util.List;
+import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
+import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
+import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
+import org.apache.beam.sdk.schemas.utils.JsonUtils;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionRowTuple;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings;
+
+/** A {@link SchemaTransformProvider} for reading from Google Cloud Firestore.
*/
+@AutoService(SchemaTransformProvider.class)
+@SuppressWarnings({
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+public class FirestoreReadSchemaTransformProvider
+ extends
TypedSchemaTransformProvider<FirestoreReadSchemaTransformConfiguration> {
+
+ private static final String OUTPUT_TAG_NAME = "output";
+ public static final TupleTag<Row> OUTPUT_TAG = new TupleTag<Row>() {};
+ public static final TupleTag<Row> ERROR_TAG = new TupleTag<Row>() {};
+
+ private static final org.apache.beam.sdk.metrics.Counter errorCounter =
+ org.apache.beam.sdk.metrics.Metrics.counter(
+ FirestoreReadSchemaTransformProvider.class,
"Firestore-read-error-counter");
+
+ @Override
+ protected SchemaTransform from(FirestoreReadSchemaTransformConfiguration
configuration) {
+ return new FirestoreReadSchemaTransform(configuration);
+ }
+
+ @Override
+ public String identifier() {
+ return "beam:schematransform:org.apache.beam:firestore_read:v1";
+ }
+
+ @Override
+ public String description() {
+ return "Reads documents from a Google Cloud Firestore collection and
outputs Beam Rows.";
+ }
+
+ @Override
+ public List<String> inputCollectionNames() {
+ return Collections.emptyList();
+ }
+
+ @Override
+ public List<String> outputCollectionNames() {
+ return Collections.singletonList(OUTPUT_TAG_NAME);
+ }
+
+ private static class FirestoreReadSchemaTransform extends SchemaTransform {
+ private final FirestoreReadSchemaTransformConfiguration configuration;
+
+ FirestoreReadSchemaTransform(FirestoreReadSchemaTransformConfiguration
configuration) {
+ configuration.validate();
+ this.configuration = configuration;
+ }
+
+ @Override
+ public PCollectionRowTuple expand(PCollectionRowTuple input) {
+ if (!input.getAll().isEmpty()) {
+ throw new IllegalStateException(
+ "Firestore read transform does not expect input PCollections.");
+ }
+
+ Schema schema =
JsonUtils.beamSchemaFromJsonSchema(configuration.getSchema());
+ String projectId = resolveProjectId(input.getPipeline());
+ String databaseId = resolveDatabaseId(input.getPipeline());
+ String parent = FirestoreUtils.documentsRoot(projectId, databaseId);
+
+ PCollection<ListDocumentsRequest> requests =
+ input
+ .getPipeline()
+ .apply("CreateCollectionId",
Create.of(configuration.getCollectionId()))
+ .apply(
+ "BuildListDocumentsRequest",
+ ParDo.of(
+ new DoFn<String, ListDocumentsRequest>() {
+ @ProcessElement
+ public void processElement(
+ @Element String collectionId,
+ OutputReceiver<ListDocumentsRequest> out) {
+ out.output(
+ ListDocumentsRequest.newBuilder()
+ .setParent(parent)
+ .setCollectionId(collectionId)
+ .build());
+ }
+ }));
+
+ FirestoreV1.ListDocuments.Builder readBuilder =
+ FirestoreIO.v1()
+ .read()
+ .listDocuments()
+ .withProjectId(projectId)
+ .withDatabaseId(databaseId);
+
+ PCollection<Document> documents = requests.apply("ReadFromFirestore",
readBuilder.build());
+
+ boolean handleErrors =
ErrorHandling.hasOutput(configuration.getErrorHandling());
+ Schema errorSchema = ErrorHandling.errorSchemaBytes();
+ String documentIdField = schema.hasField("document_id") ? "document_id"
: null;
+
+ PCollectionTuple outputTuple =
+ documents.apply(
+ "ConvertToBeamRows",
+ ParDo.of(new DocumentToRowFn(schema, documentIdField,
handleErrors, errorSchema))
+ .withOutputTags(OUTPUT_TAG, TupleTagList.of(ERROR_TAG)));
+
+ PCollection<Row> rows = outputTuple.get(OUTPUT_TAG).setRowSchema(schema);
+ PCollectionRowTuple output = PCollectionRowTuple.of(OUTPUT_TAG_NAME,
rows);
+ if (handleErrors && configuration.getErrorHandling() != null) {
+ output =
+ output.and(
+ configuration.getErrorHandling().getOutput(),
+ outputTuple.get(ERROR_TAG).setRowSchema(errorSchema));
+ }
+ return output;
+ }
+
+ private String resolveProjectId(org.apache.beam.sdk.Pipeline pipeline) {
+ if (!Strings.isNullOrEmpty(configuration.getProjectId())) {
+ return configuration.getProjectId();
+ }
+ FirestoreOptions firestoreOptions =
pipeline.getOptions().as(FirestoreOptions.class);
+ if (!Strings.isNullOrEmpty(firestoreOptions.getFirestoreProject())) {
+ return firestoreOptions.getFirestoreProject();
+ }
+ String project = pipeline.getOptions().as(GcpOptions.class).getProject();
+ if (Strings.isNullOrEmpty(project)) {
+ throw new IllegalArgumentException(
+ "Firestore project id must be set on the transform or pipeline
options.");
+ }
+ return project;
+ }
+
+ private String resolveDatabaseId(org.apache.beam.sdk.Pipeline pipeline) {
+ if (!Strings.isNullOrEmpty(configuration.getDatabaseId())) {
+ return configuration.getDatabaseId();
+ }
+ return pipeline.getOptions().as(FirestoreOptions.class).getFirestoreDb();
+ }
+ }
+
+ static class DocumentToRowFn extends DoFn<Document, Row> {
+ private final Schema schema;
+ private final @org.checkerframework.checker.nullness.qual.Nullable String
documentIdField;
+ private final boolean handleErrors;
+ private final Schema errorSchema;
+
+ DocumentToRowFn(
+ Schema schema,
+ @org.checkerframework.checker.nullness.qual.Nullable String
documentIdField,
+ boolean handleErrors,
+ Schema errorSchema) {
+ this.schema = schema;
+ this.documentIdField = documentIdField;
+ this.handleErrors = handleErrors;
+ this.errorSchema = errorSchema;
+ }
+
+ @ProcessElement
+ public void processElement(@Element Document document, MultiOutputReceiver
receiver) {
+ try {
+ receiver
+ .get(OUTPUT_TAG)
+ .output(FirestoreUtils.documentToRow(document, schema,
documentIdField));
+ } catch (Exception e) {
+ if (!handleErrors) {
+ throw new RuntimeException(
+ "Failed to convert Firestore document to Beam Row: " +
document.getName(), e);
+ }
+ errorCounter.inc();
+ receiver
+ .get(ERROR_TAG)
+ .output(
+ ErrorHandling.errorRecord(
+ errorSchema,
document.getName().getBytes(StandardCharsets.UTF_8), e));
+ }
+ }
+ }
+}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreUtils.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreUtils.java
new file mode 100644
index 00000000000..0f754846013
--- /dev/null
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreUtils.java
@@ -0,0 +1,312 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.beam.sdk.io.gcp.firestore;
+
+import com.google.firestore.v1.ArrayValue;
+import com.google.firestore.v1.Document;
+import com.google.firestore.v1.MapValue;
+import com.google.firestore.v1.Value;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.util.Timestamps;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.values.Row;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
+
+/** Utility methods for Firestore SchemaTransform providers. */
+@SuppressWarnings({
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+final class FirestoreUtils {
+
+ private FirestoreUtils() {}
+
+ static String documentsRoot(String projectId, String databaseId) {
+ return String.format("projects/%s/databases/%s/documents", projectId,
databaseId);
+ }
+
+ static String documentPath(
+ String projectId, String databaseId, String collectionId, String
documentId) {
+ return String.format(
+ "%s/%s/%s", documentsRoot(projectId, databaseId), collectionId,
documentId);
+ }
+
+ static String documentIdFromName(String documentName) {
+ int lastSlash = documentName.lastIndexOf('/');
+ if (lastSlash < 0 || lastSlash == documentName.length() - 1) {
+ throw new IllegalArgumentException("Invalid Firestore document name: " +
documentName);
+ }
+ return documentName.substring(lastSlash + 1);
+ }
+
+ static Row documentToRow(Document document, Schema schema, @Nullable String
documentIdField) {
+ Map<String, Object> values = new HashMap<>();
+ for (Map.Entry<String, Value> entry : document.getFieldsMap().entrySet()) {
+ values.put(entry.getKey(), valueToJava(entry.getValue()));
+ }
+ if (documentIdField != null && schema.hasField(documentIdField)) {
+ values.put(documentIdField, documentIdFromName(document.getName()));
+ }
+ return toRow(values, schema);
+ }
+
+ static Document rowToDocument(
+ Row row,
+ Schema schema,
+ String projectId,
+ String databaseId,
+ String collectionId,
+ String documentIdField) {
+ String documentId = row.getString(documentIdField);
+ if (documentId == null || documentId.isEmpty()) {
+ throw new IllegalArgumentException(
+ "Document id field '" + documentIdField + "' must be set on input
rows.");
+ }
+
+ Document.Builder builder =
+ Document.newBuilder()
+ .setName(documentPath(projectId, databaseId, collectionId,
documentId));
+ for (Field field : schema.getFields()) {
+ String fieldName = field.getName();
+ if (fieldName.equals(documentIdField)) {
+ continue;
+ }
+ Object fieldValue = row.getValue(fieldName);
+ if (fieldValue != null) {
+ builder.putFields(fieldName, javaToValue(fieldValue, field.getType()));
+ }
+ }
+ return builder.build();
+ }
+
+ static Row toRow(Map<String, ?> values, Schema schema) {
+ Row.Builder rowBuilder = Row.withSchema(schema);
+ for (Field field : schema.getFields()) {
+ rowBuilder.addValue(convertFromJava(values.get(field.getName()),
field.getType()));
+ }
+ return rowBuilder.build();
+ }
+
+ private static Map<String, ?> castToStringKeyMap(Map<?, ?> map) {
+ Map<String, Object> converted = new HashMap<>();
+ for (Map.Entry<?, ?> entry : map.entrySet()) {
+ converted.put(String.valueOf(entry.getKey()), entry.getValue());
+ }
+ return converted;
+ }
+
+ private static @Nullable Object valueToJava(Value value) {
+ switch (value.getValueTypeCase()) {
+ case STRING_VALUE:
+ return value.getStringValue();
+ case INTEGER_VALUE:
+ return value.getIntegerValue();
+ case DOUBLE_VALUE:
+ return value.getDoubleValue();
+ case BOOLEAN_VALUE:
+ return value.getBooleanValue();
+ case TIMESTAMP_VALUE:
+ return new Instant(Timestamps.toMillis(value.getTimestampValue()));
+ case BYTES_VALUE:
+ return value.getBytesValue().toByteArray();
+ case NULL_VALUE:
+ return null;
+ case ARRAY_VALUE:
+ List<@Nullable Object> values = new ArrayList<>();
+ for (Value element : value.getArrayValue().getValuesList()) {
+ values.add(valueToJava(element));
+ }
+ return values;
+ case MAP_VALUE:
+ Map<String, Object> map = new HashMap<>();
+ for (Map.Entry<String, Value> entry :
value.getMapValue().getFieldsMap().entrySet()) {
+ map.put(entry.getKey(), valueToJava(entry.getValue()));
+ }
+ return map;
+ case VALUETYPE_NOT_SET:
+ return null;
+ default:
+ throw new IllegalArgumentException(
+ "Unsupported Firestore value type: " + value.getValueTypeCase());
+ }
+ }
+
+ private static Value javaToValue(Object value, FieldType fieldType) {
+ if (value == null) {
+ return
Value.newBuilder().setNullValue(com.google.protobuf.NullValue.NULL_VALUE).build();
+ }
+ switch (fieldType.getTypeName()) {
+ case STRING:
+ return Value.newBuilder().setStringValue(value.toString()).build();
+ case INT64:
+ return Value.newBuilder().setIntegerValue(((Number)
value).longValue()).build();
+ case DOUBLE:
+ return Value.newBuilder().setDoubleValue(((Number)
value).doubleValue()).build();
+ case BOOLEAN:
+ return Value.newBuilder().setBooleanValue((Boolean) value).build();
+ case DATETIME:
+ Instant instant = (Instant) value;
+ return Value.newBuilder()
+ .setTimestampValue(Timestamps.fromMillis(instant.getMillis()))
+ .build();
+ case BYTES:
+ return Value.newBuilder().setBytesValue(ByteString.copyFrom((byte[])
value)).build();
+ case ARRAY:
+ case ITERABLE:
+ ArrayValue.Builder arrayBuilder = ArrayValue.newBuilder();
+ FieldType elementType = fieldType.getCollectionElementType();
+ if (elementType == null) {
+ throw new IllegalArgumentException("Collection element type cannot
be null.");
+ }
+ for (Object item : (Iterable<?>) value) {
+ arrayBuilder.addValues(
+ item == null
+ ? Value.newBuilder()
+ .setNullValue(com.google.protobuf.NullValue.NULL_VALUE)
+ .build()
+ : javaToValue(item, elementType));
+ }
+ return Value.newBuilder().setArrayValue(arrayBuilder.build()).build();
+ case MAP:
+ MapValue.Builder mapBuilder = MapValue.newBuilder();
+ FieldType valueType = fieldType.getMapValueType();
+ if (valueType == null) {
+ throw new IllegalArgumentException("Map value type cannot be null.");
+ }
+ for (Map.Entry<?, ?> entry : ((Map<?, ?>) value).entrySet()) {
+ Object mapValue = entry.getValue();
+ mapBuilder.putFields(
+ String.valueOf(entry.getKey()),
+ mapValue == null
+ ? Value.newBuilder()
+ .setNullValue(com.google.protobuf.NullValue.NULL_VALUE)
+ .build()
+ : javaToValue(mapValue, valueType));
+ }
+ return Value.newBuilder().setMapValue(mapBuilder.build()).build();
+ case ROW:
+ Schema rowSchema = fieldType.getRowSchema();
+ if (rowSchema == null) {
+ throw new IllegalArgumentException("Row schema cannot be null.");
+ }
+ if (!(value instanceof Row)) {
+ throw new IllegalArgumentException("Expected Row for nested field.");
+ }
+ MapValue.Builder nestedMapBuilder = MapValue.newBuilder();
+ Row nestedRow = (Row) value;
+ for (Field nestedField : rowSchema.getFields()) {
+ Object nestedValue = nestedRow.getValue(nestedField.getName());
+ if (nestedValue != null) {
+ nestedMapBuilder.putFields(
+ nestedField.getName(), javaToValue(nestedValue,
nestedField.getType()));
+ }
+ }
+ return
Value.newBuilder().setMapValue(nestedMapBuilder.build()).build();
+ default:
+ throw new IllegalArgumentException("Unsupported field type: " +
fieldType);
+ }
+ }
+
+ private static @Nullable Object convertFromJava(@Nullable Object value,
FieldType fieldType) {
+ if (value == null) {
+ return null;
+ }
+ switch (fieldType.getTypeName()) {
+ case BYTE:
+ return ((Number) value).byteValue();
+ case INT16:
+ return ((Number) value).shortValue();
+ case INT32:
+ return ((Number) value).intValue();
+ case INT64:
+ return ((Number) value).longValue();
+ case FLOAT:
+ return ((Number) value).floatValue();
+ case DOUBLE:
+ return ((Number) value).doubleValue();
+ case DECIMAL:
+ return value instanceof java.math.BigDecimal
+ ? value
+ : java.math.BigDecimal.valueOf(((Number) value).doubleValue());
+ case STRING:
+ return value.toString();
+ case BOOLEAN:
+ return value;
+ case DATETIME:
+ if (value instanceof Instant) {
+ return value;
+ }
+ if (value instanceof Number) {
+ return new Instant(((Number) value).longValue());
+ }
+ return Instant.parse(value.toString());
+ case BYTES:
+ if (value instanceof byte[]) {
+ return value;
+ }
+ return
value.toString().getBytes(java.nio.charset.StandardCharsets.UTF_8);
+ case ARRAY:
+ case ITERABLE:
+ if (!(value instanceof Iterable)) {
+ throw new IllegalArgumentException("Expected Iterable for array
field.");
+ }
+ FieldType elementType = fieldType.getCollectionElementType();
+ if (elementType == null) {
+ throw new IllegalArgumentException("Collection element type cannot
be null.");
+ }
+ List<@Nullable Object> rowList = new ArrayList<>();
+ for (Object item : (Iterable<?>) value) {
+ rowList.add(convertFromJava(item, elementType));
+ }
+ return rowList;
+ case MAP:
+ if (!(value instanceof Map)) {
+ throw new IllegalArgumentException("Expected Map for map field.");
+ }
+ FieldType valueType = fieldType.getMapValueType();
+ if (valueType == null) {
+ throw new IllegalArgumentException("Map value type cannot be null.");
+ }
+ Map<String, @Nullable Object> rowMap = new HashMap<>();
+ for (Map.Entry<?, ?> entry : ((Map<?, ?>) value).entrySet()) {
+ rowMap.put(String.valueOf(entry.getKey()),
convertFromJava(entry.getValue(), valueType));
+ }
+ return rowMap;
+ case ROW:
+ Schema rowSchema = fieldType.getRowSchema();
+ if (rowSchema == null) {
+ throw new IllegalArgumentException("Row schema cannot be null.");
+ }
+ if (value instanceof Map) {
+ return toRow(castToStringKeyMap((Map<?, ?>) value), rowSchema);
+ }
+ if (value instanceof Row) {
+ return value;
+ }
+ throw new IllegalArgumentException("Cannot convert value to Row.");
+ default:
+ throw new IllegalArgumentException("Unsupported field type: " +
fieldType);
+ }
+ }
+}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1.java
index 66e4eca0d87..e8da6097788 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1.java
@@ -743,7 +743,12 @@ public final class FirestoreV1 {
"listDocuments",
ParDo.of(
new ListDocumentsFn(
- clock, firestoreStatefulComponentFactory, rpcQosOptions,
readTime)))
+ clock,
+ firestoreStatefulComponentFactory,
+ rpcQosOptions,
+ readTime,
+ projectId,
+ databaseId)))
.apply(ParDo.of(new ListDocumentsResponseToDocument()))
.apply(Reshuffle.viaRandomKey());
}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1ReadFn.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1ReadFn.java
index 84e1cb1be0a..7c8f2bf345f 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1ReadFn.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1ReadFn.java
@@ -675,11 +675,12 @@ final class FirestoreV1ReadFn {
protected final @Nullable Instant readTime;
+ private final @Nullable String configuredProjectId;
+ private final @Nullable String configuredDatabaseId;
+
// transient running state information, not important to any possible
checkpointing
protected transient FirestoreStub firestoreStub;
protected transient RpcQos rpcQos;
- protected transient String projectId;
- protected transient @Nullable String databaseId;
@SuppressWarnings(
"initialization.fields.uninitialized") // allow transient fields to be
managed by component
@@ -696,10 +697,8 @@ final class FirestoreV1ReadFn {
requireNonNull(firestoreStatefulComponentFactory, "firestoreFactory
must be non null");
this.rpcQosOptions = requireNonNull(rpcQosOptions, "rpcQosOptions must
be non null");
this.readTime = readTime;
- if (projectId != null) {
- this.projectId = projectId;
- }
- this.databaseId = databaseId;
+ this.configuredProjectId = projectId;
+ this.configuredDatabaseId = databaseId;
}
/** {@inheritDoc} */
@@ -712,33 +711,30 @@ final class FirestoreV1ReadFn {
@Override
public final void startBundle(StartBundleContext c) {
String project =
- this.projectId != null
- ? this.projectId
+ configuredProjectId != null
+ ? configuredProjectId
:
c.getPipelineOptions().as(FirestoreOptions.class).getFirestoreProject();
if (project == null) {
project = c.getPipelineOptions().as(GcpOptions.class).getProject();
}
- projectId =
- requireNonNull(
- project,
- "project must be defined on FirestoreOptions or GcpOptions of
PipelineOptions");
- databaseId =
- this.databaseId != null
- ? this.databaseId
+ String databaseId =
+ configuredDatabaseId != null
+ ? configuredDatabaseId
:
c.getPipelineOptions().as(FirestoreOptions.class).getFirestoreDb();
- requireNonNull(
- databaseId, "firestoreDb must be defined on FirestoreOptions of
PipelineOptions");
firestoreStub =
firestoreStatefulComponentFactory.getFirestoreStub(
- c.getPipelineOptions(), projectId, databaseId);
+ c.getPipelineOptions(),
+ requireNonNull(
+ project,
+ "project must be defined on FirestoreOptions or GcpOptions
of PipelineOptions"),
+ requireNonNull(
+ databaseId,
+ "firestoreDb must be defined on FirestoreOptions of
PipelineOptions"));
}
/** {@inheritDoc} */
- @SuppressWarnings("nullness") // allow clearing transient fields
@Override
public void finishBundle() throws Exception {
- projectId = null;
- databaseId = null;
firestoreStub.close();
}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreWriteSchemaTransformConfiguration.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreWriteSchemaTransformConfiguration.java
new file mode 100644
index 00000000000..1f56a7b7444
--- /dev/null
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreWriteSchemaTransformConfiguration.java
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.beam.sdk.io.gcp.firestore;
+
+import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Configuration class for the Firestore Write transform. */
+@DefaultSchema(AutoValueSchema.class)
+@AutoValue
+public abstract class FirestoreWriteSchemaTransformConfiguration implements
Serializable {
+
+ @SchemaFieldDescription("GCP project id. Defaults to GcpOptions project when
unset.")
+ @Nullable
+ public abstract String getProjectId();
+
+ @SchemaFieldDescription(
+ "Firestore database id. Defaults to FirestoreOptions firestoreDb when
unset.")
+ @Nullable
+ public abstract String getDatabaseId();
+
+ @SchemaFieldDescription("Firestore collection id to write to.")
+ public abstract String getCollectionId();
+
+ @SchemaFieldDescription(
+ "Row field containing the document id. Defaults to document_id when
unset.")
+ @Nullable
+ public abstract String getDocumentIdField();
+
+ @SchemaFieldDescription(
+ "This option specifies whether and where to output unwritable rows.
Error handling is "
+ + "limited to data conversion failures before sending writes to
Firestore.")
+ @Nullable
+ public abstract ErrorHandling getErrorHandling();
+
+ public void validate() {
+ checkArgument(
+ getCollectionId() != null && !getCollectionId().isEmpty(),
+ "Firestore collection id must be specified.");
+ }
+
+ public static Builder builder() {
+ return new AutoValue_FirestoreWriteSchemaTransformConfiguration.Builder();
+ }
+
+ @AutoValue.Builder
+ public abstract static class Builder {
+ public abstract Builder setProjectId(String projectId);
+
+ public abstract Builder setDatabaseId(String databaseId);
+
+ public abstract Builder setCollectionId(String collectionId);
+
+ public abstract Builder setDocumentIdField(String documentIdField);
+
+ public abstract Builder setErrorHandling(ErrorHandling errorHandling);
+
+ public abstract FirestoreWriteSchemaTransformConfiguration build();
+ }
+}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreWriteSchemaTransformProvider.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreWriteSchemaTransformProvider.java
new file mode 100644
index 00000000000..4ff72477356
--- /dev/null
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreWriteSchemaTransformProvider.java
@@ -0,0 +1,201 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.beam.sdk.io.gcp.firestore;
+
+import com.google.auto.service.AutoService;
+import com.google.firestore.v1.Document;
+import com.google.firestore.v1.Write;
+import java.util.Collections;
+import java.util.List;
+import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
+import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
+import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionRowTuple;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings;
+
+/** A {@link SchemaTransformProvider} for writing to Google Cloud Firestore. */
+@AutoService(SchemaTransformProvider.class)
+@SuppressWarnings({
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+public class FirestoreWriteSchemaTransformProvider
+ extends
TypedSchemaTransformProvider<FirestoreWriteSchemaTransformConfiguration> {
+
+ private static final String INPUT_TAG = "input";
+ private static final String DEFAULT_DOCUMENT_ID_FIELD = "document_id";
+ public static final TupleTag<Write> OUTPUT_TAG = new TupleTag<Write>() {};
+ public static final TupleTag<Row> ERROR_TAG = new TupleTag<Row>() {};
+
+ private static final org.apache.beam.sdk.metrics.Counter errorCounter =
+ org.apache.beam.sdk.metrics.Metrics.counter(
+ FirestoreWriteSchemaTransformProvider.class,
"Firestore-write-error-counter");
+
+ @Override
+ protected SchemaTransform from(FirestoreWriteSchemaTransformConfiguration
configuration) {
+ return new FirestoreWriteSchemaTransform(configuration);
+ }
+
+ @Override
+ public String identifier() {
+ return "beam:schematransform:org.apache.beam:firestore_write:v1";
+ }
+
+ @Override
+ public String description() {
+ return "Writes Beam Rows to a Google Cloud Firestore collection.";
+ }
+
+ @Override
+ public List<String> inputCollectionNames() {
+ return Collections.singletonList(INPUT_TAG);
+ }
+
+ @Override
+ public List<String> outputCollectionNames() {
+ return Collections.emptyList();
+ }
+
+ private static class FirestoreWriteSchemaTransform extends SchemaTransform {
+ private final FirestoreWriteSchemaTransformConfiguration configuration;
+
+ FirestoreWriteSchemaTransform(FirestoreWriteSchemaTransformConfiguration
configuration) {
+ configuration.validate();
+ this.configuration = configuration;
+ }
+
+ @Override
+ public PCollectionRowTuple expand(PCollectionRowTuple input) {
+ PCollection<Row> rows = input.get(INPUT_TAG);
+ Schema inputSchema = rows.getSchema();
+ String projectId = resolveProjectId(input.getPipeline());
+ String databaseId = resolveDatabaseId(input.getPipeline());
+ String documentIdField =
+ Strings.isNullOrEmpty(configuration.getDocumentIdField())
+ ? DEFAULT_DOCUMENT_ID_FIELD
+ : configuration.getDocumentIdField();
+ if (!inputSchema.hasField(documentIdField)) {
+ throw new IllegalArgumentException(
+ "Input schema must contain document id field: " + documentIdField);
+ }
+
+ boolean handleErrors =
ErrorHandling.hasOutput(configuration.getErrorHandling());
+ Schema errorSchema = ErrorHandling.errorSchema(inputSchema);
+
+ PCollectionTuple outputTuple =
+ rows.apply(
+ "ConvertToFirestoreWrite",
+ ParDo.of(
+ new RowToWriteFn(
+ inputSchema,
+ projectId,
+ databaseId,
+ configuration.getCollectionId(),
+ documentIdField,
+ handleErrors,
+ errorSchema))
+ .withOutputTags(OUTPUT_TAG, TupleTagList.of(ERROR_TAG)));
+
+ FirestoreV1.Write write =
+
FirestoreIO.v1().write().withProjectId(projectId).withDatabaseId(databaseId);
+
+ outputTuple.get(OUTPUT_TAG).apply("WriteToFirestore",
write.batchWrite().build());
+
+ PCollection<Row> errorOutput =
outputTuple.get(ERROR_TAG).setRowSchema(errorSchema);
+ ErrorHandling errorHandling = configuration.getErrorHandling();
+ return PCollectionRowTuple.of(
+ (handleErrors && errorHandling != null) ? errorHandling.getOutput()
: "errors",
+ errorOutput);
+ }
+
+ private String resolveProjectId(org.apache.beam.sdk.Pipeline pipeline) {
+ if (!Strings.isNullOrEmpty(configuration.getProjectId())) {
+ return configuration.getProjectId();
+ }
+ FirestoreOptions firestoreOptions =
pipeline.getOptions().as(FirestoreOptions.class);
+ if (!Strings.isNullOrEmpty(firestoreOptions.getFirestoreProject())) {
+ return firestoreOptions.getFirestoreProject();
+ }
+ String project = pipeline.getOptions().as(GcpOptions.class).getProject();
+ if (Strings.isNullOrEmpty(project)) {
+ throw new IllegalArgumentException(
+ "Firestore project id must be set on the transform or pipeline
options.");
+ }
+ return project;
+ }
+
+ private String resolveDatabaseId(org.apache.beam.sdk.Pipeline pipeline) {
+ if (!Strings.isNullOrEmpty(configuration.getDatabaseId())) {
+ return configuration.getDatabaseId();
+ }
+ return pipeline.getOptions().as(FirestoreOptions.class).getFirestoreDb();
+ }
+ }
+
+ static class RowToWriteFn extends DoFn<Row, Write> {
+ private final Schema schema;
+ private final String projectId;
+ private final String databaseId;
+ private final String collectionId;
+ private final String documentIdField;
+ private final boolean handleErrors;
+ private final Schema errorSchema;
+
+ RowToWriteFn(
+ Schema schema,
+ String projectId,
+ String databaseId,
+ String collectionId,
+ String documentIdField,
+ boolean handleErrors,
+ Schema errorSchema) {
+ this.schema = schema;
+ this.projectId = projectId;
+ this.databaseId = databaseId;
+ this.collectionId = collectionId;
+ this.documentIdField = documentIdField;
+ this.handleErrors = handleErrors;
+ this.errorSchema = errorSchema;
+ }
+
+ @ProcessElement
+ public void processElement(@Element Row row, MultiOutputReceiver receiver)
{
+ try {
+ Document document =
+ FirestoreUtils.rowToDocument(
+ row, schema, projectId, databaseId, collectionId,
documentIdField);
+
receiver.get(OUTPUT_TAG).output(Write.newBuilder().setUpdate(document).build());
+ } catch (Exception e) {
+ if (!handleErrors) {
+ throw new RuntimeException(e);
+ }
+ errorCounter.inc();
+ receiver.get(ERROR_TAG).output(ErrorHandling.errorRecord(errorSchema,
row, e));
+ }
+ }
+ }
+}
diff --git
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreSchemaTransformProviderTest.java
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreSchemaTransformProviderTest.java
new file mode 100644
index 00000000000..7d58ce8d551
--- /dev/null
+++
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreSchemaTransformProviderTest.java
@@ -0,0 +1,179 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.beam.sdk.io.gcp.firestore;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.util.List;
+import java.util.ServiceLoader;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
+import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionRowTuple;
+import org.apache.beam.sdk.values.Row;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for Firestore SchemaTransform providers. */
+@RunWith(JUnit4.class)
+public class FirestoreSchemaTransformProviderTest {
+
+ @Rule
+ public final transient TestPipeline pipeline =
+ TestPipeline.fromOptions(PipelineOptionsFactory.create())
+ .enableAbandonedNodeEnforcement(false);
+
+ @Test
+ public void testReadFindTransform() {
+ SchemaTransformProvider provider = loadReadProvider();
+
+ assertEquals(Lists.newArrayList("output"),
provider.outputCollectionNames());
+ assertEquals(Lists.newArrayList(), provider.inputCollectionNames());
+ assertEquals("beam:schematransform:org.apache.beam:firestore_read:v1",
provider.identifier());
+ assertNotNull(provider.description());
+
+ assertEquals(
+ Sets.newHashSet("project_id", "database_id", "collection_id",
"schema", "error_handling"),
+ provider.configurationSchema().getFields().stream()
+ .map(Schema.Field::getName)
+ .collect(Collectors.toSet()));
+ }
+
+ @Test
+ public void testReadBuildTransform() {
+ FirestoreReadSchemaTransformConfiguration readConfig =
+ FirestoreReadSchemaTransformConfiguration.builder()
+ .setProjectId("test-project")
+ .setDatabaseId("(default)")
+ .setCollectionId("users")
+ .setSchema(
+ "{"
+ + "\"type\":\"object\","
+ + "\"properties\":{"
+ + "\"document_id\":{\"type\":\"string\"},"
+ + "\"name\":{\"type\":\"string\"}"
+ + "},"
+ + "\"required\":[\"document_id\",\"name\"]"
+ + "}")
+ .build();
+
+ SchemaTransform transform = new
FirestoreReadSchemaTransformProvider().from(readConfig);
+ PCollectionRowTuple output =
transform.expand(PCollectionRowTuple.empty(pipeline));
+
+ assertEquals(1, output.getAll().size());
+ assertTrue(output.has("output"));
+ assertEquals(
+
Schema.builder().addStringField("document_id").addStringField("name").build(),
+ output.get("output").getSchema());
+ }
+
+ @Test
+ public void testReadWithNonEmptyInputThrows() {
+ FirestoreReadSchemaTransformConfiguration readConfig =
+ FirestoreReadSchemaTransformConfiguration.builder()
+ .setProjectId("test-project")
+ .setCollectionId("users")
+
.setSchema("{\"type\":\"object\",\"properties\":{\"name\":{\"type\":\"string\"}}}")
+ .build();
+ SchemaTransform transform = new
FirestoreReadSchemaTransformProvider().from(readConfig);
+
+ PCollection<Row> dummyInput =
+ pipeline.apply(
+ "CreateDummy",
Create.empty(Schema.builder().addStringField("dummy").build()));
+ assertThrows(
+ IllegalStateException.class,
+ () -> transform.expand(PCollectionRowTuple.of("input", dummyInput)));
+ }
+
+ @Test
+ public void testWriteFindTransform() {
+ SchemaTransformProvider provider = loadWriteProvider();
+
+ assertEquals(Lists.newArrayList(), provider.outputCollectionNames());
+ assertEquals(Lists.newArrayList("input"), provider.inputCollectionNames());
+ assertEquals("beam:schematransform:org.apache.beam:firestore_write:v1",
provider.identifier());
+ assertNotNull(provider.description());
+
+ assertEquals(
+ Sets.newHashSet(
+ "project_id", "database_id", "collection_id", "document_id_field",
"error_handling"),
+ provider.configurationSchema().getFields().stream()
+ .map(Schema.Field::getName)
+ .collect(Collectors.toSet()));
+ }
+
+ @Test
+ public void testWriteBuildTransform() {
+ FirestoreWriteSchemaTransformConfiguration writeConfig =
+ FirestoreWriteSchemaTransformConfiguration.builder()
+ .setProjectId("test-project")
+ .setDatabaseId("(default)")
+ .setCollectionId("users")
+ .build();
+ SchemaTransform transform = new
FirestoreWriteSchemaTransformProvider().from(writeConfig);
+ Schema schema =
Schema.builder().addStringField("document_id").addStringField("name").build();
+ PCollection<Row> inputRows = pipeline.apply("CreateRows",
Create.empty(schema));
+ PCollectionRowTuple output =
transform.expand(PCollectionRowTuple.of("input", inputRows));
+ assertEquals(1, output.getAll().size());
+ assertTrue(output.has("errors"));
+ }
+
+ @Test
+ public void testWriteMissingDocumentIdFieldThrows() {
+ FirestoreWriteSchemaTransformConfiguration writeConfig =
+ FirestoreWriteSchemaTransformConfiguration.builder()
+ .setProjectId("test-project")
+ .setCollectionId("users")
+ .build();
+ SchemaTransform transform = new
FirestoreWriteSchemaTransformProvider().from(writeConfig);
+ Schema schema = Schema.builder().addStringField("name").build();
+ PCollection<Row> inputRows = pipeline.apply("CreateInvalidRows",
Create.empty(schema));
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> transform.expand(PCollectionRowTuple.of("input", inputRows)));
+ }
+
+ private static SchemaTransformProvider loadReadProvider() {
+ List<SchemaTransformProvider> providers =
+
StreamSupport.stream(ServiceLoader.load(SchemaTransformProvider.class).spliterator(),
false)
+ .filter(provider -> provider.getClass() ==
FirestoreReadSchemaTransformProvider.class)
+ .collect(Collectors.toList());
+ return providers.get(0);
+ }
+
+ private static SchemaTransformProvider loadWriteProvider() {
+ List<SchemaTransformProvider> providers =
+
StreamSupport.stream(ServiceLoader.load(SchemaTransformProvider.class).spliterator(),
false)
+ .filter(provider -> provider.getClass() ==
FirestoreWriteSchemaTransformProvider.class)
+ .collect(Collectors.toList());
+ return providers.get(0);
+ }
+}
diff --git
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreUtilsTest.java
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreUtilsTest.java
new file mode 100644
index 00000000000..b32075a00cc
--- /dev/null
+++
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreUtilsTest.java
@@ -0,0 +1,65 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.beam.sdk.io.gcp.firestore;
+
+import static org.junit.Assert.assertEquals;
+
+import com.google.firestore.v1.Document;
+import com.google.firestore.v1.Value;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.values.Row;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class FirestoreUtilsTest {
+
+ @Test
+ public void testDocumentRowRoundTrip() {
+ Schema schema =
Schema.builder().addStringField("document_id").addStringField("name").build();
+ Row input = Row.withSchema(schema).addValues("doc-1", "Alice").build();
+ Document document =
+ FirestoreUtils.rowToDocument(
+ input, schema, "test-project", "(default)", "users",
"document_id");
+ Row output = FirestoreUtils.documentToRow(document, schema, "document_id");
+
+ assertEquals("doc-1", output.getString("document_id"));
+ assertEquals("Alice", output.getString("name"));
+ assertEquals("Alice",
document.getFieldsMap().get("name").getStringValue());
+ }
+
+ @Test
+ public void testDocumentIdFromName() {
+ assertEquals(
+ "doc-1",
+
FirestoreUtils.documentIdFromName("projects/p/databases/(default)/documents/users/doc-1"));
+ }
+
+ @Test
+ public void testIntegerValueConversion() {
+ Schema schema = Schema.builder().addInt64Field("count").build();
+ Document document =
+ Document.newBuilder()
+ .setName("projects/p/databases/(default)/documents/users/doc-1")
+ .putFields("count",
Value.newBuilder().setIntegerValue(42L).build())
+ .build();
+ Row row = FirestoreUtils.documentToRow(document, schema, null);
+ assertEquals(42L, row.getInt64("count").longValue());
+ }
+}
diff --git
a/sdks/python/apache_beam/yaml/extended_tests/databases/firestore.yaml
b/sdks/python/apache_beam/yaml/extended_tests/databases/firestore.yaml
new file mode 100644
index 00000000000..19f2574d057
--- /dev/null
+++ b/sdks/python/apache_beam/yaml/extended_tests/databases/firestore.yaml
@@ -0,0 +1,84 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You 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.
+#
+
+fixtures:
+ - name: firestore_vars
+ type: "apache_beam.yaml.integration_tests.temp_firestore_collection"
+ config:
+ project: "apache-beam-testing"
+ database: "firestoredb"
+
+pipelines:
+ - pipeline:
+ type: composite
+ transforms:
+ - type: Create
+ name: CreateData
+ config:
+ elements:
+ - {document_id: "a", name: "Alice"}
+ - {document_id: "b", name: "Bob"}
+ - type: WriteToFirestore
+ name: WriteData
+ input: CreateData
+ config:
+ project: "apache-beam-testing"
+ database: "firestoredb"
+ collection: '{firestore_vars[COLLECTION]}'
+ error_handling:
+ output: write_errors
+ - type: AssertEqual
+ input: WriteData.write_errors
+ config:
+ elements: []
+ options:
+ project: "apache-beam-testing"
+ firestore_db: "firestoredb"
+
+ - pipeline:
+ type: composite
+ transforms:
+ - type: ReadFromFirestore
+ name: ReadData
+ config:
+ project: "apache-beam-testing"
+ database: "firestoredb"
+ collection: '{firestore_vars[COLLECTION]}'
+ schema: |
+ {
+ "type": "object",
+ "properties": {
+ "document_id": {"type": "string"},
+ "name": {"type": "string"}
+ },
+ "required": ["document_id", "name"]
+ }
+ error_handling:
+ output: read_errors
+ - type: AssertEqual
+ input: ReadData
+ config:
+ elements:
+ - {document_id: "a", name: "Alice"}
+ - {document_id: "b", name: "Bob"}
+ - type: AssertEqual
+ input: ReadData.read_errors
+ config:
+ elements: []
+ options:
+ project: "apache-beam-testing"
+ firestore_db: "firestoredb"
diff --git a/sdks/python/apache_beam/yaml/integration_tests.py
b/sdks/python/apache_beam/yaml/integration_tests.py
index 32794f4588d..71a2c3770c9 100644
--- a/sdks/python/apache_beam/yaml/integration_tests.py
+++ b/sdks/python/apache_beam/yaml/integration_tests.py
@@ -81,6 +81,10 @@ from google.cloud import pubsub_v1
from google.cloud.bigtable import client
from google.cloud.bigtable_admin_v2.types import instance
+try:
+ from google.cloud import firestore
+except ImportError:
+ firestore = None
try:
from google.cloud import secretmanager
except ImportError:
@@ -165,6 +169,56 @@ def temp_spanner_table(project, prefix='temp_spanner_db_'):
spanner_client._delete_database()
[email protected]
+def temp_firestore_collection(
+ project='apache-beam-testing',
+ database='firestoredb',
+ prefix='yaml_firestore_it_'):
+ """Context manager for an isolated Firestore collection used in YAML ITs.
+
+ Uses the shared Beam test project and the ``firestoredb`` database, matching
+ the Java Firestore integration tests.
+
+ Args:
+ project (str): GCP project id.
+ database (str): Firestore database id.
+ prefix (str): Prefix for the temporary collection name.
+
+ Yields:
+ dict: Keys ``PROJECT``, ``DATABASE``, and ``COLLECTION``.
+ """
+ if firestore is None:
+ raise RuntimeError("google-cloud-firestore is not installed.")
+
+ client = firestore.Client(project=project, database=database)
+ collection_id = f'{prefix}{uuid.uuid4().hex}'
+ logging.info(
+ 'Using Firestore collection %s in project %s database %s',
+ collection_id,
+ project,
+ database)
+ try:
+ yield {
+ 'PROJECT': project,
+ 'DATABASE': database,
+ 'COLLECTION': collection_id,
+ }
+ finally:
+ logging.info('Deleting documents in Firestore collection %s',
collection_id)
+ collection_ref = client.collection(collection_id)
+ batch = client.batch()
+ pending = 0
+ for doc in collection_ref.stream():
+ batch.delete(doc.reference)
+ pending += 1
+ if pending >= 400:
+ batch.commit()
+ batch = client.batch()
+ pending = 0
+ if pending:
+ batch.commit()
+
+
@contextlib.contextmanager
def temp_bigquery_table(project, prefix='yaml_bq_it_'):
"""Context manager to create and clean up a temporary BigQuery dataset.
diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml
b/sdks/python/apache_beam/yaml/standard_io.yaml
index 253573c8069..173236dc5e9 100644
--- a/sdks/python/apache_beam/yaml/standard_io.yaml
+++ b/sdks/python/apache_beam/yaml/standard_io.yaml
@@ -514,6 +514,33 @@
config:
gradle_target:
'sdks:java:io:google-cloud-platform:expansion-service:shadowJar'
+# Firestore
+- type: renaming
+ transforms:
+ 'ReadFromFirestore': 'ReadFromFirestore'
+ 'WriteToFirestore': 'WriteToFirestore'
+ config:
+ mappings:
+ 'ReadFromFirestore':
+ project: 'project_id'
+ database: 'database_id'
+ collection: 'collection_id'
+ schema: 'schema'
+ error_handling: 'error_handling'
+ 'WriteToFirestore':
+ project: 'project_id'
+ database: 'database_id'
+ collection: 'collection_id'
+ document_id_field: 'document_id_field'
+ error_handling: 'error_handling'
+ underlying_provider:
+ type: beamJar
+ transforms:
+ 'ReadFromFirestore':
'beam:schematransform:org.apache.beam:firestore_read:v1'
+ 'WriteToFirestore':
'beam:schematransform:org.apache.beam:firestore_write:v1'
+ config:
+ gradle_target:
'sdks:java:io:google-cloud-platform:expansion-service:shadowJar'
+
# TFRecord
- type: renaming
transforms:
diff --git a/sdks/python/build.gradle b/sdks/python/build.gradle
index 339d87c76cb..b445d042bbc 100644
--- a/sdks/python/build.gradle
+++ b/sdks/python/build.gradle
@@ -157,8 +157,8 @@ tasks.register("yamlIntegrationTests") {
// grep -oh 'sdk.*Jar' sdks/python/apache_beam/yaml/*.yaml | sort | uniq
dependsOn ":sdks:java:extensions:schemaio-expansion-service:shadowJar"
dependsOn ":sdks:java:extensions:sql:expansion-service:shadowJar"
- dependsOn ":sdks:java:io:expansion-service:build"
- dependsOn ":sdks:java:io:google-cloud-platform:expansion-service:build"
+ dependsOn ":sdks:java:io:expansion-service:shadowJar"
+ dependsOn ":sdks:java:io:google-cloud-platform:expansion-service:shadowJar"
dependsOn ":sdks:java:io:messaging-expansion-service:shadowJar"
doLast {
@@ -177,8 +177,8 @@ tasks.register("postCommitYamlIntegrationTests") {
// grep -oh 'sdk.*Jar' sdks/python/apache_beam/yaml/*.yaml | sort | uniq
dependsOn ":sdks:java:extensions:schemaio-expansion-service:shadowJar"
dependsOn ":sdks:java:extensions:sql:expansion-service:shadowJar"
- dependsOn ":sdks:java:io:expansion-service:build"
- dependsOn ":sdks:java:io:google-cloud-platform:expansion-service:build"
+ dependsOn ":sdks:java:io:expansion-service:shadowJar"
+ dependsOn ":sdks:java:io:google-cloud-platform:expansion-service:shadowJar"
dependsOn ":sdks:java:io:debezium:expansion-service:shadowJar"
dependsOn ":sdks:java:io:snowflake:expansion-service:shadowJar"
dependsOn ":sdks:java:io:amazon-web-services2:expansion-service:shadowJar"
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
index 5b02dc9ab0e..601b3fcffee 100644
--- a/sdks/python/setup.py
+++ b/sdks/python/setup.py
@@ -557,6 +557,10 @@ if __name__ == '__main__':
'orjson>=3.9.7,<4',
'regex>=2020.6.8',
],
+ # GCP packages used only by tests/ITs (not native Python IO).
+ 'gcp_test': [
+ 'google-cloud-firestore>=2.0.0,<3',
+ ],
'interactive': [
'facets-overview>=1.1.0,<2',
'google-cloud-dataproc>=5.0.0,<6',
