diff --git a/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json b/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json index 8ed972c9f579..9cc78c7d1c6c 100644 --- a/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json +++ b/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "revision": 3 + "revision": 4 } 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 000000000000..0a1e0e20e59b --- /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 000000000000..c035a53acda8 --- /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 { + + private static final String OUTPUT_TAG_NAME = "output"; + public static final TupleTag OUTPUT_TAG = new TupleTag() {}; + public static final TupleTag ERROR_TAG = new TupleTag() {}; + + 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 inputCollectionNames() { + return Collections.emptyList(); + } + + @Override + public List 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 requests = + input + .getPipeline() + .apply("CreateCollectionId", Create.of(configuration.getCollectionId())) + .apply( + "BuildListDocumentsRequest", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement( + @Element String collectionId, + OutputReceiver out) { + out.output( + ListDocumentsRequest.newBuilder() + .setParent(parent) + .setCollectionId(collectionId) + .build()); + } + })); + + FirestoreV1.ListDocuments.Builder readBuilder = + FirestoreIO.v1() + .read() + .listDocuments() + .withProjectId(projectId) + .withDatabaseId(databaseId); + + PCollection 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 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 { + 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 000000000000..0f754846013e --- /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 values = new HashMap<>(); + for (Map.Entry 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 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 castToStringKeyMap(Map map) { + Map 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 map = new HashMap<>(); + for (Map.Entry 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 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 3f22e636e8ab..cf0f365fdc69 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 PCollection expand(PCollection input) { "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 84e1cb1be0ac..7c8f2bf345f6 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 @@ abstract static class BaseFirestoreV1ReadFn 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 @@ protected BaseFirestoreV1ReadFn( 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 @@ public void setup() { @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 000000000000..1f56a7b74447 --- /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 000000000000..4ff724773568 --- /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 { + + private static final String INPUT_TAG = "input"; + private static final String DEFAULT_DOCUMENT_ID_FIELD = "document_id"; + public static final TupleTag OUTPUT_TAG = new TupleTag() {}; + public static final TupleTag ERROR_TAG = new TupleTag() {}; + + 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 inputCollectionNames() { + return Collections.singletonList(INPUT_TAG); + } + + @Override + public List 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 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 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 { + 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 000000000000..7d58ce8d5517 --- /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 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 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 inputRows = pipeline.apply("CreateInvalidRows", Create.empty(schema)); + assertThrows( + IllegalArgumentException.class, + () -> transform.expand(PCollectionRowTuple.of("input", inputRows))); + } + + private static SchemaTransformProvider loadReadProvider() { + List 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 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 000000000000..b32075a00cc2 --- /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 000000000000..19f2574d0575 --- /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 32794f4588d2..de2d34edc0d8 100644 --- a/sdks/python/apache_beam/yaml/integration_tests.py +++ b/sdks/python/apache_beam/yaml/integration_tests.py @@ -165,6 +165,55 @@ def temp_spanner_table(project, prefix='temp_spanner_db_'): spanner_client._delete_database() +@contextlib.contextmanager +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``. + """ + from google.cloud import firestore + + 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 253573c8069a..173236dc5e94 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 75afa5d7c968..299a578e58eb 100644 --- a/sdks/python/build.gradle +++ b/sdks/python/build.gradle @@ -149,8 +149,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 { @@ -169,8 +169,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 5b02dc9ab0ef..cd7e98b45944 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -535,6 +535,7 @@ def get_portability_package_data(): # errors raised during async flushes instead of swallowing them. 'google-cloud-bigtable>=2.42.0,<3', 'google-cloud-build>=3.35.0,<4', + 'google-cloud-firestore>=2.0.0,<3', 'google-cloud-spanner>=3.0.0,<4', # GCP Packages required by ML functionality 'google-cloud-dlp>=3.0.0,<4',