diff --git a/.github/trigger_files/beam_PreCommit_Java_Couchbase_IO_Direct.json b/.github/trigger_files/beam_PreCommit_Java_Couchbase_IO_Direct.json new file mode 100644 index 000000000000..e3d6056a5de9 --- /dev/null +++ b/.github/trigger_files/beam_PreCommit_Java_Couchbase_IO_Direct.json @@ -0,0 +1,4 @@ +{ + "comment": "Modify this file in a trivial way to cause this test suite to run", + "modification": 1 +} diff --git a/.github/workflows/README.md b/.github/workflows/README.md index 178fb6502680..9b681a4da1e0 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -215,6 +215,7 @@ PreCommit Jobs run in a schedule and also get triggered in a PR if relevant sour | [ PreCommit Java Cassandra IO Direct ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Cassandra_IO_Direct.yml) | N/A |`beam_PreCommit_Java_Cassandra_IO_Direct.json`| [![.github/workflows/beam_PreCommit_Java_Cassandra_IO_Direct.yml](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Cassandra_IO_Direct.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Cassandra_IO_Direct.yml?query=event%3Aschedule) | | [ PreCommit Java Cdap IO Direct ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Cdap_IO_Direct.yml) | N/A |`beam_PreCommit_Java_Cdap_IO_Direct.json`| [![.github/workflows/beam_PreCommit_Java_Cdap_IO_Direct.yml](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Cdap_IO_Direct.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Cdap_IO_Direct.yml?query=event%3Aschedule) | | [ PreCommit Java Clickhouse IO Direct ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Clickhouse_IO_Direct.yml) | N/A |`beam_PreCommit_Java_Clickhouse_IO_Direct.json`| [![.github/workflows/beam_PreCommit_Java_Clickhouse_IO_Direct.yml](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Clickhouse_IO_Direct.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Clickhouse_IO_Direct.yml?query=event%3Aschedule) | +| [ PreCommit Java Couchbase IO Direct ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Couchbase_IO_Direct.yml) | N/A |`beam_PreCommit_Java_Couchbase_IO_Direct.json`| [![.github/workflows/beam_PreCommit_Java_Couchbase_IO_Direct.yml](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Couchbase_IO_Direct.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Couchbase_IO_Direct.yml?query=event%3Aschedule) | | [ PreCommit Java Csv IO Direct ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Csv_IO_Direct.yml) | N/A |`beam_PreCommit_Java_Csv_IO_Direct.json`| [![.github/workflows/beam_PreCommit_Java_Csv_IO_Direct.yml](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Csv_IO_Direct.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Csv_IO_Direct.yml?query=event%3Aschedule) | | [ PreCommit Java Datadog IO Direct ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Datadog_IO_Direct.yml) | N/A |`beam_PreCommit_Java_Datadog_IO_Direct.json`| [![.github/workflows/beam_PreCommit_Java_Datadog_IO_Direct.yml](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Datadog_IO_Direct.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Datadog_IO_Direct.yml?query=event%3Aschedule) | | [ PreCommit Java Delta IO Direct ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Delta_IO_Direct.yml) | N/A |`beam_PreCommit_Java_Delta_IO_Direct.json`| [![.github/workflows/beam_PreCommit_Java_Delta_IO_Direct.yml](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Delta_IO_Direct.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Delta_IO_Direct.yml?query=event%3Aschedule) | diff --git a/.github/workflows/beam_PreCommit_Java_Couchbase_IO_Direct.yml b/.github/workflows/beam_PreCommit_Java_Couchbase_IO_Direct.yml new file mode 100644 index 000000000000..695314127b19 --- /dev/null +++ b/.github/workflows/beam_PreCommit_Java_Couchbase_IO_Direct.yml @@ -0,0 +1,115 @@ +# 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. + +name: PreCommit Java Couchbase IO Direct + +on: + push: + tags: ['v*'] + branches: ['master', 'release-*'] + paths: + - "sdks/java/io/couchbase/**" + - ".github/workflows/beam_PreCommit_Java_Couchbase_IO_Direct.yml" + pull_request: + branches: ['master', 'release-*'] + paths: + - "sdks/java/io/couchbase/**" + - 'release/trigger_all_tests.json' + - '.github/trigger_files/beam_PreCommit_Java_Couchbase_IO_Direct.json' + schedule: + - cron: '15 1/6 * * *' + workflow_dispatch: + +# Setting explicit permissions for the action +permissions: + actions: write + pull-requests: write + checks: write + contents: read + deployments: read + id-token: none + issues: write + discussions: read + packages: read + pages: read + repository-projects: read + security-events: read + statuses: read + +# This allows a subsequently queued workflow run to interrupt previous runs +concurrency: + group: '${{ github.workflow }} @ ${{ github.event.pull_request.number || github.event.pull_request.head.label || github.sha || github.head_ref || github.ref }}-${{ github.event.schedule || github.event.sender.login }}' + cancel-in-progress: true + +env: + DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }} + GRADLE_ENTERPRISE_CACHE_USERNAME: ${{ secrets.GE_CACHE_USERNAME }} + GRADLE_ENTERPRISE_CACHE_PASSWORD: ${{ secrets.GE_CACHE_PASSWORD }} + +jobs: + beam_PreCommit_Java_Couchbase_IO_Direct: + name: ${{ matrix.job_name }} + strategy: + matrix: + job_name: ["beam_PreCommit_Java_Couchbase_IO_Direct"] + timeout-minutes: 60 + if: | + github.event_name == 'push' || + github.event_name == 'pull_request' || + (github.event_name == 'schedule' && github.repository == 'apache/beam') || + github.event_name == 'workflow_dispatch' + runs-on: [self-hosted, ubuntu-24.04, main] + steps: + - uses: actions/checkout@v7 + with: + persist-credentials: false + - name: Setup repository + uses: $/.github/actions/setup-action + - name: Setup environment + uses: $/.github/actions/setup-environment-action + - name: run Couchbase IO build script + uses: $/.github/actions/gradle-command-self-hosted-action + with: + gradle-command: :sdks:java:io:couchbase:build + arguments: | + -PdisableSpotlessCheck=true \ + -PdisableCheckStyle=true \ + - name: Archive JUnit Test Results + uses: actions/upload-artifact@v7 + if: ${{ !success() }} + with: + name: JUnit Test Results + path: "**/build/reports/tests/" + - name: Publish JUnit Test Results + uses: EnricoMi/publish-unit-test-result-action@v2 + if: always() + with: + commit: '${{ env.GITHUB_SHA }}' + comment_mode: 'off' + check_run: ${{ github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository }} + files: '**/build/test-results/**/*.xml' + large_files: true + - name: Archive SpotBugs Results + uses: actions/upload-artifact@v7 + if: always() + with: + name: SpotBugs Results + path: '**/build/reports/spotbugs/*.html' + - name: Publish SpotBugs Results + uses: jwgmeligmeyling/spotbugs-github-action@v1.2 + if: always() && (github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository) + with: + name: Publish SpotBugs + path: '**/build/reports/spotbugs/*.html' diff --git a/CHANGES.md b/CHANGES.md index 8fa59e7d5ea0..76714370119f 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -64,6 +64,7 @@ ## I/Os +* Added CouchbaseIO for reading with SQL++ queries and writing documents (Java) ([#18381](https://github.com/apache/beam/issues/18381)). * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). * (Java) FirestoreIO: added `RpcQosOptions.Builder.withRampupThrottlingDisabled()`, mirroring the existing DatastoreIO option, to opt out of the 500/50/5 write ramp-up ([#27170](https://github.com/apache/beam/issues/27170)). * Added support for attributing BigQuery API quota and billing to a specific GCP project (quota project): `quota_project_id` parameter in `ReadFromBigQuery` or `--quota_project_id` pipeline option, honored by the Python BigQuery IO transforms (`WriteToBigQuery` with `FILE_LOADS` or `STREAMING_INSERTS`; not yet `STORAGE_WRITE_API`, which is Java-backed), `--bigQueryQuotaProjectId` pipeline option (Java), and `bigqueryio.WithQuotaProject` read/query option (Go) ([#37431](https://github.com/apache/beam/issues/37431)). diff --git a/build.gradle.kts b/build.gradle.kts index 1cefe1ac5b93..709ba1188149 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -349,6 +349,7 @@ tasks.register("javaioPreCommit") { dependsOn(":sdks:java:io:csv:build") dependsOn(":sdks:java:io:cdap:build") dependsOn(":sdks:java:io:clickhouse:build") + dependsOn(":sdks:java:io:couchbase:build") dependsOn(":sdks:java:io:datadog:build") dependsOn(":sdks:java:io:debezium:expansion-service:build") dependsOn(":sdks:java:io:debezium:build") diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy index 1e1c9108de76..4177780012c4 100644 --- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy +++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy @@ -743,6 +743,7 @@ class BeamModulePlugin implements Plugin { commons_lang3 : "org.apache.commons:commons-lang3:3.18.0", commons_logging : "commons-logging:commons-logging:1.2", commons_math3 : "org.apache.commons:commons-math3:3.6.1", + couchbase_java_client : "com.couchbase.client:java-client:3.12.2", dbcp2 : "org.apache.commons:commons-dbcp2:$dbcp2_version", delta_kernel_api : "io.delta:delta-kernel-api:$delta_lake_version", delta_kernel_defaults : "io.delta:delta-kernel-defaults:$delta_lake_version", @@ -934,6 +935,7 @@ class BeamModulePlugin implements Plugin { testcontainers_base : "org.testcontainers:testcontainers:$testcontainers_version", testcontainers_cassandra : "org.testcontainers:cassandra:$testcontainers_version", testcontainers_clickhouse : "org.testcontainers:clickhouse:$testcontainers_version", + testcontainers_couchbase : "org.testcontainers:couchbase:$testcontainers_version", testcontainers_elasticsearch : "org.testcontainers:elasticsearch:$testcontainers_version", testcontainers_gcloud : "org.testcontainers:gcloud:$testcontainers_version", testcontainers_jdbc : "org.testcontainers:jdbc:$testcontainers_version", diff --git a/sdks/java/io/couchbase/build.gradle b/sdks/java/io/couchbase/build.gradle new file mode 100644 index 000000000000..e3196720e6f1 --- /dev/null +++ b/sdks/java/io/couchbase/build.gradle @@ -0,0 +1,39 @@ +/* + * 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. + */ + +plugins { id 'org.apache.beam.module' } +applyJavaNature( automaticModuleName: 'org.apache.beam.sdk.io.couchbase') + +description = "Apache Beam :: SDKs :: Java :: IO :: Couchbase" +ext.summary = "IO to read and write on Couchbase." + +dependencies { + implementation project(path: ":sdks:java:core", configuration: "shadow") + implementation library.java.couchbase_java_client + // Same versions as java-client's transitive dependencies; used directly in ReadFn and WriteFn. + implementation "io.projectreactor:reactor-core:3.6.9" + implementation "org.reactivestreams:reactive-streams:1.0.4" + implementation library.java.jackson_databind + implementation library.java.vendored_guava_32_1_2_jre + testImplementation project(path: ":sdks:java:core", configuration: "shadowTest") + testImplementation library.java.junit + testImplementation library.java.hamcrest + testImplementation library.java.testcontainers_couchbase + testRuntimeOnly library.java.slf4j_jdk14 + testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow") +} diff --git a/sdks/java/io/couchbase/src/main/java/org/apache/beam/sdk/io/couchbase/CouchbaseIO.java b/sdks/java/io/couchbase/src/main/java/org/apache/beam/sdk/io/couchbase/CouchbaseIO.java new file mode 100644 index 000000000000..750a3db90855 --- /dev/null +++ b/sdks/java/io/couchbase/src/main/java/org/apache/beam/sdk/io/couchbase/CouchbaseIO.java @@ -0,0 +1,713 @@ +/* + * 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.couchbase; + +import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; + +import com.couchbase.client.java.Bucket; +import com.couchbase.client.java.Cluster; +import com.couchbase.client.java.ClusterOptions; +import com.couchbase.client.java.ReactiveCollection; +import com.couchbase.client.java.json.JsonArray; +import com.couchbase.client.java.json.JsonObject; +import com.couchbase.client.java.query.QueryOptions; +import com.couchbase.client.java.query.QueryScanConsistency; +import com.couchbase.client.java.query.ReactiveQueryResult; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.auto.value.AutoValue; +import java.io.Serializable; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.stream.Stream; +import org.apache.beam.sdk.coders.CannotProvideCoderException; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderRegistry; +import org.apache.beam.sdk.coders.RowCoder; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Redistribute; +import org.apache.beam.sdk.transforms.SerializableFunction; +import org.apache.beam.sdk.transforms.SerializableFunctions; +import org.apache.beam.sdk.transforms.display.DisplayData; +import org.apache.beam.sdk.util.RowJson.RowJsonDeserializer; +import org.apache.beam.sdk.util.RowJsonUtils; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PBegin; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PDone; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TypeDescriptors; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; +import org.checkerframework.checker.nullness.qual.MonotonicNonNull; +import org.checkerframework.checker.nullness.qual.Nullable; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +/** + * IO to read and write on Couchbase. + * + *

Supports Couchbase Server 8.0 and later, using the Couchbase Java SDK 3.x. Use a {@code + * couchbases://} connection string to connect over TLS. + * + *

Reading from Couchbase

+ * + *

{@link CouchbaseIO#read()} runs a SQL++ query and emits each result row as a {@link + * JsonObject}. If the {@link ConnectionConfiguration} sets a bucket, the query runs in its scope + * and can name collections without qualifying them: + * + * 

{@code
+ * PCollection airlines =
+ *     pipeline.apply(CouchbaseIO.read()
+ *         .withConnectionConfiguration(
+ *             CouchbaseIO.ConnectionConfiguration.create("couchbase://host")
+ *                 .withUsername(username)
+ *                 .withPassword(password)
+ *                 .withBucket("travel")
+ *                 .withScope("inventory"))
+ *         .withQuery("SELECT a.* FROM airline a WHERE a.country = $country")
+ *         .withNamedParameter("country", "United Kingdom"));
+ * }
+ * + *

Pass values with {@link Read#withNamedParameter} or {@link Read#withPositionalParameters} + * rather than building them into the query string. Use {@link Read#withScanConsistency} with {@link + * QueryScanConsistency#REQUEST_PLUS} when the read must see all prior mutations. + * + *

Use {@link Read#withParseFn} to convert each row to another type, and {@link Read#withCoder} + * to set its coder when Beam can't infer one: + * + * 

{@code
+ * PCollection names =
+ *     pipeline.apply(CouchbaseIO.read()
+ *         .withConnectionConfiguration(config)
+ *         .withQuery("SELECT a.name FROM airline a")
+ *         .withParseFn(row -> row.getString("name"))
+ *         .withCoder(StringUtf8Coder.of()));
+ * }
+ * + *

{@link CouchbaseIO#readRows(Schema)} emits {@link Row}s with the given schema, so the output + * can be used with Beam SQL and other schema-aware transforms. Fields not in the schema are + * ignored; a field missing from a row is null if the schema declares it nullable, otherwise the + * read fails. + * + *

The query runs as a single request on one worker, because the client can't split a SQL++ + * query. Use {@link Read#withRedistribute} to spread the output across workers for the steps that + * follow. If the read is retried, the query runs again from the start; with the default scan + * consistency the retry may see a different result. + * + *

The user needs the Query Select role on the collections the query reads. + * + *

Writing to Couchbase

+ * + *

{@link CouchbaseIO#write()} upserts each element of a {@link PCollection} as a JSON document + * into a collection. You provide a {@link ConnectionConfiguration}, a function that returns the + * document ID for an element, and a function that returns its content as a {@link JsonObject}: + * + * 

{@code
+ * pipeline
+ *     .apply(...)
+ *     .apply(CouchbaseIO.write()
+ *         .withConnectionConfiguration(
+ *             CouchbaseIO.ConnectionConfiguration.create("couchbase://host")
+ *                 .withUsername(username)
+ *                 .withPassword(password)
+ *                 .withBucket("travel")
+ *                 .withScope("inventory")
+ *                 .withCollection("airline"))
+ *         .withIdFn(airline -> airline.getId())
+ *         .withContentFn(airline -> JsonObject.create().put("name", airline.getName())));
+ * }
+ * + *

Documents are written with UPSERT, so a document that already exists is overwritten. If a + * bundle fails, the runner retries it; documents written by the failed attempt are overwritten with + * the same content, so retries do not produce duplicates. Writes are not exactly-once: a document + * may be written more than once, but the end state is the same as writing it once. + * + *

The user needs the Data Writer role on the target collection. + * + *

Elements are buffered and written in batches of {@value #DEFAULT_BATCH_SIZE} (see {@link + * Write#withBatchSize}), with at most {@value #DEFAULT_MAX_CONCURRENCY} upserts in flight per batch + * (see {@link Write#withMaxConcurrency}). + * + *

The password is serialized with the pipeline. Read it from a secret manager when building the + * pipeline rather than hardcoding it. + */ +public class CouchbaseIO { + + static final int DEFAULT_BATCH_SIZE = 500; + static final int DEFAULT_MAX_CONCURRENCY = 32; + + private CouchbaseIO() {} + + /** + * Returns a transform that runs a SQL++ query and emits each result row as a {@link JsonObject}. + * Use {@link Read#withParseFn} to emit another type. + */ + public static Read read() { + return new AutoValue_CouchbaseIO_Read.Builder() + .setScanConsistency(QueryScanConsistency.NOT_BOUNDED) + .setParseFn(SerializableFunctions.identity()) + .setCoder(JsonObjectCoder.of()) + .setRedistribute(false) + .build(); + } + + /** + * Returns a transform that runs a SQL++ query and converts each result row to a {@link Row} with + * the given schema. Fields not in the schema are ignored. A field missing from the row is null if + * the schema declares it nullable; otherwise the read fails. A value that doesn't match its field + * type also fails the read, with an error naming the field. + * + *

Throws if the schema has a field type that can't be read from JSON, such as {@code BYTES}. + */ + public static Read readRows(Schema schema) { + return read().withParseFn(new JsonToRowFn(schema)).withCoder(RowCoder.of(schema)); + } + + /** + * Returns a transform that upserts each element as a JSON document. Defaults: batch size {@value + * #DEFAULT_BATCH_SIZE}, max concurrency {@value #DEFAULT_MAX_CONCURRENCY}. + */ + public static Write write() { + return new AutoValue_CouchbaseIO_Write.Builder() + .setBatchSize(DEFAULT_BATCH_SIZE) + .setMaxConcurrency(DEFAULT_MAX_CONCURRENCY) + .build(); + } + + @AutoValue + public abstract static class ConnectionConfiguration implements Serializable { + abstract String connectionString(); + + abstract @Nullable String username(); + + abstract @Nullable String password(); + + abstract @Nullable String bucket(); + + abstract String scope(); // default "_default" + + abstract String collection(); // default "_default" + + abstract Builder toBuilder(); + + @AutoValue.Builder + abstract static class Builder { + abstract Builder setConnectionString(String connectionString); + + abstract Builder setUsername(String username); + + abstract Builder setPassword(String password); + + abstract Builder setBucket(String bucket); + + abstract Builder setScope(String scope); + + abstract Builder setCollection(String collection); + + abstract ConnectionConfiguration build(); + } + + public static ConnectionConfiguration create(String connectionString) { + return new AutoValue_CouchbaseIO_ConnectionConfiguration.Builder() + .setConnectionString(connectionString) + .setScope("_default") + .setCollection("_default") + .build(); + } + + /** Sets the username used to authenticate to the cluster. */ + public ConnectionConfiguration withUsername(String username) { + return toBuilder().setUsername(username).build(); + } + + /** Sets the password used to authenticate to the cluster. */ + public ConnectionConfiguration withPassword(String password) { + return toBuilder().setPassword(password).build(); + } + + /** Sets the bucket to read from or write to. */ + public ConnectionConfiguration withBucket(String bucket) { + return toBuilder().setBucket(bucket).build(); + } + + /** Sets the scope within the bucket. Defaults to {@code _default}. */ + public ConnectionConfiguration withScope(String scope) { + return toBuilder().setScope(scope).build(); + } + + /** Sets the collection within the scope. Defaults to {@code _default}. */ + public ConnectionConfiguration withCollection(String collection) { + return toBuilder().setCollection(collection).build(); + } + + // Omits the collection, which only Write uses. + void populateDisplayData(DisplayData.Builder builder) { + builder.add(DisplayData.item("connectionString", connectionString())); + builder.addIfNotNull(DisplayData.item("username", username())); + builder.addIfNotNull(DisplayData.item("bucket", bucket())); + builder.add(DisplayData.item("scope", scope())); + } + + // Hand-written so the password is never logged; AutoValue's generated toString includes it. + @Override + public final String toString() { + return MoreObjects.toStringHelper(this) + .add("connectionString", connectionString()) + .add("username", username()) + .add("bucket", bucket()) + .add("scope", scope()) + .add("collection", collection()) + .toString(); + } + + Cluster connect() { + String user = checkStateNotNull(username(), "username is required"); + String pass = checkStateNotNull(password(), "password is required"); + return Cluster.connect(connectionString(), ClusterOptions.clusterOptions(user, pass)); + } + } + + /** A {@link PTransform} that runs a SQL++ query and emits one element per result row. */ + @AutoValue + public abstract static class Read extends PTransform> { + abstract @Nullable ConnectionConfiguration connectionConfiguration(); + + abstract @Nullable String query(); + + abstract @Nullable JsonObject namedParameters(); + + abstract @Nullable JsonArray positionalParameters(); + + abstract QueryScanConsistency scanConsistency(); + + abstract SerializableFunction parseFn(); + + abstract @Nullable Coder coder(); + + abstract boolean redistribute(); + + abstract @Nullable Integer redistributeNumKeys(); + + abstract Builder toBuilder(); + + @AutoValue.Builder + abstract static class Builder { + abstract Builder setConnectionConfiguration(@Nullable ConnectionConfiguration config); + + abstract Builder setQuery(@Nullable String query); + + abstract Builder setNamedParameters(@Nullable JsonObject namedParameters); + + abstract Builder setPositionalParameters(@Nullable JsonArray positionalParameters); + + abstract Builder setScanConsistency(QueryScanConsistency scanConsistency); + + abstract Builder setParseFn(SerializableFunction parseFn); + + abstract Builder setCoder(@Nullable Coder coder); + + abstract Builder setRedistribute(boolean redistribute); + + abstract Builder setRedistributeNumKeys(@Nullable Integer redistributeNumKeys); + + abstract Read build(); + } + + /** + * Sets the cluster and credentials. If a bucket is set, the query runs in the configured scope, + * so it can name collections without qualifying them; otherwise it runs at cluster level. The + * configured collection is not used: the query names what it reads. + */ + public Read withConnectionConfiguration(ConnectionConfiguration config) { + return toBuilder().setConnectionConfiguration(config).build(); + } + + /** + * Sets the SQL++ query. Include {@code META().id} in the projection if the document ID is + * needed. + */ + public Read withQuery(String query) { + return toBuilder().setQuery(query).build(); + } + + /** Adds a value for the named parameter {@code $name} in the query. */ + public Read withNamedParameter(String name, Object value) { + JsonObject current = namedParameters(); + JsonObject parameters = + current == null ? JsonObject.create() : JsonObject.from(current.toMap()); + return toBuilder().setNamedParameters(parameters.put(name, value)).build(); + } + + /** Sets the values for the positional parameters {@code $1}, {@code $2}, ... in the query. */ + public Read withPositionalParameters(Object... values) { + return toBuilder().setPositionalParameters(JsonArray.from(values)).build(); + } + + /** + * Sets the scan consistency. Defaults to {@link QueryScanConsistency#NOT_BOUNDED}; use {@link + * QueryScanConsistency#REQUEST_PLUS} when the read must see all prior mutations. + */ + public Read withScanConsistency(QueryScanConsistency scanConsistency) { + return toBuilder().setScanConsistency(scanConsistency).build(); + } + + /** + * Sets the function that converts each result row to the output type. It runs inside the read, + * before any redistribute. Resets any coder set with {@link #withCoder}, because that coder was + * for the previous output type. + */ + public Read withParseFn(SerializableFunction parseFn) { + return new AutoValue_CouchbaseIO_Read.Builder() + .setConnectionConfiguration(connectionConfiguration()) + .setQuery(query()) + .setNamedParameters(namedParameters()) + .setPositionalParameters(positionalParameters()) + .setScanConsistency(scanConsistency()) + .setParseFn(parseFn) + .setRedistribute(redistribute()) + .setRedistributeNumKeys(redistributeNumKeys()) + .build(); + } + + /** + * Sets the output coder. Required when it can't be inferred from {@link #withParseFn}, for + * example when the parsing function is a lambda. + */ + public Read withCoder(Coder coder) { + return toBuilder().setCoder(coder).build(); + } + + /** + * Redistributes the output after parsing, so the shuffle carries the parsed type. + * Redistributing the default {@link JsonObject} output is costly; use {@link #withParseFn} + * first. + */ + public Read withRedistribute() { + return toBuilder().setRedistribute(true).build(); + } + + /** Sets the number of keys to redistribute into. Requires {@link #withRedistribute()}. */ + public Read withRedistributeNumKeys(int redistributeNumKeys) { + checkArgument( + redistributeNumKeys > 0, "redistributeNumKeys must be > 0, got %s", redistributeNumKeys); + return toBuilder().setRedistributeNumKeys(redistributeNumKeys).build(); + } + + @Override + public PCollection expand(PBegin input) { + ConnectionConfiguration config = + checkArgumentNotNull( + connectionConfiguration(), "withConnectionConfiguration() is required"); + checkArgument(query() != null, "withQuery() is required"); + checkArgument( + config.username() != null, "ConnectionConfiguration.withUsername() is required"); + checkArgument( + config.password() != null, "ConnectionConfiguration.withPassword() is required"); + checkArgument( + namedParameters() == null || positionalParameters() == null, + "withNamedParameter() and withPositionalParameters() cannot be combined"); + checkArgument( + redistributeNumKeys() == null || redistribute(), + "withRedistributeNumKeys() requires withRedistribute()"); + + PCollection output = + input + .apply("Impulse", Impulse.create()) + .apply("Query", ParDo.of(new ReadFn<>(this))) + .setCoder(outputCoder(input.getPipeline().getCoderRegistry())); + if (redistribute()) { + output = + output.apply( + "Redistribute", + Redistribute.arbitrarily().withNumBuckets(redistributeNumKeys())); + } + return output; + } + + private Coder outputCoder(CoderRegistry registry) { + Coder coder = coder(); + if (coder != null) { + return coder; + } + try { + return registry.getCoder(TypeDescriptors.outputOf(parseFn())); + } catch (CannotProvideCoderException e) { + throw new IllegalArgumentException( + "Unable to infer a coder for the output of withParseFn(); set one with withCoder()", e); + } + } + + @Override + public void populateDisplayData(DisplayData.Builder builder) { + super.populateDisplayData(builder); + ConnectionConfiguration config = connectionConfiguration(); + if (config != null) { + config.populateDisplayData(builder); + } + builder.addIfNotNull(DisplayData.item("query", query())); + builder.add(DisplayData.item("scanConsistency", scanConsistency().name())); + builder.add(DisplayData.item("redistribute", redistribute())); + builder.addIfNotNull(DisplayData.item("redistributeNumKeys", redistributeNumKeys())); + } + } + + // Runs the whole query in one call: a SQL++ query can't be split by the client. A retried + // bundle runs it again from the start. + static class ReadFn extends DoFn { + private final ConnectionConfiguration config; + private final String query; + private final @Nullable JsonObject namedParameters; + private final @Nullable JsonArray positionalParameters; + private final QueryScanConsistency scanConsistency; + private final SerializableFunction parseFn; + + private transient @MonotonicNonNull Cluster cluster; + + ReadFn(Read spec) { + this.config = checkStateNotNull(spec.connectionConfiguration()); + this.query = checkStateNotNull(spec.query()); + this.namedParameters = spec.namedParameters(); + this.positionalParameters = spec.positionalParameters(); + this.scanConsistency = spec.scanConsistency(); + this.parseFn = spec.parseFn(); + } + + @Setup + public void setup() { + cluster = config.connect(); + } + + @ProcessElement + public void processElement(OutputReceiver out) { + Cluster connected = checkStateNotNull(cluster); + QueryOptions options = QueryOptions.queryOptions().scanConsistency(scanConsistency); + if (namedParameters != null) { + options.parameters(namedParameters); + } else if (positionalParameters != null) { + options.parameters(positionalParameters); + } + String bucket = config.bucket(); + Mono result = + bucket == null + ? connected.reactive().query(query, options) + : connected.bucket(bucket).scope(config.scope()).reactive().query(query, options); + // toStream() pulls rows with backpressure on this thread, so the full result set is never + // buffered; closing it cancels the query if parsing fails part way. + try (Stream rows = + result.flatMapMany(ReactiveQueryResult::rowsAsObject).toStream()) { + rows.forEach(row -> out.output(parseFn.apply(row))); + } + } + + @Teardown + public void teardown() { + if (cluster != null) { + cluster.disconnect(); + } + } + } + + static class JsonToRowFn implements SerializableFunction { + private final RowJsonDeserializer deserializer; + + // ObjectMapper isn't serializable, so each worker builds its own. + private transient @MonotonicNonNull ObjectMapper mapper; + + JsonToRowFn(Schema schema) { + this.deserializer = RowJsonDeserializer.forSchema(schema); + } + + @Override + public Row apply(JsonObject json) { + ObjectMapper m = mapper; + if (m == null) { + m = RowJsonUtils.newObjectMapperWith(deserializer); + mapper = m; + } + return RowJsonUtils.jsonToRow(m, json.toString()); + } + } + + /** A {@link PTransform} that upserts elements into a Couchbase collection. */ + @AutoValue + public abstract static class Write extends PTransform, PDone> { + abstract @Nullable ConnectionConfiguration connectionConfiguration(); + + abstract @Nullable SerializableFunction idFn(); + + abstract @Nullable SerializableFunction contentFn(); + + abstract int batchSize(); + + abstract int maxConcurrency(); + + abstract Builder toBuilder(); + + @AutoValue.Builder + abstract static class Builder { + abstract Builder setConnectionConfiguration(ConnectionConfiguration config); + + abstract Builder setIdFn(SerializableFunction idFn); + + abstract Builder setContentFn(SerializableFunction contentFn); + + abstract Builder setBatchSize(int batchSize); + + abstract Builder setMaxConcurrency(int maxConcurrency); + + abstract Write build(); + } + + /** Sets the cluster, credentials and target collection. */ + public Write withConnectionConfiguration(ConnectionConfiguration config) { + return toBuilder().setConnectionConfiguration(config).build(); + } + + /** Sets the function that returns the document ID for an element. */ + public Write withIdFn(SerializableFunction idFn) { + return toBuilder().setIdFn(idFn).build(); + } + + /** + * Sets the function that returns the document content for an element. It must return a new + * {@link JsonObject} on each call, because the object is buffered and written later. + */ + public Write withContentFn(SerializableFunction contentFn) { + return toBuilder().setContentFn(contentFn).build(); + } + + /** Sets the number of documents buffered before they are written. Must be positive. */ + public Write withBatchSize(int batchSize) { + checkArgument(batchSize > 0, "batchSize must be > 0, got %s", batchSize); + return toBuilder().setBatchSize(batchSize).build(); + } + + /** Sets the maximum number of in-flight upserts per batch. Must be positive. */ + public Write withMaxConcurrency(int maxConcurrency) { + checkArgument(maxConcurrency > 0, "maxConcurrency must be > 0, got %s", maxConcurrency); + return toBuilder().setMaxConcurrency(maxConcurrency).build(); + } + + @Override + public PDone expand(PCollection input) { + ConnectionConfiguration config = + checkArgumentNotNull( + connectionConfiguration(), "withConnectionConfiguration() is required"); + checkArgument(idFn() != null, "withIdFn() is required"); + checkArgument(contentFn() != null, "withContentFn() is required"); + checkArgument(config.bucket() != null, "ConnectionConfiguration.withBucket() is required"); + checkArgument( + config.username() != null, "ConnectionConfiguration.withUsername() is required"); + checkArgument( + config.password() != null, "ConnectionConfiguration.withPassword() is required"); + input.apply(ParDo.of(new WriteFn<>(this))); + return PDone.in(input.getPipeline()); + } + + @Override + public void populateDisplayData(DisplayData.Builder builder) { + super.populateDisplayData(builder); + ConnectionConfiguration config = connectionConfiguration(); + if (config != null) { + config.populateDisplayData(builder); + builder.add(DisplayData.item("collection", config.collection())); + } + builder.add(DisplayData.item("batchSize", batchSize())); + builder.add(DisplayData.item("maxConcurrency", maxConcurrency())); + } + } + + static class WriteFn extends DoFn { + private final ConnectionConfiguration config; + private final SerializableFunction idFn; + private final SerializableFunction contentFn; + private final int batchSize; + private final int maxConcurrency; + + private static final Duration BUCKET_READY_TIMEOUT = Duration.ofSeconds(30); + + private transient @MonotonicNonNull Cluster cluster; + private transient @MonotonicNonNull ReactiveCollection collection; + private transient @MonotonicNonNull List> batch; + + WriteFn(Write spec) { + this.config = checkStateNotNull(spec.connectionConfiguration()); + this.idFn = checkStateNotNull(spec.idFn()); + this.contentFn = checkStateNotNull(spec.contentFn()); + this.batchSize = spec.batchSize(); + this.maxConcurrency = spec.maxConcurrency(); + } + + @Setup + public void setup() { + Cluster connected = config.connect(); + cluster = connected; + Bucket bucket = connected.bucket(checkStateNotNull(config.bucket())); + // Opening a bucket is asynchronous; without this, the first upserts can time out with + // BUCKET_OPEN_IN_PROGRESS on a freshly started worker. + bucket.waitUntilReady(BUCKET_READY_TIMEOUT); + collection = bucket.scope(config.scope()).collection(config.collection()).reactive(); + } + + @StartBundle + public void startBundle() { + batch = new ArrayList<>(batchSize); + } + + @ProcessElement + public void processElement(@Element T element) { + List> buffer = checkStateNotNull(batch); + buffer.add(KV.of(idFn.apply(element), contentFn.apply(element))); + if (buffer.size() >= batchSize) { + flush(buffer); + } + } + + @FinishBundle + public void finishBundle() { + flush(checkStateNotNull(batch)); + } + + @Teardown + public void teardown() { + if (cluster != null) { + cluster.disconnect(); + } + } + + // The first failed upsert fails the bundle and the runner retries it. Upsert makes the + // retry idempotent: documents already written are overwritten with the same content. + private void flush(List> buffer) { + if (buffer.isEmpty()) { + return; + } + ReactiveCollection target = checkStateNotNull(collection); + Flux.fromIterable(buffer) + .flatMap(kv -> target.upsert(kv.getKey(), kv.getValue()), maxConcurrency) + .then() + .block(); + buffer.clear(); + } + } +} diff --git a/sdks/java/io/couchbase/src/main/java/org/apache/beam/sdk/io/couchbase/JsonObjectCoder.java b/sdks/java/io/couchbase/src/main/java/org/apache/beam/sdk/io/couchbase/JsonObjectCoder.java new file mode 100644 index 000000000000..0f574ee44842 --- /dev/null +++ b/sdks/java/io/couchbase/src/main/java/org/apache/beam/sdk/io/couchbase/JsonObjectCoder.java @@ -0,0 +1,64 @@ +/* + * 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.couchbase; + +import com.couchbase.client.java.json.JsonObject; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import org.apache.beam.sdk.coders.AtomicCoder; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.StringUtf8Coder; + +/** + * A {@link Coder} that encodes a {@link JsonObject} as its JSON string in UTF-8. + * + *

The encoding is 2 to 3 times smaller than {@code SerializableCoder}, which Beam would + * otherwise infer, and does not depend on Java serialization of Couchbase SDK classes. + * + *

Decoding preserves the JSON content but not always the Java number type: numbers are decoded + * as {@code Integer} when they fit, so a {@code Long} value of {@code 10} decodes as an {@code + * Integer}. Documents read from Couchbase already use the decoded types. + */ +public class JsonObjectCoder extends AtomicCoder { + + private static final JsonObjectCoder INSTANCE = new JsonObjectCoder(); + private static final StringUtf8Coder STRING_CODER = StringUtf8Coder.of(); + + public static JsonObjectCoder of() { + return INSTANCE; + } + + private JsonObjectCoder() {} + + @Override + public void encode(JsonObject value, OutputStream outStream) throws IOException { + STRING_CODER.encode(value.toString(), outStream); + } + + @Override + public JsonObject decode(InputStream inStream) throws IOException { + return JsonObject.fromJson(STRING_CODER.decode(inStream)); + } + + /** Field order in the encoded JSON is not stable, so equal objects can encode differently. */ + @Override + public void verifyDeterministic() throws NonDeterministicException { + throw new NonDeterministicException(this, "JSON field order is not deterministic"); + } +} diff --git a/sdks/java/io/couchbase/src/main/java/org/apache/beam/sdk/io/couchbase/package-info.java b/sdks/java/io/couchbase/src/main/java/org/apache/beam/sdk/io/couchbase/package-info.java new file mode 100644 index 000000000000..bec8c494afc2 --- /dev/null +++ b/sdks/java/io/couchbase/src/main/java/org/apache/beam/sdk/io/couchbase/package-info.java @@ -0,0 +1,24 @@ +/* + * 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. + */ + +/** + * Transforms for reading from and writing to Couchbase. + * + * @see org.apache.beam.sdk.io.couchbase.CouchbaseIO + */ +package org.apache.beam.sdk.io.couchbase; diff --git a/sdks/java/io/couchbase/src/test/java/org/apache/beam/sdk/io/couchbase/CouchbaseIOReadTest.java b/sdks/java/io/couchbase/src/test/java/org/apache/beam/sdk/io/couchbase/CouchbaseIOReadTest.java new file mode 100644 index 000000000000..ed6bc29a748b --- /dev/null +++ b/sdks/java/io/couchbase/src/test/java/org/apache/beam/sdk/io/couchbase/CouchbaseIOReadTest.java @@ -0,0 +1,289 @@ +/* + * 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.couchbase; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; +import static org.junit.Assert.assertThrows; + +import com.couchbase.client.java.Bucket; +import com.couchbase.client.java.Cluster; +import com.couchbase.client.java.Collection; +import com.couchbase.client.java.json.JsonObject; +import com.couchbase.client.java.query.QueryScanConsistency; +import java.time.Duration; +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.io.couchbase.CouchbaseIO.ConnectionConfiguration; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Count; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.testcontainers.couchbase.BucketDefinition; +import org.testcontainers.couchbase.CouchbaseContainer; + +/** Reads from a single-node Couchbase Server started by Testcontainers. */ +@RunWith(JUnit4.class) +public class CouchbaseIOReadTest { + + private static final String BUCKET = "beam"; + private static final String SCOPE = "app"; + private static final String COLLECTION = "airlines"; + + @ClassRule + public static final CouchbaseContainer COUCHBASE = + new CouchbaseContainer("couchbase/server:enterprise-8.0.3") + // The container creates its index straight after the bucket, which can fail while the + // new cluster is still rebalancing; setUpClass creates the indexes once the bucket is + // ready instead. + .withBucket(new BucketDefinition(BUCKET).withPrimaryIndex(false)); + + private static Cluster seedCluster; + private static Bucket seedBucket; + + @Rule public final transient TestPipeline p = TestPipeline.create(); + + // Only for the round-trip test, which runs a write pipeline before the read pipeline. + @Rule public final transient TestPipeline writePipeline = TestPipeline.create(); + + @BeforeClass + public static void setUpClass() { + seedCluster = + Cluster.connect( + COUCHBASE.getConnectionString(), COUCHBASE.getUsername(), COUCHBASE.getPassword()); + seedBucket = seedCluster.bucket(BUCKET); + seedBucket.waitUntilReady(Duration.ofSeconds(60)); + seedBucket.collections().createScope(SCOPE); + seedBucket.collections().createCollection(SCOPE, COLLECTION); + seedCluster.query(String.format("CREATE PRIMARY INDEX ON `%s`.`_default`.`_default`", BUCKET)); + seedCluster.query( + String.format("CREATE PRIMARY INDEX ON `%s`.`%s`.`%s`", BUCKET, SCOPE, COLLECTION)); + } + + @AfterClass + public static void tearDownClass() { + if (seedCluster != null) { + seedCluster.disconnect(); + } + } + + private static ConnectionConfiguration config() { + return ConnectionConfiguration.create(COUCHBASE.getConnectionString()) + .withUsername(COUCHBASE.getUsername()) + .withPassword(COUCHBASE.getPassword()) + .withBucket(BUCKET); + } + + /** All tests share one bucket, so each tags its documents and queries only its own tag. */ + private static JsonObject doc(String test, int n) { + return JsonObject.create().put("test", test).put("n", n); + } + + private static List seed(Collection collection, String test, int count) { + List docs = + IntStream.range(0, count).mapToObj(n -> doc(test, n)).collect(Collectors.toList()); + for (JsonObject d : docs) { + collection.upsert(test + "-" + d.getInt("n"), d); + } + return docs; + } + + private static CouchbaseIO.Read readTagged(String collection) { + return CouchbaseIO.read() + .withConnectionConfiguration(config()) + .withQuery("SELECT d.* FROM `" + collection + "` d WHERE d.test = $test") + .withScanConsistency(QueryScanConsistency.REQUEST_PLUS); + } + + @Test + public void testReadAsJsonObject() { + List docs = seed(seedBucket.defaultCollection(), "json", 10); + + PCollection read = + p.apply(readTagged("_default").withNamedParameter("test", "json")); + + PAssert.that(read).containsInAnyOrder(docs); + p.run().waitUntilFinish(); + } + + @Test + public void testReadEmptyResult() { + PCollection read = + p.apply(readTagged("_default").withNamedParameter("test", "no-such-test")); + + PAssert.that(read).empty(); + p.run().waitUntilFinish(); + } + + @Test + public void testReadWithParseFnAndPositionalParameters() { + seed(seedBucket.scope(SCOPE).collection(COLLECTION), "parse", 5); + + PCollection ids = + p.apply( + CouchbaseIO.read() + .withConnectionConfiguration(config().withScope(SCOPE)) + .withQuery("SELECT META(a).id AS id FROM " + COLLECTION + " a WHERE a.test = $1") + .withPositionalParameters("parse") + .withScanConsistency(QueryScanConsistency.REQUEST_PLUS) + .withParseFn(row -> row.getString("id")) + .withCoder(StringUtf8Coder.of())); + + PAssert.that(ids).containsInAnyOrder("parse-0", "parse-1", "parse-2", "parse-3", "parse-4"); + p.run().waitUntilFinish(); + } + + @Test + public void testReadAtClusterLevelWithoutBucket() { + List docs = seed(seedBucket.scope(SCOPE).collection(COLLECTION), "cluster", 5); + ConnectionConfiguration noBucket = + ConnectionConfiguration.create(COUCHBASE.getConnectionString()) + .withUsername(COUCHBASE.getUsername()) + .withPassword(COUCHBASE.getPassword()); + + PCollection read = + p.apply( + CouchbaseIO.read() + .withConnectionConfiguration(noBucket) + .withQuery( + String.format( + "SELECT d.* FROM `%s`.`%s`.`%s` d WHERE d.test = $test", + BUCKET, SCOPE, COLLECTION)) + .withNamedParameter("test", "cluster") + .withScanConsistency(QueryScanConsistency.REQUEST_PLUS)); + + PAssert.that(read).containsInAnyOrder(docs); + p.run().waitUntilFinish(); + } + + @Test + public void testReadWithRedistribute() { + seed(seedBucket.defaultCollection(), "redistribute", 20); + + PCollection count = + p.apply( + readTagged("_default") + .withNamedParameter("test", "redistribute") + .withParseFn(row -> row.getInt("n").toString()) + .withCoder(StringUtf8Coder.of()) + .withRedistribute() + .withRedistributeNumKeys(4)) + .apply(Count.globally()); + + PAssert.thatSingleton(count).isEqualTo(20L); + p.run().waitUntilFinish(); + } + + @Test + public void testWriteThenReadRoundTrip() { + List docs = + IntStream.range(0, 10).mapToObj(n -> doc("roundtrip", n)).collect(Collectors.toList()); + writePipeline + .apply(Create.of(docs).withCoder(JsonObjectCoder.of())) + .apply( + CouchbaseIO.write() + .withConnectionConfiguration(config()) + .withIdFn(d -> "roundtrip-" + d.getInt("n")) + .withContentFn(d -> d)); + writePipeline.run().waitUntilFinish(); + + PCollection read = + p.apply(readTagged("_default").withNamedParameter("test", "roundtrip")); + + PAssert.that(read).containsInAnyOrder(docs); + p.run().waitUntilFinish(); + } + + @Test + public void testParseFnFailureFailsPipeline() { + seed(seedBucket.defaultCollection(), "parse-failure", 3); + p.apply( + readTagged("_default") + .withNamedParameter("test", "parse-failure") + .withParseFn( + row -> { + throw new IllegalStateException("cannot parse row " + row.getInt("n")); + }) + .withCoder(StringUtf8Coder.of())); + + RuntimeException e = assertThrows(RuntimeException.class, () -> p.run().waitUntilFinish()); + String chain = + Throwables.getCausalChain(e).stream() + .map(String::valueOf) + .collect(Collectors.joining("\n")); + assertThat(chain, containsString("cannot parse row")); + } + + /** The documents also carry {@code test}, which the schema doesn't name. */ + @Test + public void testReadRows() { + seed(seedBucket.defaultCollection(), "rows", 3); + Schema schema = + Schema.builder() + .addInt32Field("n") + .addNullableField("missing", Schema.FieldType.STRING) + .build(); + + PCollection rows = + p.apply( + CouchbaseIO.readRows(schema) + .withConnectionConfiguration(config()) + .withQuery("SELECT d.* FROM `_default` d WHERE d.test = $test") + .withNamedParameter("test", "rows") + .withScanConsistency(QueryScanConsistency.REQUEST_PLUS)); + + PAssert.that(rows) + .containsInAnyOrder( + IntStream.range(0, 3) + .mapToObj(n -> Row.withSchema(schema).addValues(n, null).build()) + .collect(Collectors.toList())); + p.run().waitUntilFinish(); + } + + @Test + public void testReadRowsTypeMismatchFailsPipeline() { + seed(seedBucket.defaultCollection(), "rows-mismatch", 1); + Schema schema = Schema.builder().addBooleanField("n").build(); + p.apply( + CouchbaseIO.readRows(schema) + .withConnectionConfiguration(config()) + .withQuery("SELECT d.* FROM `_default` d WHERE d.test = $test") + .withNamedParameter("test", "rows-mismatch") + .withScanConsistency(QueryScanConsistency.REQUEST_PLUS)); + + RuntimeException e = assertThrows(RuntimeException.class, () -> p.run().waitUntilFinish()); + String chain = + Throwables.getCausalChain(e).stream() + .map(String::valueOf) + .collect(Collectors.joining("\n")); + assertThat(chain, containsString("'n'")); + } +} diff --git a/sdks/java/io/couchbase/src/test/java/org/apache/beam/sdk/io/couchbase/CouchbaseIOTest.java b/sdks/java/io/couchbase/src/test/java/org/apache/beam/sdk/io/couchbase/CouchbaseIOTest.java new file mode 100644 index 000000000000..8fcdfb9d734b --- /dev/null +++ b/sdks/java/io/couchbase/src/test/java/org/apache/beam/sdk/io/couchbase/CouchbaseIOTest.java @@ -0,0 +1,440 @@ +/* + * 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.couchbase; + +import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.not; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import com.couchbase.client.java.json.JsonObject; +import com.couchbase.client.java.query.QueryScanConsistency; +import java.util.ArrayList; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.io.couchbase.CouchbaseIO.ConnectionConfiguration; +import org.apache.beam.sdk.runners.TransformHierarchy; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.display.DisplayData; +import org.apache.beam.sdk.util.RowJson.UnsupportedRowJsonException; +import org.apache.beam.sdk.util.SerializableUtils; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link CouchbaseIO} configuration and validation. No Couchbase server needed. */ +@RunWith(JUnit4.class) +public class CouchbaseIOTest { + + private static final String PASSWORD = "s3cr3t-do-not-log"; + + // These tests only apply transforms to check validation; none of the pipelines are run. + @Rule + public final transient TestPipeline p = + TestPipeline.create().enableAbandonedNodeEnforcement(false); + + private static ConnectionConfiguration fullConfig() { + return ConnectionConfiguration.create("couchbase://localhost") + .withUsername("user") + .withPassword(PASSWORD) + .withBucket("beam"); + } + + private static CouchbaseIO.Write fullWrite() { + return CouchbaseIO.write() + .withConnectionConfiguration(fullConfig()) + .withIdFn((String s) -> s) + .withContentFn((String s) -> JsonObject.create().put("v", s)); + } + + private IllegalArgumentException applyAndExpectFailure(CouchbaseIO.Write write) { + return assertThrows(IllegalArgumentException.class, () -> p.apply(Create.of("a")).apply(write)); + } + + @Test + public void testConnectionConfigurationDefaults() { + ConnectionConfiguration config = ConnectionConfiguration.create("couchbase://localhost"); + assertEquals("_default", config.scope()); + assertEquals("_default", config.collection()); + } + + @Test + public void testWriteDefaults() { + CouchbaseIO.Write write = CouchbaseIO.write(); + assertEquals(CouchbaseIO.DEFAULT_BATCH_SIZE, write.batchSize()); + assertEquals(CouchbaseIO.DEFAULT_MAX_CONCURRENCY, write.maxConcurrency()); + } + + @Test + public void testToStringOmitsPassword() { + String s = fullConfig().toString(); + assertThat(s, containsString("beam")); + assertThat(s, not(containsString(PASSWORD))); + } + + @Test + public void testDisplayDataOmitsPassword() { + DisplayData displayData = DisplayData.from(fullWrite().withBatchSize(100)); + assertThat(displayData, hasDisplayItem("connectionString", "couchbase://localhost")); + assertThat(displayData, hasDisplayItem("username", "user")); + assertThat(displayData, hasDisplayItem("bucket", "beam")); + assertThat(displayData, hasDisplayItem("scope", "_default")); + assertThat(displayData, hasDisplayItem("collection", "_default")); + assertThat(displayData, hasDisplayItem("batchSize", 100)); + assertThat(displayData, not(hasDisplayItem("password"))); + assertThat(displayData.toString(), not(containsString(PASSWORD))); + } + + @Test + public void testValidWriteApplies() { + p.apply(Create.of("a")).apply(fullWrite()); + } + + @Test + public void testMissingConnectionConfiguration() { + CouchbaseIO.Write write = + CouchbaseIO.write() + .withIdFn((String s) -> s) + .withContentFn((String s) -> JsonObject.create()); + assertThat( + applyAndExpectFailure(write).getMessage(), + containsString("withConnectionConfiguration() is required")); + } + + @Test + public void testMissingIdFn() { + CouchbaseIO.Write write = + CouchbaseIO.write() + .withConnectionConfiguration(fullConfig()) + .withContentFn((String s) -> JsonObject.create()); + assertThat(applyAndExpectFailure(write).getMessage(), containsString("withIdFn() is required")); + } + + @Test + public void testMissingContentFn() { + CouchbaseIO.Write write = + CouchbaseIO.write() + .withConnectionConfiguration(fullConfig()) + .withIdFn((String s) -> s); + assertThat( + applyAndExpectFailure(write).getMessage(), containsString("withContentFn() is required")); + } + + @Test + public void testMissingBucket() { + ConnectionConfiguration config = + ConnectionConfiguration.create("couchbase://localhost") + .withUsername("user") + .withPassword(PASSWORD); + assertThat( + applyAndExpectFailure(fullWrite().withConnectionConfiguration(config)).getMessage(), + containsString("withBucket() is required")); + } + + @Test + public void testMissingUsername() { + ConnectionConfiguration config = + ConnectionConfiguration.create("couchbase://localhost") + .withPassword(PASSWORD) + .withBucket("beam"); + assertThat( + applyAndExpectFailure(fullWrite().withConnectionConfiguration(config)).getMessage(), + containsString("withUsername() is required")); + } + + @Test + public void testMissingPassword() { + ConnectionConfiguration config = + ConnectionConfiguration.create("couchbase://localhost") + .withUsername("user") + .withBucket("beam"); + assertThat( + applyAndExpectFailure(fullWrite().withConnectionConfiguration(config)).getMessage(), + containsString("withPassword() is required")); + } + + @Test + public void testNonPositiveBatchSizeRejected() { + assertThrows(IllegalArgumentException.class, () -> CouchbaseIO.write().withBatchSize(0)); + } + + @Test + public void testNonPositiveMaxConcurrencyRejected() { + assertThrows(IllegalArgumentException.class, () -> CouchbaseIO.write().withMaxConcurrency(0)); + } + + private static CouchbaseIO.Read fullRead() { + return CouchbaseIO.read() + .withConnectionConfiguration(fullConfig()) + .withQuery("SELECT a.* FROM airline a"); + } + + private IllegalArgumentException applyAndExpectFailure(CouchbaseIO.Read read) { + return assertThrows(IllegalArgumentException.class, () -> p.apply(read)); + } + + @Test + public void testReadDefaults() { + CouchbaseIO.Read read = CouchbaseIO.read(); + assertEquals(QueryScanConsistency.NOT_BOUNDED, read.scanConsistency()); + assertEquals(JsonObjectCoder.of(), read.coder()); + assertFalse(read.redistribute()); + } + + @Test + public void testValidReadApplies() { + PCollection docs = p.apply(fullRead()); + assertEquals(JsonObjectCoder.of(), docs.getCoder()); + } + + @Test + public void testReadMissingConnectionConfiguration() { + assertThat( + applyAndExpectFailure(CouchbaseIO.read().withQuery("SELECT 1")).getMessage(), + containsString("withConnectionConfiguration() is required")); + } + + @Test + public void testReadMissingQuery() { + assertThat( + applyAndExpectFailure(CouchbaseIO.read().withConnectionConfiguration(fullConfig())) + .getMessage(), + containsString("withQuery() is required")); + } + + @Test + public void testReadMissingPassword() { + ConnectionConfiguration config = + ConnectionConfiguration.create("couchbase://localhost").withUsername("user"); + assertThat( + applyAndExpectFailure(fullRead().withConnectionConfiguration(config)).getMessage(), + containsString("withPassword() is required")); + } + + @Test + public void testReadWithoutBucketApplies() { + ConnectionConfiguration config = + ConnectionConfiguration.create("couchbase://localhost") + .withUsername("user") + .withPassword(PASSWORD); + p.apply(fullRead().withConnectionConfiguration(config)); + } + + @Test + public void testNamedAndPositionalParametersRejected() { + CouchbaseIO.Read read = + fullRead().withNamedParameter("country", "UK").withPositionalParameters("UK"); + assertThat(applyAndExpectFailure(read).getMessage(), containsString("cannot be combined")); + } + + @Test + public void testWithNamedParameterAccumulatesWithoutMutatingEarlierTransform() { + CouchbaseIO.Read one = fullRead().withNamedParameter("country", "UK"); + CouchbaseIO.Read two = one.withNamedParameter("limit", 10); + assertEquals(JsonObject.create().put("country", "UK"), one.namedParameters()); + assertEquals(JsonObject.create().put("country", "UK").put("limit", 10), two.namedParameters()); + } + + @Test + public void testWithNamedParameterRejectsUnsupportedType() { + assertThrows(RuntimeException.class, () -> fullRead().withNamedParameter("when", new Object())); + } + + @Test + public void testRedistributeNumKeysRequiresRedistribute() { + assertThat( + applyAndExpectFailure(fullRead().withRedistributeNumKeys(4)).getMessage(), + containsString("requires withRedistribute()")); + } + + @Test + public void testNonPositiveRedistributeNumKeysRejected() { + assertThrows( + IllegalArgumentException.class, () -> CouchbaseIO.read().withRedistributeNumKeys(0)); + } + + @Test + public void testWithParseFnLambdaRequiresCoder() { + CouchbaseIO.Read read = fullRead().withParseFn(row -> row.getString("name")); + assertThat(applyAndExpectFailure(read).getMessage(), containsString("withCoder()")); + } + + @Test + public void testWithParseFnResetsCoder() { + CouchbaseIO.Read read = + fullRead().withCoder(JsonObjectCoder.of()).withParseFn(row -> row.getString("name")); + assertNull(read.coder()); + } + + @Test + public void testWithParseFnAndCoderApplies() { + PCollection names = + p.apply( + fullRead().withParseFn(row -> row.getString("name")).withCoder(StringUtf8Coder.of())); + assertEquals(StringUtf8Coder.of(), names.getCoder()); + } + + /** The shuffle must carry the parsed type, so Redistribute has to come after the query DoFn. */ + @Test + public void testRedistributeAppliedAfterParsing() { + PCollection names = + p.apply( + fullRead() + .withParseFn(row -> row.getString("name")) + .withCoder(StringUtf8Coder.of()) + .withRedistribute() + .withRedistributeNumKeys(4)); + + List transforms = new ArrayList<>(); + p.traverseTopologically( + new Pipeline.PipelineVisitor.Defaults() { + @Override + public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) { + transforms.add(node.getFullName()); + return CompositeBehavior.ENTER_TRANSFORM; + } + + @Override + public void visitPrimitiveTransform(TransformHierarchy.Node node) { + transforms.add(node.getFullName()); + } + }); + assertThat( + transforms.indexOf("CouchbaseIO.Read/Redistribute"), + greaterThan(transforms.indexOf("CouchbaseIO.Read/Query"))); + assertTrue(transforms.contains("CouchbaseIO.Read/Query")); + assertEquals(StringUtf8Coder.of(), names.getCoder()); + } + + @Test + public void testReadDisplayDataOmitsPasswordAndParameters() { + DisplayData displayData = + DisplayData.from( + fullRead() + .withNamedParameter("customer", "secret-customer-id") + .withScanConsistency(QueryScanConsistency.REQUEST_PLUS) + .withRedistribute()); + assertThat(displayData, hasDisplayItem("query", "SELECT a.* FROM airline a")); + assertThat(displayData, hasDisplayItem("scanConsistency", "REQUEST_PLUS")); + assertThat(displayData, hasDisplayItem("redistribute", true)); + assertThat(displayData, hasDisplayItem("bucket", "beam")); + assertThat(displayData, not(hasDisplayItem("collection"))); + assertThat(displayData.toString(), not(containsString(PASSWORD))); + assertThat(displayData.toString(), not(containsString("secret-customer-id"))); + } + + private static final Schema HUB_SCHEMA = Schema.builder().addStringField("code").build(); + + private static final Schema AIRLINE_SCHEMA = + Schema.builder() + .addStringField("name") + .addNullableField("rating", Schema.FieldType.DOUBLE) + .addRowField("hub", HUB_SCHEMA) + .build(); + + private static String causalChain(Throwable t) { + return Throwables.getCausalChain(t).stream() + .map(String::valueOf) + .collect(Collectors.joining("\n")); + } + + @Test + public void testReadRowsSetsSchema() { + PCollection rows = + p.apply( + CouchbaseIO.readRows(AIRLINE_SCHEMA) + .withConnectionConfiguration(fullConfig()) + .withQuery("SELECT a.* FROM airline a")); + assertTrue(rows.hasSchema()); + assertEquals(AIRLINE_SCHEMA, rows.getSchema()); + } + + @Test + public void testReadRowsRejectsUnsupportedSchema() { + Schema schema = Schema.builder().addByteArrayField("payload").build(); + assertThrows(UnsupportedRowJsonException.class, () -> CouchbaseIO.readRows(schema)); + } + + @Test + public void testJsonToRowIgnoresExtraFields() { + JsonObject json = + JsonObject.create() + .put("name", "A") + .put("rating", 4.5) + .put("hub", JsonObject.create().put("code", "LHR").put("terminal", 5)) + .put("country", "UK"); + Row expected = + Row.withSchema(AIRLINE_SCHEMA) + .addValues("A", 4.5, Row.withSchema(HUB_SCHEMA).addValue("LHR").build()) + .build(); + assertEquals(expected, new CouchbaseIO.JsonToRowFn(AIRLINE_SCHEMA).apply(json)); + } + + @Test + public void testJsonToRowMissingNullableFieldIsNull() { + JsonObject json = + JsonObject.create().put("name", "A").put("hub", JsonObject.create().put("code", "LHR")); + assertNull(new CouchbaseIO.JsonToRowFn(AIRLINE_SCHEMA).apply(json).getValue("rating")); + } + + @Test + public void testJsonToRowMissingRequiredFieldFails() { + JsonObject json = JsonObject.create().put("hub", JsonObject.create().put("code", "LHR")); + RuntimeException e = + assertThrows( + RuntimeException.class, () -> new CouchbaseIO.JsonToRowFn(AIRLINE_SCHEMA).apply(json)); + assertThat(causalChain(e), containsString("'name'")); + } + + @Test + public void testJsonToRowTypeMismatchFails() { + JsonObject json = + JsonObject.create() + .put("name", "A") + .put("rating", "high") + .put("hub", JsonObject.create().put("code", "LHR")); + RuntimeException e = + assertThrows( + RuntimeException.class, () -> new CouchbaseIO.JsonToRowFn(AIRLINE_SCHEMA).apply(json)); + assertThat(causalChain(e), containsString("'rating'")); + } + + /** The mapper is transient, so a deserialized function has to build its own. */ + @Test + public void testJsonToRowWorksAfterSerialization() { + CouchbaseIO.JsonToRowFn fn = + SerializableUtils.ensureSerializable(new CouchbaseIO.JsonToRowFn(AIRLINE_SCHEMA)); + JsonObject json = + JsonObject.create().put("name", "A").put("hub", JsonObject.create().put("code", "LHR")); + assertEquals("A", fn.apply(json).getString("name")); + } +} diff --git a/sdks/java/io/couchbase/src/test/java/org/apache/beam/sdk/io/couchbase/CouchbaseIOWriteTest.java b/sdks/java/io/couchbase/src/test/java/org/apache/beam/sdk/io/couchbase/CouchbaseIOWriteTest.java new file mode 100644 index 000000000000..bf3906fbc5b7 --- /dev/null +++ b/sdks/java/io/couchbase/src/test/java/org/apache/beam/sdk/io/couchbase/CouchbaseIOWriteTest.java @@ -0,0 +1,186 @@ +/* + * 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.couchbase; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; +import static org.hamcrest.Matchers.not; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +import com.couchbase.client.java.Bucket; +import com.couchbase.client.java.Cluster; +import com.couchbase.client.java.Collection; +import com.couchbase.client.java.json.JsonObject; +import java.time.Duration; +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import org.apache.beam.sdk.io.couchbase.CouchbaseIO.ConnectionConfiguration; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.testcontainers.couchbase.BucketDefinition; +import org.testcontainers.couchbase.CouchbaseContainer; + +/** Writes to a single-node Couchbase Server started by Testcontainers. */ +@RunWith(JUnit4.class) +public class CouchbaseIOWriteTest { + + private static final String BUCKET = "beam"; + private static final String SCOPE = "app"; + private static final String COLLECTION = "airlines"; + + @ClassRule + public static final CouchbaseContainer COUCHBASE = + new CouchbaseContainer("couchbase/server:enterprise-8.0.3") + .withBucket(new BucketDefinition(BUCKET).withPrimaryIndex(false)); + + // Independent of the pipeline's connection, so results are checked through a separate client. + private static Cluster verifyCluster; + private static Bucket verifyBucket; + + @Rule public final transient TestPipeline p = TestPipeline.create(); + + // Only for the idempotency test, which runs the same write a second time. + @Rule public final transient TestPipeline rerun = TestPipeline.create(); + + @BeforeClass + public static void setUpClass() { + verifyCluster = + Cluster.connect( + COUCHBASE.getConnectionString(), COUCHBASE.getUsername(), COUCHBASE.getPassword()); + verifyBucket = verifyCluster.bucket(BUCKET); + verifyBucket.waitUntilReady(Duration.ofSeconds(60)); + verifyBucket.collections().createScope(SCOPE); + verifyBucket.collections().createCollection(SCOPE, COLLECTION); + } + + @AfterClass + public static void tearDownClass() { + if (verifyCluster != null) { + verifyCluster.disconnect(); + } + } + + private static ConnectionConfiguration config() { + return ConnectionConfiguration.create(COUCHBASE.getConnectionString()) + .withUsername(COUCHBASE.getUsername()) + .withPassword(COUCHBASE.getPassword()) + .withBucket(BUCKET); + } + + /** Each test uses its own ID prefix because all tests share one bucket. */ + private static List ids(String prefix, int n) { + return IntStream.range(0, n).mapToObj(i -> prefix + "-" + i).collect(Collectors.toList()); + } + + private static JsonObject content(String id) { + return JsonObject.create().put("id", id).put("source", "beam"); + } + + private static CouchbaseIO.Write write(ConnectionConfiguration config) { + return CouchbaseIO.write() + .withConnectionConfiguration(config) + .withIdFn((String id) -> id) + .withContentFn(CouchbaseIOWriteTest::content); + } + + private static void assertWritten(Collection collection, List ids) { + for (String id : ids) { + assertEquals(content(id), collection.get(id).contentAsObject()); + } + } + + @Test + public void testWrite() { + List ids = ids("basic", 10); + p.apply(Create.of(ids)).apply(write(config())); + p.run().waitUntilFinish(); + + assertWritten(verifyBucket.defaultCollection(), ids); + } + + @Test + public void testRerunningWriteIsIdempotent() { + List ids = ids("rerun", 10); + p.apply(Create.of(ids)).apply(write(config())); + p.run().waitUntilFinish(); + rerun.apply(Create.of(ids)).apply(write(config())); + rerun.run().waitUntilFinish(); + + assertWritten(verifyBucket.defaultCollection(), ids); + } + + /** Simulates a retried bundle: some documents already exist from a partially applied attempt. */ + @Test + public void testWriteOverwritesExistingDocuments() { + List ids = ids("retry", 10); + Collection collection = verifyBucket.defaultCollection(); + for (String id : ids.subList(0, 5)) { + collection.upsert(id, JsonObject.create().put("stale", true)); + } + + p.apply(Create.of(ids)).apply(write(config())); + p.run().waitUntilFinish(); + + assertWritten(collection, ids); + } + + @Test + public void testWriteToNonDefaultCollection() { + List ids = ids("collection", 10); + p.apply(Create.of(ids)).apply(write(config().withScope(SCOPE).withCollection(COLLECTION))); + p.run().waitUntilFinish(); + + assertWritten(verifyBucket.scope(SCOPE).collection(COLLECTION), ids); + } + + /** 10 documents with batch size 3 exercises both flush-when-full and flush in FinishBundle. */ + @Test + public void testWriteWithSmallBatchSize() { + List ids = ids("batch", 10); + p.apply(Create.of(ids)).apply(write(config()).withBatchSize(3).withMaxConcurrency(2)); + p.run().waitUntilFinish(); + + assertWritten(verifyBucket.defaultCollection(), ids); + } + + @Test + public void testWrongPasswordFailsWithoutLeakingIt() { + String wrongPassword = "wrong-s3cr3t-do-not-log"; + p.apply(Create.of(ids("auth", 1))).apply(write(config().withPassword(wrongPassword))); + + // The wrapper type depends on the runner and on which DoFn method failed, so assert on the + // chain. + RuntimeException e = assertThrows(RuntimeException.class, () -> p.run().waitUntilFinish()); + String chain = + Throwables.getCausalChain(e).stream() + .map(String::valueOf) + .collect(Collectors.joining("\n")); + assertThat(chain, containsString("authentication failed")); + assertThat(chain, not(containsString(wrongPassword))); + } +} diff --git a/sdks/java/io/couchbase/src/test/java/org/apache/beam/sdk/io/couchbase/JsonObjectCoderTest.java b/sdks/java/io/couchbase/src/test/java/org/apache/beam/sdk/io/couchbase/JsonObjectCoderTest.java new file mode 100644 index 000000000000..816f50cf15ba --- /dev/null +++ b/sdks/java/io/couchbase/src/test/java/org/apache/beam/sdk/io/couchbase/JsonObjectCoderTest.java @@ -0,0 +1,60 @@ +/* + * 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.couchbase; + +import static org.junit.Assert.assertThrows; + +import com.couchbase.client.java.json.JsonArray; +import com.couchbase.client.java.json.JsonObject; +import org.apache.beam.sdk.coders.Coder.NonDeterministicException; +import org.apache.beam.sdk.testing.CoderProperties; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class JsonObjectCoderTest { + + private static final JsonObjectCoder CODER = JsonObjectCoder.of(); + + @Test + public void testDecodeEncodeEqual() throws Exception { + CoderProperties.coderDecodeEncodeEqual(CODER, JsonObject.create()); + CoderProperties.coderDecodeEncodeEqual( + CODER, + JsonObject.create() + .put("name", "Air Ünïcode ✈") + .put("id", 10) + .put("passengers", 3_000_000_000L) + .put("rating", 4.5) + .put("active", true) + .putNull("callsign") + .put("routes", JsonArray.from("LHR", "JFK")) + .put("hub", JsonObject.create().put("code", "LHR").put("terminals", 5))); + } + + @Test + public void testCoderSerializable() { + CoderProperties.coderSerializable(CODER); + } + + @Test + public void testNotDeterministic() { + assertThrows(NonDeterministicException.class, CODER::verifyDeterministic); + } +} diff --git a/settings.gradle.kts b/settings.gradle.kts index b45c3f9ffa4c..99d89acb6c42 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -238,6 +238,7 @@ include(":sdks:java:io:clickhouse") include(":sdks:java:io:common") include(":sdks:java:io:components") include(":sdks:java:io:contextualtextio") +include(":sdks:java:io:couchbase") include(":sdks:java:io:debezium") include(":sdks:java:io:debezium:expansion-service") include(":sdks:java:io:elasticsearch") diff --git a/website/www/site/content/en/documentation/io/connectors.md b/website/www/site/content/en/documentation/io/connectors.md index b54bdd2edc05..9205936c2e24 100644 --- a/website/www/site/content/en/documentation/io/connectors.md +++ b/website/www/site/content/en/documentation/io/connectors.md @@ -856,6 +856,22 @@ This table provides a consolidated, at-a-glance overview of the available built- ✘ ✘ + + CouchbaseIO + ✔ + ✔ + + ✔ + native + + Not available + Not available + Not available + Not available + ✔ + ✘ + ✘ + RedisIO ✔