getExecutionTimeMonitoringData(ShortIdMap shortId
for (SimpleExecutionState state : executionStates) {
if (state.getTotalMillis() != 0) {
String shortId = state.getTotalMillisShortId(shortIds);
- if (result.containsKey(shortId)) {
- // This can happen due to flatten unzipping.
- result.put(shortId, state.mergeTotalMillisPayload(result.get(shortId)));
- } else {
- result.put(shortId, state.getTotalMillisPayload());
- }
+ result.compute(
+ shortId,
+ (String k, @Nullable ByteString existing) -> {
+ if (existing != null) {
+ // This can happen due to flatten unzipping.
+ return state.mergeTotalMillisPayload(existing);
+ } else {
+ return state.getTotalMillisPayload();
+ }
+ });
}
}
return result;
diff --git a/runners/flink/flink_runner.gradle b/runners/flink/flink_runner.gradle
index 2b6392f46b9bc..bc0cc188f0632 100644
--- a/runners/flink/flink_runner.gradle
+++ b/runners/flink/flink_runner.gradle
@@ -290,7 +290,7 @@ tasks.register('validatesRunner') {
// Generates :runners:flink:1.13:runQuickstartJavaFlinkLocal
createJavaExamplesArchetypeValidationTask(type: 'Quickstart', runner: 'FlinkLocal')
-task examplesIntegrationTest(type: Test) {
+tasks.register("examplesIntegrationTest", Test) {
group = "Verification"
// Disable gradle cache
outputs.upToDateWhen { false }
diff --git a/runners/flink/job-server/flink_job_server.gradle b/runners/flink/job-server/flink_job_server.gradle
index 981a2cd870b16..2b3bbdfceb337 100644
--- a/runners/flink/job-server/flink_job_server.gradle
+++ b/runners/flink/job-server/flink_job_server.gradle
@@ -212,7 +212,7 @@ project.ext.validatesPortableRunnerBatch = portableValidatesRunnerTask("Batch",
project.ext.validatesPortableRunnerStreaming = portableValidatesRunnerTask("Streaming", true, false, false)
project.ext.validatesPortableRunnerStreamingCheckpoint = portableValidatesRunnerTask("StreamingCheckpointing", true, true, false)
-task validatesPortableRunner() {
+tasks.register("validatesPortableRunner") {
dependsOn validatesPortableRunnerDocker
dependsOn validatesPortableRunnerBatch
dependsOn validatesPortableRunnerStreaming
@@ -265,7 +265,8 @@ createCrossLanguageValidatesRunnerTask(
)
// miniCluster jar starts an embedded Flink cluster intended for use in testing.
-task miniCluster(type: Jar, dependsOn: shadowJar) {
+tasks.register("miniCluster", Jar) {
+ dependsOn shadowJar
archiveBaseName = "${project.archivesBaseName}-mini-cluster"
dependencies {
runtimeOnly project(path: flinkRunnerProject, configuration: "miniCluster")
diff --git a/runners/google-cloud-dataflow-java/build.gradle b/runners/google-cloud-dataflow-java/build.gradle
index 89d185bf53512..355fc4547dfdc 100644
--- a/runners/google-cloud-dataflow-java/build.gradle
+++ b/runners/google-cloud-dataflow-java/build.gradle
@@ -367,7 +367,13 @@ task printRunnerV2PipelineOptions {
task validatesRunner {
group = "Verification"
description "Validates Dataflow runner"
- dependsOn(createLegacyWorkerValidatesRunnerTest(name: 'validatesRunnerLegacyWorkerTest'))
+ dependsOn(createLegacyWorkerValidatesRunnerTest(
+ name: 'validatesRunnerLegacyWorkerTest',
+ excludedTests: [
+ // TODO(BEAM-13952)
+ 'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testAfterProcessingTimeContinuationTriggerUsingState',
+ ]
+ ))
}
task validatesRunnerStreaming {
@@ -382,6 +388,10 @@ task validatesRunnerStreaming {
'org.apache.beam.sdk.testing.UsesRequiresTimeSortedInput',
'org.apache.beam.sdk.testing.UsesSetState',
],
+ excludedTests: [
+ // TODO(BEAM-13952)
+ 'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testAfterProcessingTimeContinuationTriggerUsingState'
+ ]
))
}
@@ -473,6 +483,9 @@ task validatesRunnerV2 {
'org.apache.beam.sdk.transforms.GroupByKeyTest$WindowTests.testRewindowWithTimestampCombiner',
'org.apache.beam.sdk.transforms.FlattenTest.testFlattenWithDifferentInputAndOutputCoders2',
+
+ // TODO(BEAM-13952)
+ 'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testAfterProcessingTimeContinuationTriggerUsingState',
]
))
}
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineDebugOptions.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineDebugOptions.java
index 264827f33aa69..2995a485527cd 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineDebugOptions.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineDebugOptions.java
@@ -177,6 +177,30 @@ public Dataflow create(PipelineOptions options) {
void setNumberOfWorkerHarnessThreads(int value);
+ /**
+ * Maximum number of bundles outstanding from windmill before the worker stops requesting.
+ *
+ * If <= 0, use the default value of 100 + getNumberOfWorkerHarnessThreads()
+ */
+ @Description(
+ "Maximum number of bundles outstanding from windmill before the worker stops requesting.")
+ @Default.Integer(0)
+ int getMaxBundlesFromWindmillOutstanding();
+
+ void setMaxBundlesFromWindmillOutstanding(int value);
+
+ /**
+ * Maximum number of bytes outstanding from windmill before the worker stops requesting.
+ *
+ *
If <= 0, use the default value of 50% of jvm memory.
+ */
+ @Description(
+ "Maximum number of bytes outstanding from windmill before the worker stops requesting. If <= 0, use the default value of 50% of jvm memory.")
+ @Default.Long(0)
+ long getMaxBytesFromWindmillOutstanding();
+
+ void setMaxBytesFromWindmillOutstanding(long value);
+
/**
* If {@literal true}, save a heap dump before killing a thread or process which is GC thrashing
* or out of memory. The location of the heap file will either be echoed back to the user, or the
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
index 236e47715a367..c49865aa7511c 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
@@ -191,10 +191,6 @@ public class StreamingDataflowWorker {
// Maximum number of threads for processing. Currently each thread processes one key at a time.
static final int MAX_PROCESSING_THREADS = 300;
static final long THREAD_EXPIRATION_TIME_SEC = 60;
- // Maximum work units retrieved from Windmill and queued before processing. Limiting this delays
- // retrieving extra work from Windmill without working on it, leading to better
- // prioritization / utilization.
- static final int MAX_WORK_UNITS_QUEUED = 100;
static final long TARGET_COMMIT_BUNDLE_BYTES = 32 << 20;
static final int MAX_COMMIT_QUEUE_BYTES = 500 << 20; // 500MB
static final int NUM_COMMIT_STREAMS = 1;
@@ -210,9 +206,6 @@ public class StreamingDataflowWorker {
// Matches kWindmillCounterUpdate in workflow_worker_service_multi_hubs.cc.
private static final String WINDMILL_COUNTER_UPDATE_WORK_ID = "3";
- /** Maximum number of items to return in a GetWork request. */
- private static final long MAX_GET_WORK_ITEMS = MAX_WORK_UNITS_QUEUED + MAX_PROCESSING_THREADS;
-
/** Maximum number of failure stacktraces to report in each update sent to backend. */
private static final int MAX_FAILURES_TO_REPORT_IN_UPDATE = 1000;
@@ -666,7 +659,8 @@ public static StreamingDataflowWorker fromDataflowWorkerHarnessOptions(
chooseMaximumNumberOfThreads(),
THREAD_EXPIRATION_TIME_SEC,
TimeUnit.SECONDS,
- MAX_WORK_UNITS_QUEUED,
+ chooseMaximumBundlesOutstanding(),
+ chooseMaximumBytesOutstanding(),
threadFactory);
maxSinkBytes =
@@ -785,6 +779,22 @@ private int chooseMaximumNumberOfThreads() {
return MAX_PROCESSING_THREADS;
}
+ private int chooseMaximumBundlesOutstanding() {
+ int maxBundles = options.getMaxBundlesFromWindmillOutstanding();
+ if (maxBundles > 0) {
+ return maxBundles;
+ }
+ return chooseMaximumNumberOfThreads() + 100;
+ }
+
+ private long chooseMaximumBytesOutstanding() {
+ long maxMem = options.getMaxBytesFromWindmillOutstanding();
+ if (maxMem > 0) {
+ return maxMem;
+ }
+ return Runtime.getRuntime().maxMemory() / 2;
+ }
+
void addStateNameMappings(Map nameMap) {
stateNameMap.putAll(nameMap);
}
@@ -804,7 +814,7 @@ public void setMaxWorkItemCommitBytes(int maxWorkItemCommitBytes) {
@VisibleForTesting
public boolean workExecutorIsEmpty() {
- return workUnitExecutor.getQueue().isEmpty();
+ return workUnitExecutor.executorQueueIsEmpty();
}
public void start() {
@@ -949,9 +959,6 @@ public void stop() {
memoryMonitor.stop();
memoryMonitorThread.join();
workUnitExecutor.shutdown();
- if (!workUnitExecutor.awaitTermination(5, TimeUnit.MINUTES)) {
- throw new RuntimeException("Work executor did not terminate within 5 minutes");
- }
for (ComputationState state : computationMap.values()) {
state.close();
}
@@ -1065,7 +1072,7 @@ void streamingDispatchLoop() {
windmillServer.getWorkStream(
Windmill.GetWorkRequest.newBuilder()
.setClientId(clientId)
- .setMaxItems(MAX_GET_WORK_ITEMS)
+ .setMaxItems(chooseMaximumBundlesOutstanding())
.setMaxBytes(MAX_GET_WORK_FETCH_BYTES)
.build(),
(String computation,
@@ -1236,7 +1243,8 @@ private void callFinalizeCallbacks(Windmill.WorkItem work) {
} catch (Throwable t) {
LOG.error("Source checkpoint finalization failed:", t);
}
- });
+ },
+ 0);
}
}
}
@@ -1548,7 +1556,7 @@ private void process(
if (retryLocally) {
// Try again after some delay and at the end of the queue to avoid a tight loop.
sleep(retryLocallyDelayMs);
- workUnitExecutor.forceExecute(work);
+ workUnitExecutor.forceExecute(work, work.getWorkItem().getSerializedSize());
} else {
// Consider the item invalid. It will eventually be retried by Windmill if it still needs to
// be processed.
@@ -1726,7 +1734,7 @@ private Windmill.GetWorkResponse getWork() {
return windmillServer.getWork(
Windmill.GetWorkRequest.newBuilder()
.setClientId(clientId)
- .setMaxItems(MAX_GET_WORK_ITEMS)
+ .setMaxItems(chooseMaximumBundlesOutstanding())
.setMaxBytes(MAX_GET_WORK_FETCH_BYTES)
.build());
}
@@ -2285,7 +2293,7 @@ public boolean activateWork(ShardedKey shardedKey, Work work) {
// Fall through to execute without the lock held.
}
}
- executor.execute(work);
+ executor.execute(work, work.getWorkItem().getSerializedSize());
return true;
}
@@ -2327,7 +2335,7 @@ public void completeWork(ShardedKey shardedKey, long workToken) {
}
}
if (nextWork != null) {
- executor.forceExecute(nextWork);
+ executor.forceExecute(nextWork, nextWork.getWorkItem().getSerializedSize());
}
}
@@ -2511,27 +2519,9 @@ public void doGet(HttpServletRequest request, HttpServletResponse response) thro
}
private class MetricsDataProvider implements StatusDataProvider {
-
@Override
public void appendSummaryHtml(PrintWriter writer) {
- writer.println(
- "Worker Threads: "
- + workUnitExecutor.getPoolSize()
- + "/"
- + workUnitExecutor.getMaximumPoolSize()
- + "(),
- threadFactory);
- this.semaphore = new ReducableSemaphore(maximumQueueSize);
- allowCoreThreadTimeOut(true);
+ executor =
+ new ThreadPoolExecutor(
+ maximumPoolSize,
+ maximumPoolSize,
+ keepAliveTime,
+ unit,
+ new LinkedBlockingQueue<>(),
+ threadFactory);
+ executor.allowCoreThreadTimeOut(true);
+ this.maximumElementsOutstanding = maximumElementsOutstanding;
+ this.maximumBytesOutstanding = maximumBytesOutstanding;
}
- // Before adding a Runnable to the queue, acquire the semaphore.
- @Override
- public void execute(Runnable r) {
- semaphore.acquireUninterruptibly();
- super.execute(r);
+ // Before adding a Work to the queue, check that there are enough bytes of space or no other
+ // outstanding elements of work.
+ public void execute(Runnable work, long workBytes) {
+ monitor.enterWhenUninterruptibly(
+ new Guard(monitor) {
+ @Override
+ public boolean isSatisfied() {
+ return elementsOutstanding == 0
+ || (bytesAvailable() >= workBytes
+ && elementsOutstanding < maximumElementsOutstanding);
+ }
+ });
+ executeLockHeld(work, workBytes);
}
// Forcibly add something to the queue, ignoring the length limit.
- public void forceExecute(Runnable r) {
- semaphore.reducePermits(1);
- super.execute(r);
+ public void forceExecute(Runnable work, long workBytes) {
+ monitor.enter();
+ executeLockHeld(work, workBytes);
+ }
+
+ public void shutdown() throws InterruptedException {
+ executor.shutdown();
+ if (!executor.awaitTermination(5, TimeUnit.MINUTES)) {
+ throw new RuntimeException("Work executor did not terminate within 5 minutes");
+ }
+ }
+
+ public boolean executorQueueIsEmpty() {
+ return executor.getQueue().isEmpty();
+ }
+
+ public String summaryHtml() {
+ monitor.enter();
+ try {
+ StringBuilder builder = new StringBuilder();
+ builder.append("Worker Threads: ");
+ builder.append(executor.getPoolSize());
+ builder.append("/");
+ builder.append(executor.getMaximumPoolSize());
+ builder.append(" {
+ let checkStatus: CheckStatus = {
+ completed: true,
+ succeeded: true,
+ };
+ const mostRecentChecks = await getMostRecentChecks(owner, repo, checkSha);
+ for (let i = 0; i < mostRecentChecks.length; i++) {
+ if (mostRecentChecks[i].status != "completed") {
+ checkStatus.completed = false;
+ }
+ if (mostRecentChecks[i].conclusion != "success") {
+ checkStatus.succeeded = false;
+ }
+ }
+
+ return checkStatus;
+}
+
+async function getMostRecentChecks(
+ owner: string,
+ repo: string,
+ checkSha: string
+): Promise {
+ let mostRecentChecks: any[] = [];
+ const checksByName = await getChecksByName(owner, repo, checkSha);
+
+ const checkNames = Object.keys(checksByName);
+ for (let i = 0; i < checkNames.length; i++) {
+ let checks = checksByName[checkNames[i]];
+ let mostRecent = checks.sort((a, b) =>
+ a.completionTime > b.completionTime ? 1 : -1
+ )[0];
+ mostRecentChecks.push(mostRecent);
+ }
+
+ return mostRecentChecks;
+}
+
+async function getChecksByName(
+ owner: string,
+ repo: string,
+ checkSha: string
+): Promise {
+ const githubClient = getGitHubClient();
+ const allChecks = (
+ await githubClient.rest.checks.listForRef({
+ owner: owner,
+ repo: repo,
+ ref: checkSha,
+ })
+ ).data.check_runs;
+ let checksByName = {};
+ allChecks.forEach((checkRun) => {
+ if (!shouldExcludeCheck(checkRun)) {
+ let name = checkRun.name;
+ let check = {
+ status: checkRun.status,
+ conclusion: checkRun.conclusion,
+ completionTime: checkRun.completed_at,
+ };
+ if (!checksByName[name]) {
+ checksByName[name] = [check];
+ } else {
+ checksByName[name].push(check);
+ }
+ }
+ });
+
+ return checksByName;
+}
+
+// Returns checks we should exclude because they are flaky or not always predictive of pr mergability.
+// Currently just excludes codecov.
+function shouldExcludeCheck(check): boolean {
+ if (check.name.toLowerCase().indexOf("codecov") != -1) {
+ return true;
+ }
+ return false;
+}
diff --git a/scripts/ci/pr-bot/shared/commentStrings.ts b/scripts/ci/pr-bot/shared/commentStrings.ts
new file mode 100644
index 0000000000000..e70a721c95fa7
--- /dev/null
+++ b/scripts/ci/pr-bot/shared/commentStrings.ts
@@ -0,0 +1,72 @@
+/*
+ * 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.
+ */
+
+export function allChecksPassed(reviewersToNotify: string[]): string {
+ return `All checks have passed: @${reviewersToNotify.join(" ")}`;
+}
+
+export function assignCommitter(committer: string): string {
+ return `R: @${committer} for final approval`;
+}
+
+export function assignReviewer(labelToReviewerMapping: any): string {
+ let commentString =
+ "Assigning reviewers. If you would like to opt out of this review, comment `assign to next reviewer`:\n\n";
+
+ for (let label in labelToReviewerMapping) {
+ let reviewer = labelToReviewerMapping[label];
+ if (label === "no-matching-label") {
+ commentString += `R: @${reviewer} added as fallback since no labels match configuration\n`;
+ } else {
+ commentString += `R: @${reviewer} for label ${label}.\n`;
+ }
+ }
+
+ commentString += `
+Available commands:
+- \`stop reviewer notifications\` - opt out of the automated review tooling
+- \`remind me after tests pass\` - tag the comment author after tests pass
+- \`waiting on author\` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)`;
+ return commentString;
+}
+
+export function failingChecksCantAssign(): string {
+ return "Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment `assign set of reviewers`";
+}
+
+export function someChecksFailing(reviewersToNotify: string[]): string {
+ return `Some checks have failed: @${reviewersToNotify.join(" ")}`;
+}
+
+export function stopNotifications(reason: string): string {
+ return `Stopping reviewer notifications for this pull request: ${reason}`;
+}
+
+export function remindReviewerAfterTestsPass(requester: string): string {
+ return `Ok - I'll remind @${requester} after tests pass`;
+}
+
+export function reviewersAlreadyAssigned(reviewers: string[]): string {
+ return `Reviewers are already assigned to this PR: ${reviewers
+ .map((reviewer) => "@" + reviewer)
+ .join(" ")}`;
+}
+
+export function noLegalReviewers(): string {
+ return "No reviewers could be found from any of the labels on the PR or in the fallback reviewers list. Check the config file to make sure reviewers are configured";
+}
diff --git a/scripts/ci/pr-bot/shared/constants.ts b/scripts/ci/pr-bot/shared/constants.ts
new file mode 100644
index 0000000000000..859bf8589e623
--- /dev/null
+++ b/scripts/ci/pr-bot/shared/constants.ts
@@ -0,0 +1,30 @@
+/*
+ * 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.
+ */
+
+const path = require("path");
+
+export const REPO_OWNER = "apache";
+export const REPO = "beam";
+export const PATH_TO_CONFIG_FILE = path.join(
+ __dirname,
+ "../../../../../.github/REVIEWERS.yml"
+);
+export const PATH_TO_METRICS_CSV = path.resolve(
+ path.join(__dirname, "../../metrics.csv")
+);
+export const BOT_NAME = "github-actions";
diff --git a/scripts/ci/pr-bot/shared/githubUtils.ts b/scripts/ci/pr-bot/shared/githubUtils.ts
new file mode 100644
index 0000000000000..e650ca52b0d09
--- /dev/null
+++ b/scripts/ci/pr-bot/shared/githubUtils.ts
@@ -0,0 +1,96 @@
+/*
+ * 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.
+ */
+
+const { Octokit } = require("@octokit/rest");
+const { REPO_OWNER, REPO } = require("./constants");
+
+export interface Label {
+ name: string;
+}
+
+export function getGitHubClient() {
+ let auth = process.env["GITHUB_TOKEN"];
+ if (!auth) {
+ throw new Error(
+ "No github token provided - process.env['GITHUB_TOKEN'] must be set."
+ );
+ }
+ return new Octokit({ auth });
+}
+
+export async function addPrComment(pullNumber: number, body: string) {
+ await getGitHubClient().rest.issues.createComment({
+ owner: REPO_OWNER,
+ repo: REPO,
+ issue_number: pullNumber,
+ body,
+ });
+}
+
+export async function nextActionReviewers(
+ pullNumber: number,
+ existingLabels: Label[]
+) {
+ let newLabels = removeNextActionLabel(existingLabels);
+ newLabels.push("Next Action: Reviewers");
+ await getGitHubClient().rest.issues.setLabels({
+ owner: REPO_OWNER,
+ repo: REPO,
+ issue_number: pullNumber,
+ labels: newLabels,
+ });
+}
+
+export async function nextActionAuthor(
+ pullNumber: number,
+ existingLabels: Label[]
+) {
+ let newLabels = removeNextActionLabel(existingLabels);
+ newLabels.push("Next Action: Author");
+ await getGitHubClient().rest.issues.setLabels({
+ owner: REPO_OWNER,
+ repo: REPO,
+ issue_number: pullNumber,
+ labels: newLabels,
+ });
+}
+
+export async function checkIfCommitter(username: string): Promise {
+ const permissionLevel = (
+ await getGitHubClient().rest.repos.getCollaboratorPermissionLevel({
+ owner: REPO_OWNER,
+ repo: REPO,
+ username,
+ })
+ ).data;
+
+ return (
+ permissionLevel.permission === "write" ||
+ permissionLevel.permission === "admin"
+ );
+}
+
+function removeNextActionLabel(existingLabels: Label[]): string[] {
+ return existingLabels
+ .filter(
+ (label) =>
+ label.name != "Next Action: Reviewers" &&
+ label.name != "Next Action: Author"
+ )
+ .map((label) => label.name);
+}
diff --git a/scripts/ci/pr-bot/shared/persistentState.ts b/scripts/ci/pr-bot/shared/persistentState.ts
new file mode 100644
index 0000000000000..f7fade16f9b93
--- /dev/null
+++ b/scripts/ci/pr-bot/shared/persistentState.ts
@@ -0,0 +1,121 @@
+/*
+ * 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.
+ */
+
+const exec = require("@actions/exec");
+const fs = require("fs");
+const path = require("path");
+const { Pr } = require("./pr");
+const { ReviewersForLabel } = require("./reviewersForLabel");
+const { BOT_NAME } = require("./constants");
+
+function getPrFileName(prNumber) {
+ return `pr-${prNumber}.json`.toLowerCase();
+}
+
+function getReviewersForLabelFileName(label) {
+ return `reviewers-for-label-${label}.json`.toLowerCase();
+}
+
+async function commitStateToRepo() {
+ await exec.exec("git pull origin pr-bot-state");
+ await exec.exec("git add state/*");
+ await exec.exec(`git commit -m "Updating config from bot" --allow-empty`);
+ await exec.exec("git push origin pr-bot-state");
+}
+
+export class PersistentState {
+ private switchedBranch = false;
+
+ // Returns a Pr object representing the current saved state of the pr.
+ async getPrState(prNumber: number): Promise {
+ var fileName = getPrFileName(prNumber);
+ return new Pr(await this.getState(fileName, "state/pr-state"));
+ }
+
+ // Writes a Pr object representing the current saved state of the pr to persistent storage.
+ async writePrState(prNumber: number, newState: any) {
+ var fileName = getPrFileName(prNumber);
+ await this.writeState(fileName, "state/pr-state", new Pr(newState));
+ }
+
+ // Returns a ReviewersForLabel object representing the current saved state of which reviewers have reviewed recently.
+ async getReviewersForLabelState(
+ label: string
+ ): Promise {
+ var fileName = getReviewersForLabelFileName(label);
+ return new ReviewersForLabel(label, await this.getState(fileName, "state"));
+ }
+
+ // Writes a ReviewersForLabel object representing the current saved state of which reviewers have reviewed recently.
+ async writeReviewersForLabelState(label: string, newState: any) {
+ var fileName = getReviewersForLabelFileName(label);
+ await this.writeState(
+ fileName,
+ "state",
+ new ReviewersForLabel(label, newState)
+ );
+ }
+
+ private async getState(fileName, baseDirectory) {
+ await this.ensureCorrectBranch();
+ fileName = path.join(baseDirectory, fileName);
+ if (!fs.existsSync(fileName)) {
+ return null;
+ }
+ return JSON.parse(fs.readFileSync(fileName, { encoding: "utf-8" }));
+ }
+
+ private async writeState(fileName, baseDirectory, state) {
+ await this.ensureCorrectBranch();
+ fileName = path.join(baseDirectory, fileName);
+ if (!fs.existsSync(baseDirectory)) {
+ fs.mkdirSync(baseDirectory, { recursive: true });
+ }
+ fs.writeFileSync(fileName, JSON.stringify(state, null, 2), {
+ encoding: "utf-8",
+ });
+ await commitStateToRepo();
+ }
+
+ private async ensureCorrectBranch() {
+ if (this.switchedBranch) {
+ return;
+ }
+ console.log(
+ "Switching to branch pr-bot-state for reading/storing persistent state between runs"
+ );
+ try {
+ await exec.exec(`git config user.name ${BOT_NAME}`);
+ await exec.exec(`git config user.email ${BOT_NAME}@github.com`);
+ await exec.exec("git config pull.rebase false");
+ await exec.exec("git fetch origin pr-bot-state");
+ await exec.exec("git checkout pr-bot-state");
+ } catch {
+ console.log(
+ "Couldnt find branch pr-bot-state in origin, trying to create it"
+ );
+ try {
+ await exec.exec("git checkout -b pr-bot-state");
+ } catch {
+ console.log("Creating branch failed, trying a simple checkout.");
+ await exec.exec("git checkout pr-bot-state");
+ }
+ }
+ this.switchedBranch = true;
+ }
+}
diff --git a/scripts/ci/pr-bot/shared/pr.ts b/scripts/ci/pr-bot/shared/pr.ts
new file mode 100644
index 0000000000000..c1419b8419cee
--- /dev/null
+++ b/scripts/ci/pr-bot/shared/pr.ts
@@ -0,0 +1,89 @@
+/*
+ * 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.
+ */
+
+const github = require("./githubUtils");
+
+export class Pr {
+ public commentedAboutFailingChecks: boolean;
+ public reviewersAssignedForLabels: { [key: string]: string };
+ public nextAction: string;
+ public stopReviewerNotifications: boolean;
+ public remindAfterTestsPass: string[];
+
+ constructor(propertyDictionary) {
+ this.commentedAboutFailingChecks = false;
+ this.reviewersAssignedForLabels = {}; // map of label to reviewer
+ this.nextAction = "Author";
+ this.stopReviewerNotifications = false;
+ this.remindAfterTestsPass = []; // List of handles
+
+ if (!propertyDictionary) {
+ return;
+ }
+ if (propertyDictionary) {
+ if ("commentedAboutFailingChecks" in propertyDictionary) {
+ this.commentedAboutFailingChecks =
+ propertyDictionary["commentedAboutFailingChecks"];
+ }
+ if ("reviewersAssignedForLabels" in propertyDictionary) {
+ this.reviewersAssignedForLabels =
+ propertyDictionary["reviewersAssignedForLabels"];
+ }
+ if ("nextAction" in propertyDictionary) {
+ this.nextAction = propertyDictionary["nextAction"];
+ }
+ if ("stopReviewerNotifications" in propertyDictionary) {
+ this.stopReviewerNotifications =
+ propertyDictionary["stopReviewerNotifications"];
+ }
+ if ("remindAfterTestsPass" in propertyDictionary) {
+ this.remindAfterTestsPass = propertyDictionary["remindAfterTestsPass"];
+ }
+ }
+ }
+
+ // Returns a label that the reviewer is assigned for.
+ // If none, returns an empty string
+ getLabelForReviewer(reviewer: string): string {
+ const labels = Object.keys(this.reviewersAssignedForLabels);
+ for (let i = 0; i < labels.length; i++) {
+ let label = labels[i];
+ if (this.reviewersAssignedForLabels[label] === reviewer) {
+ return label;
+ }
+ }
+
+ return "";
+ }
+
+ // Returns whether any of the assigned reviewers are committers
+ async isAnyAssignedReviewerCommitter(): Promise {
+ const labels = Object.keys(this.reviewersAssignedForLabels);
+ for (let i = 0; i < labels.length; i++) {
+ if (
+ await github.checkIfCommitter(
+ this.reviewersAssignedForLabels[labels[i]]
+ )
+ ) {
+ return true;
+ }
+ }
+
+ return false;
+ }
+}
diff --git a/scripts/ci/pr-bot/shared/reviewerConfig.ts b/scripts/ci/pr-bot/shared/reviewerConfig.ts
new file mode 100644
index 0000000000000..f8300adbdfaf3
--- /dev/null
+++ b/scripts/ci/pr-bot/shared/reviewerConfig.ts
@@ -0,0 +1,97 @@
+/*
+ * 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.
+ */
+
+const yaml = require("js-yaml");
+const fs = require("fs");
+import { Label } from "./githubUtils";
+
+export class ReviewerConfig {
+ private config: any;
+ constructor(pathToConfigFile) {
+ this.config = yaml.load(
+ fs.readFileSync(pathToConfigFile, { encoding: "utf-8" })
+ );
+ }
+
+ // Given a list of labels and an exclusion list of reviewers not to include (e.g. the author)
+ // returns all possible reviewers for each label
+ getReviewersForLabels(
+ labels: Label[],
+ exclusionList: string[]
+ ): { [key: string]: string[] } {
+ let reviewersFound = false;
+ let labelToReviewerMapping = {};
+ labels.forEach((label) => {
+ let reviewers = this.getReviewersForLabel(label.name, exclusionList);
+ if (reviewers.length > 0) {
+ labelToReviewerMapping[label.name] = reviewers;
+ reviewersFound = true;
+ }
+ });
+ if (!reviewersFound) {
+ const fallbackReviewers = this.getFallbackReviewers(exclusionList);
+ if (fallbackReviewers.length > 0) {
+ labelToReviewerMapping["no-matching-label"] =
+ this.getFallbackReviewers(exclusionList);
+ }
+ }
+ return labelToReviewerMapping;
+ }
+
+ // Get possible reviewers excluding the author.
+ getReviewersForLabel(label: string, exclusionList: string[]): string[] {
+ var labelObjects = this.config.labels;
+ const labelObject = labelObjects.find(
+ (labelObject) => labelObject.name.toLowerCase() === label.toLowerCase()
+ );
+ if (!labelObject) {
+ return [];
+ }
+
+ return this.excludeFromReviewers(labelObject.reviewers, exclusionList);
+ }
+
+ getExclusionListForLabel(label: string): string[] {
+ var labelObjects = this.config.labels;
+ const labelObject = labelObjects.find(
+ (labelObject) => labelObject.name.toLowerCase() === label.toLowerCase()
+ );
+ return labelObject?.exclusionList ?? [];
+ }
+
+ // Get fallback reviewers excluding the author.
+ getFallbackReviewers(exclusionList: string[]): string[] {
+ return this.excludeFromReviewers(
+ this.config.fallbackReviewers,
+ exclusionList
+ );
+ }
+
+ private excludeFromReviewers(
+ reviewers: string[],
+ exclusionList: string[]
+ ): string[] {
+ if (!exclusionList) {
+ return reviewers;
+ }
+
+ return reviewers.filter(
+ (reviewer) => exclusionList.indexOf(reviewer) == -1
+ );
+ }
+}
diff --git a/scripts/ci/pr-bot/shared/reviewersForLabel.ts b/scripts/ci/pr-bot/shared/reviewersForLabel.ts
new file mode 100644
index 0000000000000..971f3f1cd7a53
--- /dev/null
+++ b/scripts/ci/pr-bot/shared/reviewersForLabel.ts
@@ -0,0 +1,102 @@
+/*
+ * 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.
+ */
+
+const github = require("./githubUtils");
+
+export class ReviewersForLabel {
+ public label: string;
+ public dateOfLastReviewAssignment: { [key: string]: number };
+
+ constructor(
+ label: string,
+ propertyDictionary: {
+ dateOfLastReviewAssignment: { [key: string]: number };
+ }
+ ) {
+ this.label = label;
+ this.dateOfLastReviewAssignment = {}; // map of reviewer to date
+
+ if (!propertyDictionary) {
+ return;
+ }
+ if ("dateOfLastReviewAssignment" in propertyDictionary) {
+ this.dateOfLastReviewAssignment =
+ propertyDictionary["dateOfLastReviewAssignment"];
+ }
+ }
+
+ // Given a list of available reviewers,
+ // returns the next reviewer up based on who has reviewed least recently.
+ // Updates this object to reflect their assignment.
+ assignNextReviewer(availableReviewers: string[]): string {
+ if (availableReviewers.length === 0) {
+ throw new Error(`No reviewers available for label ${this.label}`);
+ }
+
+ if (!this.dateOfLastReviewAssignment[availableReviewers[0]]) {
+ this.dateOfLastReviewAssignment[availableReviewers[0]] = Date.now();
+ return availableReviewers[0];
+ }
+
+ let earliestDate = this.dateOfLastReviewAssignment[availableReviewers[0]];
+ let earliestReviewer = availableReviewers[0];
+
+ for (let i = 0; i < availableReviewers.length; i++) {
+ let availableReviewer = availableReviewers[i];
+ if (!this.dateOfLastReviewAssignment[availableReviewer]) {
+ this.dateOfLastReviewAssignment[availableReviewer] = Date.now();
+ return availableReviewer;
+ }
+ if (earliestDate > this.dateOfLastReviewAssignment[availableReviewer]) {
+ earliestDate = this.dateOfLastReviewAssignment[availableReviewer];
+ earliestReviewer = availableReviewer;
+ }
+ }
+
+ this.dateOfLastReviewAssignment[earliestReviewer] = Date.now();
+ return earliestReviewer;
+ }
+
+ // Given the up to date list of available reviewers (excluding the author),
+ // returns the next reviewer up based on who has reviewed least recently.
+ // Updates this object to reflect their assignment.
+ async assignNextCommitter(availableReviewers: string[]): Promise {
+ let earliestDate = Date.now();
+ let earliestCommitter: string = "";
+
+ for (let i = 0; i < availableReviewers.length; i++) {
+ let availableReviewer = availableReviewers[i];
+ if (await github.checkIfCommitter(availableReviewer)) {
+ if (!this.dateOfLastReviewAssignment[availableReviewer]) {
+ this.dateOfLastReviewAssignment[availableReviewer] = Date.now();
+ return availableReviewer;
+ }
+ if (earliestDate > this.dateOfLastReviewAssignment[availableReviewer]) {
+ earliestDate = this.dateOfLastReviewAssignment[availableReviewer];
+ earliestCommitter = availableReviewer;
+ }
+ }
+ }
+
+ if (!earliestCommitter) {
+ throw new Error(`No committers available for label ${this.label}`);
+ }
+ this.dateOfLastReviewAssignment[earliestCommitter] = Date.now();
+ return earliestCommitter;
+ }
+}
diff --git a/scripts/ci/pr-bot/shared/userCommand.ts b/scripts/ci/pr-bot/shared/userCommand.ts
new file mode 100644
index 0000000000000..34e2741a3f6d4
--- /dev/null
+++ b/scripts/ci/pr-bot/shared/userCommand.ts
@@ -0,0 +1,231 @@
+/*
+ * 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.
+ */
+
+const github = require("./githubUtils");
+const commentStrings = require("./commentStrings");
+const { BOT_NAME } = require("./constants");
+const { StateClient } = require("./persistentState");
+const { ReviewerConfig } = require("./reviewerConfig");
+
+// Reads the comment and processes the command if one is contained in it.
+// Returns true if it runs a command, false otherwise.
+export async function processCommand(
+ payload: any,
+ commentAuthor: string,
+ commentText: string,
+ stateClient: typeof StateClient,
+ reviewerConfig: typeof ReviewerConfig
+) {
+ // Don't process any commands from our bot.
+ if (commentAuthor === BOT_NAME) {
+ return false;
+ }
+ console.log(commentAuthor);
+
+ const pullNumber = payload.issue?.number || payload.pull_request?.number;
+ commentText = commentText.toLowerCase();
+ if (commentText.indexOf("r: @") > -1) {
+ await manuallyAssignedToReviewer(pullNumber, stateClient);
+ } else if (commentText.indexOf("assign to next reviewer") > -1) {
+ await assignToNextReviewer(
+ payload,
+ commentAuthor,
+ pullNumber,
+ stateClient,
+ reviewerConfig
+ );
+ } else if (commentText.indexOf("stop reviewer notifications") > -1) {
+ await stopReviewerNotifications(
+ pullNumber,
+ stateClient,
+ "requested by reviewer"
+ );
+ } else if (commentText.indexOf("remind me after tests pass") > -1) {
+ await remindAfterTestsPass(pullNumber, commentAuthor, stateClient);
+ } else if (commentText.indexOf("waiting on author") > -1) {
+ await waitOnAuthor(payload, pullNumber, stateClient);
+ } else if (commentText.indexOf("assign set of reviewers") > -1) {
+ await assignReviewerSet(payload, pullNumber, stateClient, reviewerConfig);
+ } else {
+ return false;
+ }
+
+ return true;
+}
+
+async function assignToNextReviewer(
+ payload: any,
+ commentAuthor: string,
+ pullNumber: number,
+ stateClient: typeof StateClient,
+ reviewerConfig: typeof ReviewerConfig
+) {
+ let prState = await stateClient.getPrState(pullNumber);
+ let labelOfReviewer = prState.getLabelForReviewer(payload.sender.login);
+ if (labelOfReviewer) {
+ let reviewersState = await stateClient.getReviewersForLabelState(
+ labelOfReviewer
+ );
+ const pullAuthor =
+ payload.issue?.user?.login || payload.pull_request?.user?.login;
+ let availableReviewers = reviewerConfig.getReviewersForLabel(
+ labelOfReviewer,
+ [commentAuthor, pullAuthor]
+ );
+ let chosenReviewer = reviewersState.assignNextReviewer(availableReviewers);
+ prState.reviewersAssignedForLabels[labelOfReviewer] = chosenReviewer;
+
+ // Comment assigning reviewer
+ console.log(`Assigning ${chosenReviewer}`);
+ await github.addPrComment(
+ pullNumber,
+ commentStrings.assignReviewer(prState.reviewersAssignedForLabels)
+ );
+
+ // Set next action to reviewer
+ const existingLabels =
+ payload.issue?.labels || payload.pull_request?.labels;
+ await github.nextActionReviewers(pullNumber, existingLabels);
+ prState.nextAction = "Reviewers";
+
+ // Persist state
+ await stateClient.writePrState(pullNumber, prState);
+ await stateClient.writeReviewersForLabelState(
+ labelOfReviewer,
+ reviewersState
+ );
+ }
+}
+
+// If they've manually assigned a reviewer, just silence notifications and ignore this pr going forward.
+// TODO(damccorm) - we could try to do something more intelligent here like figuring out which label that reviewer belongs to.
+async function manuallyAssignedToReviewer(
+ pullNumber: number,
+ stateClient: typeof StateClient
+) {
+ await stopReviewerNotifications(
+ pullNumber,
+ stateClient,
+ "review requested by someone other than the bot, ceding control"
+ );
+}
+
+async function stopReviewerNotifications(
+ pullNumber: number,
+ stateClient: typeof StateClient,
+ reason: string
+) {
+ let prState = await stateClient.getPrState(pullNumber);
+ prState.stopReviewerNotifications = true;
+ await stateClient.writePrState(pullNumber, prState);
+
+ // Comment acknowledging command
+ await github.addPrComment(
+ pullNumber,
+ commentStrings.stopNotifications(reason)
+ );
+}
+
+async function remindAfterTestsPass(
+ pullNumber: number,
+ username: string,
+ stateClient: typeof StateClient
+) {
+ let prState = await stateClient.getPrState(pullNumber);
+ prState.remindAfterTestsPass.push(username);
+ await stateClient.writePrState(pullNumber, prState);
+
+ // Comment acknowledging command
+ await github.addPrComment(
+ pullNumber,
+ commentStrings.remindReviewerAfterTestsPass(username)
+ );
+}
+
+async function waitOnAuthor(
+ payload: any,
+ pullNumber: number,
+ stateClient: typeof StateClient
+) {
+ const existingLabels = payload.issue?.labels || payload.pull_request?.labels;
+ await github.nextActionAuthor(pullNumber, existingLabels);
+ let prState = await stateClient.getPrState(pullNumber);
+ prState.nextAction = "Author";
+ await stateClient.writePrState(pullNumber, prState);
+}
+
+async function assignReviewerSet(
+ payload: any,
+ pullNumber: number,
+ stateClient: typeof StateClient,
+ reviewerConfig: typeof ReviewerConfig
+) {
+ let prState = await stateClient.getPrState(pullNumber);
+ if (Object.values(prState.reviewersAssignedForLabels).length > 0) {
+ await github.addPrComment(
+ pullNumber,
+ commentStrings.reviewersAlreadyAssigned(
+ Object.values(prState.reviewersAssignedForLabels)
+ )
+ );
+ return;
+ }
+
+ const existingLabels = payload.issue?.labels || payload.pull_request?.labels;
+ const pullAuthor =
+ payload.issue?.user?.login || payload.pull_request?.user?.login;
+ const reviewersForLabels = reviewerConfig.getReviewersForLabels(
+ existingLabels,
+ [pullAuthor]
+ );
+ let reviewerStateToUpdate = {};
+ var labels = Object.keys(reviewersForLabels);
+ if (!labels || labels.length == 0) {
+ await github.addPrComment(
+ pullNumber,
+ commentStrings.noLegalReviewers(existingLabels)
+ );
+ return;
+ }
+ for (let i = 0; i < labels.length; i++) {
+ let label = labels[i];
+ let availableReviewers = reviewersForLabels[label];
+ let reviewersState = await stateClient.getReviewersForLabelState(label);
+ let chosenReviewer = reviewersState.assignNextReviewer(availableReviewers);
+ reviewerStateToUpdate[label] = reviewersState;
+ prState.reviewersAssignedForLabels[label] = chosenReviewer;
+ }
+ console.log(`Assigning reviewers for pr ${pullNumber}`);
+ await github.addPrComment(
+ pullNumber,
+ commentStrings.assignReviewer(prState.reviewersAssignedForLabels)
+ );
+
+ github.nextActionReviewers(pullNumber, existingLabels);
+ prState.nextAction = "Reviewers";
+
+ await stateClient.writePrState(pullNumber, prState);
+ let labelsToUpdate = Object.keys(reviewerStateToUpdate);
+ for (let i = 0; i < labelsToUpdate.length; i++) {
+ let label = labelsToUpdate[i];
+ await stateClient.writeReviewersForLabelState(
+ label,
+ reviewerStateToUpdate[label]
+ );
+ }
+}
diff --git a/scripts/ci/pr-bot/test/prTest.ts b/scripts/ci/pr-bot/test/prTest.ts
new file mode 100644
index 0000000000000..b771852584a30
--- /dev/null
+++ b/scripts/ci/pr-bot/test/prTest.ts
@@ -0,0 +1,45 @@
+/*
+ * 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.
+ */
+
+var assert = require("assert");
+const { Pr } = require("../shared/pr");
+describe("Pr", function () {
+ describe("getLabelForReviewer()", function () {
+ it("should return the label a reviewer is assigned to", function () {
+ let testPr = new Pr({});
+ testPr.reviewersAssignedForLabels = {
+ Go: "testReviewer1",
+ Java: "testReviewer2",
+ Python: "testReviewer3",
+ };
+ assert.equal("Go", testPr.getLabelForReviewer("testReviewer1"));
+ assert.equal("Java", testPr.getLabelForReviewer("testReviewer2"));
+ assert.equal("Python", testPr.getLabelForReviewer("testReviewer3"));
+ });
+
+ it("should return an empty string when a reviewer is not assigned", function () {
+ let testPr = new Pr({});
+ testPr.reviewersAssignedForLabels = {
+ Go: "testReviewer1",
+ Java: "testReviewer2",
+ Python: "testReviewer3",
+ };
+ assert.equal("", testPr.getLabelForReviewer("testReviewer4"));
+ });
+ });
+});
diff --git a/scripts/ci/pr-bot/test/reviewerConfigTest.ts b/scripts/ci/pr-bot/test/reviewerConfigTest.ts
new file mode 100644
index 0000000000000..589aab635bc95
--- /dev/null
+++ b/scripts/ci/pr-bot/test/reviewerConfigTest.ts
@@ -0,0 +1,238 @@
+/*
+ * 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.
+ */
+
+var assert = require("assert");
+var fs = require("fs");
+const { ReviewerConfig } = require("../shared/reviewerConfig");
+const configPath = "test-config.yml";
+const configContents = `labels:
+- name: "Go"
+ reviewers: ["testReviewer1", "testReviewer2"]
+ exclusionList: ["testReviewer3"] # These users will never be suggested as reviewers
+# I don't know the other areas well enough to assess who the normal committers/contributors who might want to be reviewers are
+- name: "Java"
+ reviewers: ["testReviewer3", "testReviewer2"]
+ exclusionList: [] # These users will never be suggested as reviewers
+- name: "Python"
+ reviewers: ["testReviewer4"]
+ exclusionList: [] # These users will never be suggested as reviewers
+fallbackReviewers: ["testReviewer5", "testReviewer1", "testReviewer3"] # List of committers to use when no label matches
+`;
+describe("ReviewerConfig", function () {
+ before(function () {
+ if (fs.existsSync(configPath)) {
+ fs.rmSync(configPath);
+ }
+ fs.writeFileSync(configPath, configContents);
+ });
+
+ after(function () {
+ fs.rmSync(configPath);
+ });
+
+ describe("getReviewersForLabels()", function () {
+ it("should return all reviewers configured for all labels", function () {
+ const config = new ReviewerConfig(configPath);
+ const reviewersForLabels = config.getReviewersForLabels(
+ [{ name: "Go" }, { name: "Java" }],
+ []
+ );
+ assert(
+ reviewersForLabels["Go"].find(
+ (reviewer) => reviewer === "testReviewer1"
+ ),
+ "Return value for Go label should include testReviewer1"
+ );
+ assert(
+ reviewersForLabels["Go"].find(
+ (reviewer) => reviewer === "testReviewer2"
+ ),
+ "Return value for Go label should include testReviewer2"
+ );
+ assert(
+ !reviewersForLabels["Go"].find(
+ (reviewer) => reviewer === "testReviewer3"
+ ),
+ "Return value for Go label should not include testReviewer3"
+ );
+ assert(
+ !reviewersForLabels["Go"].find(
+ (reviewer) => reviewer === "testReviewer4"
+ ),
+ "Return value for Go label should not include testReviewer4"
+ );
+
+ assert(
+ reviewersForLabels["Java"].find(
+ (reviewer) => reviewer === "testReviewer3"
+ ),
+ "Return value for Java label should include testReviewer3"
+ );
+ assert(
+ reviewersForLabels["Java"].find(
+ (reviewer) => reviewer === "testReviewer2"
+ ),
+ "Return value for Java label should include testReviewer2"
+ );
+ assert(
+ !reviewersForLabels["Java"].find(
+ (reviewer) => reviewer === "testReviewer4"
+ ),
+ "Return value for Java label should not include testReviewer4"
+ );
+ assert(
+ !reviewersForLabels["Java"].find(
+ (reviewer) => reviewer === "testReviewer1"
+ ),
+ "Return value for Java label should not include testReviewer1"
+ );
+ assert(
+ !reviewersForLabels["Java"].find(
+ (reviewer) => reviewer === "testReviewer5"
+ ),
+ "Return value for Java label should not include testReviewer5"
+ );
+
+ assert(
+ Object.keys(reviewersForLabels).indexOf("Python") == -1,
+ "No reviewers should be included for python"
+ );
+ });
+
+ it("should return no entry if a label is not configured", function () {
+ const config = new ReviewerConfig(configPath);
+ const reviewersForLabels = config.getReviewersForLabels(
+ [{ name: "FakeLabel" }],
+ []
+ );
+
+ assert(
+ !("FakeLabel" in reviewersForLabels),
+ "FakeLabel should not be included in the returned label map"
+ );
+ });
+
+ it("should exclude any reviewers who are passed into the exlusionList", function () {
+ const config = new ReviewerConfig(configPath);
+ const reviewersForLabels = config.getReviewersForLabels(
+ [{ name: "Go" }, { name: "Java" }],
+ ["testReviewer1"]
+ );
+
+ assert(
+ !reviewersForLabels["Go"].find(
+ (reviewer) => reviewer === "testReviewer1"
+ ),
+ "testReviewer1 should have been excluded from the result set"
+ );
+ });
+ });
+
+ describe("getReviewersForLabel()", function () {
+ it("should return all reviewers configured for a label", function () {
+ const config = new ReviewerConfig(configPath);
+ const reviewersForGo = config.getReviewersForLabel("Go", []);
+
+ assert(
+ reviewersForGo.find((reviewer) => reviewer === "testReviewer1"),
+ "Return value for Go label should include testReviewer1"
+ );
+ assert(
+ reviewersForGo.find((reviewer) => reviewer === "testReviewer2"),
+ "Return value for Go label should include testReviewer2"
+ );
+ assert(
+ !reviewersForGo.find((reviewer) => reviewer === "testReviewer3"),
+ "Return value for Go label should not include testReviewer3"
+ );
+ assert(
+ !reviewersForGo.find((reviewer) => reviewer === "testReviewer4"),
+ "Return value for Go label should not include testReviewer4"
+ );
+ });
+
+ it("should return an empty list if a label is not configured", function () {
+ const config = new ReviewerConfig(configPath);
+ const reviewersForFakeLabel = config.getReviewersForLabel(
+ "FakeLabel",
+ []
+ );
+ assert.equal(0, reviewersForFakeLabel.length);
+ });
+
+ it("should exclude any reviewers who are passed into the exlusionList", function () {
+ const config = new ReviewerConfig(configPath);
+ const reviewersForGo = config.getReviewersForLabel("Go", [
+ "testReviewer1",
+ ]);
+
+ assert(
+ !reviewersForGo.find((reviewer) => reviewer === "testReviewer1"),
+ "Return value for Go label should not include testReviewer1"
+ );
+ });
+ });
+
+ describe("getExclusionListForLabel()", function () {
+ it("should get the exclusion list configured for a label", function () {
+ const config = new ReviewerConfig(configPath);
+ const goExclusionList = config.getExclusionListForLabel("Go");
+
+ assert(
+ goExclusionList.find((reviewer) => reviewer === "testReviewer3"),
+ "Return value for Go label should include testReviewer3"
+ );
+ assert.equal(1, goExclusionList.length);
+ });
+ });
+
+ describe("getFallbackReviewers()", function () {
+ it("should get the configured falback list", function () {
+ const config = new ReviewerConfig(configPath);
+ const fallbackReviewers = config.getFallbackReviewers([]);
+
+ assert.equal(3, fallbackReviewers.length);
+ assert(
+ fallbackReviewers.find((reviewer) => reviewer === "testReviewer5"),
+ "Fallback reviewers should include testReviewer5"
+ );
+ assert(
+ fallbackReviewers.find((reviewer) => reviewer === "testReviewer1"),
+ "Fallback reviewers should include testReviewer1"
+ );
+ assert(
+ fallbackReviewers.find((reviewer) => reviewer === "testReviewer3"),
+ "Fallback reviewers should include testReviewer3"
+ );
+ });
+
+ it("should not include excluded reviewers", function () {
+ const config = new ReviewerConfig(configPath);
+ const fallbackReviewers = config.getFallbackReviewers([
+ "testReviewer1",
+ "testReviewer3",
+ ]);
+
+ assert.equal(1, fallbackReviewers.length);
+ assert(
+ fallbackReviewers.find((reviewer) => reviewer === "testReviewer5"),
+ "Fallback reviewers should only include testReviewer5"
+ );
+ });
+ });
+});
diff --git a/scripts/ci/pr-bot/test/reviewersForLabelTest.ts b/scripts/ci/pr-bot/test/reviewersForLabelTest.ts
new file mode 100644
index 0000000000000..f425750505927
--- /dev/null
+++ b/scripts/ci/pr-bot/test/reviewersForLabelTest.ts
@@ -0,0 +1,135 @@
+/*
+ * 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.
+ */
+
+var assert = require("assert");
+const { ReviewersForLabel } = require("../shared/reviewersForLabel");
+describe("ReviewersForLabel", function () {
+ describe("assignNextReviewer()", function () {
+ it("should repeatedly assign the reviewer who reviewed least recently", function () {
+ const dateOfLastReviewAssignment = {
+ testReviewer1: 1,
+ testReviewer2: 3,
+ testReviewer3: 4,
+ testReviewer4: 2,
+ };
+ let reviewersForGo = new ReviewersForLabel("Go", {
+ dateOfLastReviewAssignment: dateOfLastReviewAssignment,
+ });
+
+ assert.equal(
+ "testReviewer1",
+ reviewersForGo.assignNextReviewer([
+ "testReviewer1",
+ "testReviewer2",
+ "testReviewer3",
+ "testReviewer4",
+ ])
+ );
+ assert.equal(
+ "testReviewer4",
+ reviewersForGo.assignNextReviewer([
+ "testReviewer1",
+ "testReviewer2",
+ "testReviewer3",
+ "testReviewer4",
+ ])
+ );
+ assert.equal(
+ "testReviewer2",
+ reviewersForGo.assignNextReviewer([
+ "testReviewer1",
+ "testReviewer2",
+ "testReviewer3",
+ "testReviewer4",
+ ])
+ );
+ assert.equal(
+ "testReviewer3",
+ reviewersForGo.assignNextReviewer([
+ "testReviewer1",
+ "testReviewer2",
+ "testReviewer3",
+ "testReviewer4",
+ ])
+ );
+ assert.equal(
+ "testReviewer1",
+ reviewersForGo.assignNextReviewer([
+ "testReviewer1",
+ "testReviewer2",
+ "testReviewer3",
+ "testReviewer4",
+ ])
+ );
+ });
+
+ it("should assign a reviewer who hasnt reviewed before", function () {
+ const dateOfLastReviewAssignment = {
+ testReviewer1: 1,
+ testReviewer2: 2,
+ testReviewer3: 3,
+ testReviewer4: 4,
+ };
+ let reviewersForGo = new ReviewersForLabel("Go", {
+ dateOfLastReviewAssignment: dateOfLastReviewAssignment,
+ });
+
+ assert.equal(
+ "testReviewer5",
+ reviewersForGo.assignNextReviewer([
+ "testReviewer1",
+ "testReviewer2",
+ "testReviewer3",
+ "testReviewer4",
+ "testReviewer5",
+ ])
+ );
+ });
+
+ it("should only assign reviewers in the availableReviewers list", function () {
+ const dateOfLastReviewAssignment = {
+ testReviewer1: 1,
+ testReviewer2: 2,
+ testReviewer3: 3,
+ testReviewer4: 4,
+ };
+ let reviewersForGo = new ReviewersForLabel("Go", {
+ dateOfLastReviewAssignment: dateOfLastReviewAssignment,
+ });
+
+ assert.equal(
+ "testReviewer2",
+ reviewersForGo.assignNextReviewer(["testReviewer4", "testReviewer2"])
+ );
+ });
+
+ it("should throw if no reviewer available", function () {
+ const dateOfLastReviewAssignment = {
+ testReviewer1: 1,
+ testReviewer2: 2,
+ testReviewer3: 3,
+ testReviewer4: 4,
+ };
+ let reviewersForGo = new ReviewersForLabel("Go", {
+ dateOfLastReviewAssignment: dateOfLastReviewAssignment,
+ });
+
+ assert.throws(() => reviewersForGo.assignNextReviewer([]));
+ });
+ });
+});
diff --git a/scripts/ci/pr-bot/tsconfig.json b/scripts/ci/pr-bot/tsconfig.json
new file mode 100644
index 0000000000000..719f5dff8b9f5
--- /dev/null
+++ b/scripts/ci/pr-bot/tsconfig.json
@@ -0,0 +1,12 @@
+{
+ "compilerOptions": {
+ "target": "es6" /* Specify ECMAScript target version: 'ES3' (default), 'ES5', 'ES2015', 'ES2016', 'ES2017', 'ES2018', 'ES2019' or 'ESNEXT'. */,
+ "module": "commonjs" /* Specify module code generation: 'none', 'commonjs', 'amd', 'system', 'umd', 'es2015', or 'ESNext'. */,
+ "outDir": "./lib" /* Redirect output structure to the directory. */,
+ "rootDir": "./" /* Specify the root directory of input files. Use to control the output directory structure with --outDir. */,
+ "strict": true /* Enable all strict type-checking options. */,
+ "noImplicitAny": false /* Raise error on expressions and declarations with an implied 'any' type. */,
+ "esModuleInterop": true /* Enables emit interoperability between CommonJS and ES Modules via creation of namespace objects for all imports. Implies 'allowSyntheticDefaultImports'. */
+ },
+ "exclude": ["node_modules"]
+}
diff --git a/sdks/go/pkg/beam/core/graph/mtime/time.go b/sdks/go/pkg/beam/core/graph/mtime/time.go
index 25ea9ce567f72..bdbb0c441af80 100644
--- a/sdks/go/pkg/beam/core/graph/mtime/time.go
+++ b/sdks/go/pkg/beam/core/graph/mtime/time.go
@@ -37,7 +37,7 @@ const (
// EndOfGlobalWindowTime is the timestamp at the end of the global window. It
// is a day before the max timestamp.
- // TODO Use GLOBAL_WINDOW_MAX_TIMESTAMP_MILLIS from the Runner API constants
+ // TODO(BEAM-4179) Use GLOBAL_WINDOW_MAX_TIMESTAMP_MILLIS from the Runner API constants
EndOfGlobalWindowTime = MaxTimestamp - 24*60*60*1000
// ZeroTimestamp is the default zero value time. It corresponds to the unix epoch.
@@ -65,7 +65,8 @@ func FromDuration(d time.Duration) Time {
// FromTime returns a milli-second precision timestamp from a time.Time.
func FromTime(t time.Time) Time {
- return Normalize(Time(n2m(t.UnixNano())))
+ // TODO(BEAM-13988): Replace t.UnixNano() with t.UnixMilli() for Go 1.17 or higher.
+ return Normalize(Time(t.UnixNano() / 1e6))
}
// Milliseconds returns the number of milli-seconds since the Unix epoch.
@@ -73,14 +74,18 @@ func (t Time) Milliseconds() int64 {
return int64(t)
}
-// Add returns the time plus the duration.
+// Add returns the time plus the duration. Input Durations of less than one
+// millisecond will not increment the time due to a loss of precision when
+// converting to milliseconds.
func (t Time) Add(d time.Duration) Time {
- return Normalize(Time(int64(t) + n2m(d.Nanoseconds())))
+ return Normalize(Time(int64(t) + d.Milliseconds()))
}
-// Subtract returns the time minus the duration.
+// Subtract returns the time minus the duration. Input Durations of less than one
+// millisecond will not increment the time due to a loss of precision when
+// converting to milliseconds.
func (t Time) Subtract(d time.Duration) Time {
- return Normalize(Time(int64(t) - n2m(d.Nanoseconds())))
+ return Normalize(Time(int64(t) - d.Milliseconds()))
}
func (t Time) String() string {
@@ -116,8 +121,3 @@ func Max(a, b Time) Time {
func Normalize(t Time) Time {
return Min(Max(t, MinTimestamp), MaxTimestamp)
}
-
-// n2m converts nanoseconds to milliseconds.
-func n2m(v int64) int64 {
- return v / 1e6
-}
diff --git a/sdks/go/pkg/beam/core/graph/mtime/time_test.go b/sdks/go/pkg/beam/core/graph/mtime/time_test.go
new file mode 100644
index 0000000000000..7012184c18ce7
--- /dev/null
+++ b/sdks/go/pkg/beam/core/graph/mtime/time_test.go
@@ -0,0 +1,215 @@
+// 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 mtime
+
+import (
+ "math"
+ "testing"
+ "time"
+)
+
+func TestAdd(t *testing.T) {
+ tests := []struct {
+ name string
+ baseTime Time
+ addition time.Duration
+ expOut Time
+ }{
+ {
+ "insignificant addition small",
+ Time(1000),
+ 1 * time.Nanosecond,
+ Time(1000),
+ },
+ {
+ "insignificant addition large",
+ Time(1000),
+ 999999 * time.Nanosecond,
+ Time(1000),
+ },
+ {
+ "significant addition small",
+ Time(1000),
+ 1 * time.Millisecond,
+ Time(1001),
+ },
+ {
+ "significant addition large",
+ Time(1000),
+ 10 * time.Second,
+ Time(11000),
+ },
+ {
+ "add past max timestamp",
+ MaxTimestamp,
+ 1 * time.Minute,
+ MaxTimestamp,
+ },
+ {
+ "add across max boundary",
+ Time(int64(MaxTimestamp) - 10000),
+ 10 * time.Minute,
+ MaxTimestamp,
+ },
+ }
+ for _, test := range tests {
+ t.Run(test.name, func(t *testing.T) {
+ if got, want := test.baseTime.Add(test.addition), test.expOut; got != want {
+ t.Errorf("(%v).Add(%v), got time %v, want %v", test.baseTime, test.addition, got, want)
+ }
+ })
+ }
+}
+
+func TestSubtract(t *testing.T) {
+ tests := []struct {
+ name string
+ baseTime Time
+ subtraction time.Duration
+ expOut Time
+ }{
+ {
+ "insignificant subtraction small",
+ Time(1000),
+ 1 * time.Nanosecond,
+ Time(1000),
+ },
+ {
+ "insignificant subtraction large",
+ Time(1000),
+ 999999 * time.Nanosecond,
+ Time(1000),
+ },
+ {
+ "significant subtraction small",
+ Time(1000),
+ 1 * time.Millisecond,
+ Time(999),
+ },
+ {
+ "significant subtraction large",
+ Time(1000),
+ 10 * time.Second,
+ Time(-9000),
+ },
+ {
+ "subtract past min timestamp",
+ MinTimestamp,
+ 1 * time.Minute,
+ MinTimestamp,
+ },
+ {
+ "subtract across min boundary",
+ Time(int64(MinTimestamp) + 10000),
+ 10 * time.Minute,
+ MinTimestamp,
+ },
+ }
+ for _, test := range tests {
+ t.Run(test.name, func(t *testing.T) {
+ if got, want := test.baseTime.Subtract(test.subtraction), test.expOut; got != want {
+ t.Errorf("(%v).Subtract(%v), got time %v, want %v", test.baseTime, test.subtraction, got, want)
+ }
+ })
+ }
+}
+
+func TestNormalize(t *testing.T) {
+ tests := []struct {
+ name string
+ in Time
+ expOut Time
+ }{
+ {
+ "min timestamp",
+ MinTimestamp,
+ MinTimestamp,
+ },
+ {
+ "max timestamp",
+ MaxTimestamp,
+ MaxTimestamp,
+ },
+ {
+ "end of global window",
+ EndOfGlobalWindowTime,
+ EndOfGlobalWindowTime,
+ },
+ {
+ "beyond max timestamp",
+ Time(math.MaxInt64),
+ MaxTimestamp,
+ },
+ {
+ "below min timestamp",
+ Time(math.MinInt64),
+ MinTimestamp,
+ },
+ {
+ "normal value",
+ Time(int64(20000)),
+ Time(int64(20000)),
+ },
+ }
+ for _, test := range tests {
+ t.Run(test.name, func(t *testing.T) {
+ if got, want := Normalize(test.in), test.expOut; got != want {
+ t.Errorf("Normalize(%v), got Time %v, want %v", test.in, got, want)
+ }
+ })
+ }
+}
+
+func TestFromTime(t *testing.T) {
+ tests := []struct {
+ name string
+ input time.Time
+ expOut Time
+ }{
+ {
+ "zero unix",
+ time.Unix(0, 0).UTC(),
+ Time(0),
+ },
+ {
+ "behind unix",
+ time.Unix(-1, 0).UTC(),
+ Time(-1000),
+ },
+ {
+ "ahead of unix",
+ time.Unix(1, 0).UTC(),
+ Time(1000),
+ },
+ {
+ "insignificant time small",
+ time.Unix(0, 1),
+ Time(0),
+ },
+ {
+ "insignificant time large",
+ time.Unix(0, 999999),
+ Time(0),
+ },
+ }
+ for _, test := range tests {
+ t.Run(test.name, func(t *testing.T) {
+ if got, want := FromTime(test.input), test.expOut; got != want {
+ t.Errorf("FromTime(%v), got %v, want %v", test.input, got, want)
+ }
+ })
+ }
+}
diff --git a/sdks/go/pkg/beam/core/graph/window/trigger/trigger.go b/sdks/go/pkg/beam/core/graph/window/trigger/trigger.go
index c12c03aa5e893..148a672fc695b 100644
--- a/sdks/go/pkg/beam/core/graph/window/trigger/trigger.go
+++ b/sdks/go/pkg/beam/core/graph/window/trigger/trigger.go
@@ -66,6 +66,9 @@ func (t *AfterCountTrigger) ElementCount() int32 {
// AfterCount constructs a trigger that fires after
// at least `count` number of elements are processed.
func AfterCount(count int32) *AfterCountTrigger {
+ if count < 1 {
+ panic(fmt.Errorf("trigger.AfterCount(%v) must be a positive integer", count))
+ }
return &AfterCountTrigger{elementCount: count}
}
diff --git a/sdks/go/pkg/beam/core/graph/window/trigger/trigger_test.go b/sdks/go/pkg/beam/core/graph/window/trigger/trigger_test.go
new file mode 100644
index 0000000000000..a23de354bfd61
--- /dev/null
+++ b/sdks/go/pkg/beam/core/graph/window/trigger/trigger_test.go
@@ -0,0 +1,73 @@
+// 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 trigger
+
+import (
+ "testing"
+ "time"
+
+ "github.com/google/go-cmp/cmp"
+)
+
+func TestAfterCountTrigger(t *testing.T) {
+ tr := AfterCount(1)
+ want := int32(1)
+ if got := tr.ElementCount(); got != want {
+ t.Errorf("element count not configured correctly. got %v, want %v", got, want)
+ }
+}
+
+func TestAfterProcessingTimeTrigger(t *testing.T) {
+ tests := []struct {
+ tr *AfterProcessingTimeTrigger
+ tt []TimestampTransform
+ }{
+ {
+ tr: AfterProcessingTime().PlusDelay(time.Millisecond),
+ tt: []TimestampTransform{DelayTransform{Delay: 1}},
+ },
+ {
+ tr: AfterProcessingTime().PlusDelay(time.Millisecond).AlignedTo(time.Millisecond, time.Time{}),
+ tt: []TimestampTransform{DelayTransform{Delay: 1}, AlignToTransform{Period: 1, Offset: 0}},
+ },
+ }
+ for _, test := range tests {
+ if diff := cmp.Diff(test.tr.TimestampTransforms(), test.tt); diff != "" {
+ t.Errorf("timestamp transforms are not equal: %v", diff)
+ }
+ }
+}
+
+func TestRepeatTrigger(t *testing.T) {
+ subTr := AfterCount(2)
+ tr := Repeat(subTr)
+
+ if got, ok := tr.SubTrigger().(*AfterCountTrigger); ok && got != subTr {
+ t.Errorf("subtrigger not configured correctly. got %v, want %v", got, subTr)
+ }
+}
+
+func TestAfterEndOfWindowTrigger(t *testing.T) {
+ earlyTr := AfterCount(50)
+ lateTr := Always()
+ tr := AfterEndOfWindow().EarlyFiring(earlyTr).LateFiring(lateTr)
+
+ if got, ok := tr.Early().(*AfterCountTrigger); ok && got != earlyTr {
+ t.Errorf("early firing trigger not configured correctly. got %v, want %v", got, earlyTr)
+ }
+ if got, ok := tr.Late().(*AlwaysTrigger); ok && got != lateTr {
+ t.Errorf("late firing trigger not configured correctly. got %v, want %v", got, lateTr)
+ }
+}
diff --git a/sdks/go/pkg/beam/core/runtime/exec/coder_test.go b/sdks/go/pkg/beam/core/runtime/exec/coder_test.go
index 25812aca4e560..02d1f81da7e05 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/coder_test.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/coder_test.go
@@ -25,6 +25,7 @@ import (
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/coderx"
)
@@ -54,6 +55,9 @@ func TestCoders(t *testing.T) {
return &coder.Coder{Kind: coder.Custom, Custom: c, T: typex.New(reflectx.String)}
}(),
val: &FullValue{Elm: "myString"},
+ }, {
+ coder: &coder.Coder{Kind: coder.LP, Components: []*coder.Coder{coder.NewString()}},
+ val: &FullValue{Elm: "myString"},
}, {
coder: coder.NewKV([]*coder.Coder{coder.NewVarInt(), coder.NewBool()}),
val: &FullValue{Elm: int64(72), Elm2: false},
@@ -64,6 +68,18 @@ func TestCoders(t *testing.T) {
coder.NewDouble(),
coder.NewBool()})}),
val: &FullValue{Elm: int64(42), Elm2: &FullValue{Elm: float64(3.14), Elm2: true}},
+ }, {
+ coder: &coder.Coder{Kind: coder.Window, Window: coder.NewGlobalWindow()},
+ val: &FullValue{Windows: []typex.Window{window.GlobalWindow{}}},
+ }, {
+ coder: &coder.Coder{Kind: coder.Window, Window: coder.NewIntervalWindow()},
+ val: &FullValue{Windows: []typex.Window{window.IntervalWindow{Start: 0, End: 100}}},
+ }, {
+ coder: coder.NewW(coder.NewVarInt(), coder.NewGlobalWindow()),
+ val: &FullValue{Elm: int64(13), Windows: []typex.Window{window.GlobalWindow{}}},
+ }, {
+ coder: coder.NewPW(coder.NewString(), coder.NewGlobalWindow()),
+ val: &FullValue{Elm: "myString" /*Windowing info isn't encoded for PW so we can omit it here*/},
},
} {
t.Run(fmt.Sprintf("%v", test.coder), func(t *testing.T) {
@@ -116,4 +132,140 @@ func compareFV(t *testing.T, got *FullValue, want *FullValue) {
t.Errorf("got %v [type: %s], want %v [type %s]",
got, reflect.TypeOf(got), wantFv, reflect.TypeOf(wantFv))
}
+
+ // Check if the desired FV has windowing information
+ if want.Windows != nil {
+ if gotLen, wantLen := len(got.Windows), len(want.Windows); gotLen != wantLen {
+ t.Fatalf("got %d windows in FV, want %v", gotLen, wantLen)
+ }
+ for i := range want.Windows {
+ if gotWin, wantWin := got.Windows[i], want.Windows[i]; !wantWin.Equals(gotWin) {
+ t.Errorf("got window %v at position %d, want %v", gotWin, i, wantWin)
+ }
+ }
+ }
+}
+
+func TestIterableCoder(t *testing.T) {
+ cod := coder.NewI(coder.NewVarInt())
+ wantVals := []int64{8, 24, 72}
+ val := &FullValue{Elm: wantVals}
+
+ var buf bytes.Buffer
+ enc := MakeElementEncoder(cod)
+ if err := enc.Encode(val, &buf); err != nil {
+ t.Fatalf("Couldn't encode value: %v", err)
+ }
+
+ dec := MakeElementDecoder(cod)
+ result, err := dec.Decode(&buf)
+ if err != nil {
+ t.Fatalf("Couldn't decode value: %v", err)
+ }
+
+ gotVals, ok := result.Elm.([]int64)
+ if !ok {
+ t.Fatalf("got output element %v, want []int64", result.Elm)
+ }
+
+ if got, want := len(gotVals), len(wantVals); got != want {
+ t.Errorf("got %d elements in iterable, want %d", got, want)
+ }
+
+ for i := range gotVals {
+ if got, want := gotVals[i], wantVals[i]; got != want {
+ t.Errorf("got %d at position %d, want %d", got, i, want)
+ }
+ }
+}
+
+// TODO(BEAM-10660): Update once proper timer support is added
+func TestTimerCoder(t *testing.T) {
+ var buf bytes.Buffer
+ tCoder := coder.NewT(coder.NewVarInt(), coder.NewGlobalWindow())
+ wantVal := &FullValue{Elm: int64(13)}
+
+ enc := MakeElementEncoder(tCoder)
+ if err := enc.Encode(wantVal, &buf); err != nil {
+ t.Fatalf("Couldn't encode value: %v", err)
+ }
+
+ dec := MakeElementDecoder(tCoder)
+ result, err := dec.Decode(&buf)
+ if err != nil {
+ t.Fatalf("Couldn't decode value: %v", err)
+ }
+
+ compareFV(t, result, wantVal)
+}
+
+type namedTypeForTest struct {
+ A, B int64
+ C string
+}
+
+func TestRowCoder(t *testing.T) {
+ var buf bytes.Buffer
+ rCoder := coder.NewR(typex.New(reflect.TypeOf((*namedTypeForTest)(nil))))
+ wantStruct := &namedTypeForTest{A: int64(8), B: int64(24), C: "myString"}
+ wantVal := &FullValue{Elm: wantStruct}
+
+ enc := MakeElementEncoder(rCoder)
+ if err := enc.Encode(wantVal, &buf); err != nil {
+ t.Fatalf("Couldn't encode value: %v", err)
+ }
+
+ dec := MakeElementDecoder(rCoder)
+ result, err := dec.Decode(&buf)
+ if err != nil {
+ t.Fatalf("Couldn't decode value: %v", err)
+ }
+ gotPtr, ok := result.Elm.(*namedTypeForTest)
+ gotStruct := *gotPtr
+ if !ok {
+ t.Fatalf("got %v, want namedTypeForTest struct", result.Elm)
+ }
+ if got, want := gotStruct.A, wantStruct.A; got != want {
+ t.Errorf("got A field value %d, want %d", got, want)
+ }
+ if got, want := gotStruct.B, wantStruct.B; got != want {
+ t.Errorf("got B field value %d, want %d", got, want)
+ }
+ if got, want := gotStruct.C, wantStruct.C; got != want {
+ t.Errorf("got C field value %v, want %v", got, want)
+ }
+}
+
+func TestPaneCoder(t *testing.T) {
+ pn := coder.NewPane(0x04)
+ val := &FullValue{Pane: pn}
+ cod := &coder.Coder{Kind: coder.PaneInfo}
+
+ var buf bytes.Buffer
+ enc := MakeElementEncoder(cod)
+ if err := enc.Encode(val, &buf); err != nil {
+ t.Fatalf("Couldn't encode value: %v", err)
+ }
+
+ dec := MakeElementDecoder(cod)
+ result, err := dec.Decode(&buf)
+ if err != nil {
+ t.Fatalf("Couldn't decode value: %v", err)
+ }
+
+ if got, want := result.Pane.Timing, pn.Timing; got != want {
+ t.Errorf("got pane timing %v, want %v", got, want)
+ }
+ if got, want := result.Pane.IsFirst, pn.IsFirst; got != want {
+ t.Errorf("got IsFirst %v, want %v", got, want)
+ }
+ if got, want := result.Pane.IsLast, pn.IsLast; got != want {
+ t.Errorf("got IsLast %v, want %v", got, want)
+ }
+ if got, want := result.Pane.Index, pn.Index; got != want {
+ t.Errorf("got pane index %v, want %v", got, want)
+ }
+ if got, want := result.Pane.NonSpeculativeIndex, pn.NonSpeculativeIndex; got != want {
+ t.Errorf("got pane non-speculative index %v, want %v", got, want)
+ }
}
diff --git a/sdks/go/pkg/beam/core/util/stringx/bytes.go b/sdks/go/pkg/beam/core/util/stringx/bytes.go
index b2110fa1072ce..bf4cd111dddec 100644
--- a/sdks/go/pkg/beam/core/util/stringx/bytes.go
+++ b/sdks/go/pkg/beam/core/util/stringx/bytes.go
@@ -15,6 +15,9 @@
// Package stringx contains utilities for working with strings. It
// complements the standard "strings" package.
+//
+// Deprecated: the utilities in this package are unused within the code base
+// and will be removed in a future Beam release.
package stringx
// ToBytes converts a string to a byte slice.
diff --git a/sdks/go/pkg/beam/internal/errors/errors_test.go b/sdks/go/pkg/beam/internal/errors/errors_test.go
index 614609ff80b89..26e0057a74679 100644
--- a/sdks/go/pkg/beam/internal/errors/errors_test.go
+++ b/sdks/go/pkg/beam/internal/errors/errors_test.go
@@ -57,6 +57,9 @@ func TestWrap(t *testing.T) {
}, {
err: Wrap(Wrap(New(base), msg1), msg2),
want: errorStructure{{ERROR, msg2}, {ERROR, msg1}, {ERROR, base}},
+ }, {
+ err: Wrap(nil, msg1),
+ want: nil,
},
}
for _, test := range tests {
@@ -76,6 +79,13 @@ func TestWrapf(t *testing.T) {
}
}
+func TestWrapf_NilErr(t *testing.T) {
+ err := Wrapf(nil, "%s %d", "ten", 10)
+ if err != nil {
+ t.Errorf(`Wrapf(nil, "%%s %%d", "ten", 10). Want: nil, Got: %q`, err)
+ }
+}
+
func TestContext(t *testing.T) {
tests := []struct {
err error
@@ -90,6 +100,9 @@ func TestContext(t *testing.T) {
}, {
err: Wrap(WithContext(WithContext(Wrap(New(base), msg1), ctx1), ctx2), msg2),
want: errorStructure{{ERROR, msg2}, {CONTEXT, ctx2}, {CONTEXT, ctx1}, {ERROR, msg1}, {ERROR, base}},
+ }, {
+ err: WithContext(nil, ctx1),
+ want: nil,
},
}
for _, test := range tests {
@@ -100,6 +113,13 @@ func TestContext(t *testing.T) {
}
}
+func TestWithContextf_NilErr(t *testing.T) {
+ err := WithContextf(nil, "%s %d", "ten", 10)
+ if err != nil {
+ t.Errorf(`WithContextf(nil, "%%s %%d", "ten", 10). Want: nil, Got: %q`, err)
+ }
+}
+
func TestWithContextf(t *testing.T) {
want := fmt.Sprintf("%s %d", "ten", 10)
err := WithContextf(New(base), "%s %d", "ten", 10)
@@ -129,6 +149,9 @@ func TestTopLevelMsg(t *testing.T) {
}, {
err: Wrap(SetTopLevelMsg(WithContext(SetTopLevelMsg(New(base), top1), ctx1), top2), msg1),
want: top2,
+ }, {
+ err: SetTopLevelMsg(nil, top1),
+ want: "",
},
}
for _, test := range tests {
@@ -147,10 +170,51 @@ func TestSetTopLevelMsgf(t *testing.T) {
}
}
+func TestSetTopLevelMsgf_NilErr(t *testing.T) {
+ want := ""
+ err := SetTopLevelMsgf(nil, "%s %d", "ten", 10)
+ if getTop(err) != want {
+ t.Errorf("Incorrect formatting. Want: %q, Got: %q", want, getTop(err))
+ }
+}
+
+func TestError(t *testing.T) {
+ tests := []struct {
+ err error
+ want string
+ }{
+ {
+ err: Wrap(New(base), msg1),
+ want: "message 1\n\tcaused by:\nbase",
+ },
+ {
+ err: SetTopLevelMsg(New(base), top1),
+ want: "top level message 1\nFull error:\nbase",
+ },
+ {
+ err: SetTopLevelMsg(Wrap(Wrap(New(base), msg1), msg2), top1),
+ want: "top level message 1\nFull error:\nmessage 2\n\tcaused by:\nmessage 1\n\tcaused by:\nbase",
+ },
+ }
+
+ for _, test := range tests {
+ if be, ok := test.err.(*beamError); ok {
+ if got, want := be.Error(), test.want; got != want {
+ t.Errorf("Incorrect formatting. Want: %q, Got: %q", want, got)
+ }
+ } else {
+ t.Errorf("Error should be type *beamError, got: %q", test.err)
+ }
+ }
+}
+
// getStructure extracts the structure of an error, outputting a slice that
// represents the nested messages in that error in the order they are output
// and with the type of message (context or error) described.
func getStructure(e error) errorStructure {
+ if e == nil {
+ return nil
+ }
var structure errorStructure
for {
@@ -175,6 +239,10 @@ func getStructure(e error) errorStructure {
}
func equalStructure(left errorStructure, right errorStructure) bool {
+ if left == nil || right == nil {
+ return left == nil && right == nil
+ }
+
if len(left) != len(right) {
return false
}
diff --git a/sdks/go/pkg/beam/runners/dataflow/dataflow.go b/sdks/go/pkg/beam/runners/dataflow/dataflow.go
index 4cd877f5afde0..1a7da076755a0 100644
--- a/sdks/go/pkg/beam/runners/dataflow/dataflow.go
+++ b/sdks/go/pkg/beam/runners/dataflow/dataflow.go
@@ -159,8 +159,59 @@ func Execute(ctx context.Context, p *beam.Pipeline) (beam.PipelineResult, error)
panic("Beam has not been initialized. Call beam.Init() before pipeline construction.")
}
- // (1) Gather job options
+ beam.PipelineOptions.LoadOptionsFromFlags(flagFilter)
+ opts, err := getJobOptions(ctx)
+ if err != nil {
+ return nil, err
+ }
+
+ // (1) Build and submit
+ // NOTE(herohde) 10/8/2018: the last segment of the names must be "worker" and "dataflow-worker.jar".
+ id := fmt.Sprintf("go-%v-%v", atomic.AddInt32(&unique, 1), time.Now().UnixNano())
+
+ modelURL := gcsx.Join(*stagingLocation, id, "model")
+ workerURL := gcsx.Join(*stagingLocation, id, "worker")
+ jarURL := gcsx.Join(*stagingLocation, id, "dataflow-worker.jar")
+ xlangURL := gcsx.Join(*stagingLocation, id, "xlang")
+
+ edges, _, err := p.Build()
+ if err != nil {
+ return nil, err
+ }
+ artifactURLs, err := dataflowlib.ResolveXLangArtifacts(ctx, edges, opts.Project, xlangURL)
+ if err != nil {
+ return nil, errors.WithContext(err, "resolving cross-language artifacts")
+ }
+ opts.ArtifactURLs = artifactURLs
+ environment, err := graphx.CreateEnvironment(ctx, jobopts.GetEnvironmentUrn(ctx), getContainerImage)
+ if err != nil {
+ return nil, errors.WithContext(err, "creating environment for model pipeline")
+ }
+ model, err := graphx.Marshal(edges, &graphx.Options{Environment: environment})
+ if err != nil {
+ return nil, errors.WithContext(err, "generating model pipeline")
+ }
+ err = pipelinex.ApplySdkImageOverrides(model, jobopts.GetSdkImageOverrides())
+ if err != nil {
+ return nil, errors.WithContext(err, "applying container image overrides")
+ }
+
+ if *dryRun {
+ log.Info(ctx, "Dry-run: not submitting job!")
+
+ log.Info(ctx, proto.MarshalTextString(model))
+ job, err := dataflowlib.Translate(ctx, model, opts, workerURL, jarURL, modelURL)
+ if err != nil {
+ return nil, err
+ }
+ dataflowlib.PrintJob(ctx, job)
+ return nil, nil
+ }
+ return dataflowlib.Execute(ctx, model, opts, workerURL, jarURL, modelURL, *endpoint, *executeAsync)
+}
+
+func getJobOptions(ctx context.Context) (*dataflowlib.JobOptions, error) {
project := gcpopts.GetProjectFromFlagOrEnvironment(ctx)
if project == "" {
return nil, errors.New("no Google Cloud project specified. Use --project=")
@@ -254,51 +305,9 @@ func Execute(ctx context.Context, p *beam.Pipeline) (beam.PipelineResult, error)
opts.TempLocation = gcsx.Join(*stagingLocation, "tmp")
}
- // (1) Build and submit
- // NOTE(herohde) 10/8/2018: the last segment of the names must be "worker" and "dataflow-worker.jar".
- id := fmt.Sprintf("go-%v-%v", atomic.AddInt32(&unique, 1), time.Now().UnixNano())
-
- modelURL := gcsx.Join(*stagingLocation, id, "model")
- workerURL := gcsx.Join(*stagingLocation, id, "worker")
- jarURL := gcsx.Join(*stagingLocation, id, "dataflow-worker.jar")
- xlangURL := gcsx.Join(*stagingLocation, id, "xlang")
-
- edges, _, err := p.Build()
- if err != nil {
- return nil, err
- }
- artifactURLs, err := dataflowlib.ResolveXLangArtifacts(ctx, edges, opts.Project, xlangURL)
- if err != nil {
- return nil, errors.WithContext(err, "resolving cross-language artifacts")
- }
- opts.ArtifactURLs = artifactURLs
- environment, err := graphx.CreateEnvironment(ctx, jobopts.GetEnvironmentUrn(ctx), getContainerImage)
- if err != nil {
- return nil, errors.WithContext(err, "creating environment for model pipeline")
- }
- model, err := graphx.Marshal(edges, &graphx.Options{Environment: environment})
- if err != nil {
- return nil, errors.WithContext(err, "generating model pipeline")
- }
- err = pipelinex.ApplySdkImageOverrides(model, jobopts.GetSdkImageOverrides())
- if err != nil {
- return nil, errors.WithContext(err, "applying container image overrides")
- }
-
- if *dryRun {
- log.Info(ctx, "Dry-run: not submitting job!")
-
- log.Info(ctx, proto.MarshalTextString(model))
- job, err := dataflowlib.Translate(ctx, model, opts, workerURL, jarURL, modelURL)
- if err != nil {
- return nil, err
- }
- dataflowlib.PrintJob(ctx, job)
- return nil, nil
- }
-
- return dataflowlib.Execute(ctx, model, opts, workerURL, jarURL, modelURL, *endpoint, *executeAsync)
+ return opts, nil
}
+
func gcsRecorderHook(opts []string) perf.CaptureHook {
bucket, prefix, err := gcsx.ParseObject(opts[0])
if err != nil {
diff --git a/sdks/go/pkg/beam/runners/dataflow/dataflow_test.go b/sdks/go/pkg/beam/runners/dataflow/dataflow_test.go
index 1e2844630c4b3..568860194ec18 100644
--- a/sdks/go/pkg/beam/runners/dataflow/dataflow_test.go
+++ b/sdks/go/pkg/beam/runners/dataflow/dataflow_test.go
@@ -15,7 +15,12 @@
package dataflow
-import "testing"
+import (
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/options/jobopts"
+ "sort"
+ "testing"
+)
func TestDontUseFlagAsPipelineOption(t *testing.T) {
f := "dummy_flag"
@@ -27,3 +32,131 @@ func TestDontUseFlagAsPipelineOption(t *testing.T) {
t.Fatalf("%q should be in the filter, but isn't set", f)
}
}
+
+func TestGetJobOptions(t *testing.T) {
+ *labels = `{"label1": "val1", "label2": "val2"}`
+ *stagingLocation = "gs://testStagingLocation"
+ *autoscalingAlgorithm = "NONE"
+ *minCPUPlatform = "testPlatform"
+
+ *gcpopts.Project = "testProject"
+ *gcpopts.Region = "testRegion"
+
+ *jobopts.Experiments = "use_runner_v2,use_portable_job_submission"
+ *jobopts.JobName = "testJob"
+
+ opts, err := getJobOptions(nil)
+ if err != nil {
+ t.Fatalf("getJobOptions() returned error %q, want %q", err, "nil")
+ }
+ if got, want := opts.Name, "testJob"; got != want {
+ t.Errorf("getJobOptions().Name = %q, want %q", got, want)
+ }
+ if got, want := len(opts.Experiments), 3; got != want {
+ t.Errorf("len(getJobOptions().Experiments) = %q, want %q", got, want)
+ } else {
+ sort.Strings(opts.Experiments)
+ expectedExperiments := []string{"min_cpu_platform=testPlatform", "use_portable_job_submission", "use_runner_v2"}
+ for i := 0; i < 3; i++ {
+ if got, want := opts.Experiments[i], expectedExperiments[i]; got != want {
+ t.Errorf("getJobOptions().Experiments = %q, want %q", got, want)
+ }
+ }
+ }
+ if got, want := opts.Project, "testProject"; got != want {
+ t.Errorf("getJobOptions().Project = %q, want %q", got, want)
+ }
+ if got, want := opts.Region, "testRegion"; got != want {
+ t.Errorf("getJobOptions().Region = %q, want %q", got, want)
+ }
+ if got, want := len(opts.Labels), 2; got != want {
+ t.Errorf("len(getJobOptions().Labels) = %q, want %q", got, want)
+ } else {
+ if got, want := opts.Labels["label1"], "val1"; got != want {
+ t.Errorf("getJobOptions().Labels[\"label1\"] = %q, want %q", got, want)
+ }
+ if got, want := opts.Labels["label2"], "val2"; got != want {
+ t.Errorf("getJobOptions().Labels[\"label2\"] = %q, want %q", got, want)
+ }
+ }
+ if got, want := opts.TempLocation, "gs://testStagingLocation/tmp"; got != want {
+ t.Errorf("getJobOptions().TempLocation = %q, want %q", got, want)
+ }
+}
+
+func TestGetJobOptions_NoExperimentsSet(t *testing.T) {
+ *labels = `{"label1": "val1", "label2": "val2"}`
+ *stagingLocation = "gs://testStagingLocation"
+ *autoscalingAlgorithm = "NONE"
+ *minCPUPlatform = ""
+
+ *gcpopts.Project = "testProject"
+ *gcpopts.Region = "testRegion"
+
+ *jobopts.Experiments = ""
+ *jobopts.JobName = "testJob"
+
+ opts, err := getJobOptions(nil)
+
+ if err != nil {
+ t.Fatalf("getJobOptions() returned error %q, want %q", err, "nil")
+ }
+ if got, want := len(opts.Experiments), 2; got != want {
+ t.Fatalf("len(getJobOptions().Experiments) = %q, want %q", got, want)
+ }
+ sort.Strings(opts.Experiments)
+ expectedExperiments := []string{"use_portable_job_submission", "use_unified_worker"}
+ for i := 0; i < 2; i++ {
+ if got, want := opts.Experiments[i], expectedExperiments[i]; got != want {
+ t.Errorf("getJobOptions().Experiments = %q, want %q", got, want)
+ }
+ }
+}
+
+func TestGetJobOptions_NoStagingLocation(t *testing.T) {
+ *stagingLocation = ""
+ *gcpopts.Project = "testProject"
+ *gcpopts.Region = "testRegion"
+
+ _, err := getJobOptions(nil)
+ if err == nil {
+ t.Fatalf("getJobOptions() returned error nil, want an error")
+ }
+}
+
+func TestGetJobOptions_InvalidAutoscaling(t *testing.T) {
+ *labels = `{"label1": "val1", "label2": "val2"}`
+ *stagingLocation = "gs://testStagingLocation"
+ *autoscalingAlgorithm = "INVALID"
+ *minCPUPlatform = "testPlatform"
+
+ *gcpopts.Project = "testProject"
+ *gcpopts.Region = "testRegion"
+
+ *jobopts.Experiments = "use_runner_v2,use_portable_job_submission"
+ *jobopts.JobName = "testJob"
+
+ _, err := getJobOptions(nil)
+ if err == nil {
+ t.Fatalf("getJobOptions() returned error nil, want an error")
+ }
+}
+
+func TestGetJobOptions_DockerNoImage(t *testing.T) {
+ *jobopts.EnvironmentType = "docker"
+ *jobopts.EnvironmentConfig = "testContainerImage"
+
+ if got, want := getContainerImage(nil), "testContainerImage"; got != want {
+ t.Fatalf("getContainerImage() = %q, want %q", got, want)
+ }
+}
+
+func TestGetJobOptions_DockerWithImage(t *testing.T) {
+ *jobopts.EnvironmentType = "docker"
+ *jobopts.EnvironmentConfig = "testContainerImage"
+ *image = "testContainerImageOverride"
+
+ if got, want := getContainerImage(nil), "testContainerImageOverride"; got != want {
+ t.Fatalf("getContainerImage() = %q, want %q", got, want)
+ }
+}
diff --git a/sdks/go/pkg/beam/util/errorx/guarded.go b/sdks/go/pkg/beam/util/errorx/guarded.go
index cc0b07b4eee1f..186885b717c62 100644
--- a/sdks/go/pkg/beam/util/errorx/guarded.go
+++ b/sdks/go/pkg/beam/util/errorx/guarded.go
@@ -39,7 +39,7 @@ func (g *GuardedError) TrySetError(err error) bool {
g.mu.Lock()
defer g.mu.Unlock()
- upd := g.err == nil
+ upd := (g.err == nil)
if upd {
g.err = err
}
diff --git a/sdks/go/pkg/beam/util/errorx/guarded_test.go b/sdks/go/pkg/beam/util/errorx/guarded_test.go
new file mode 100644
index 0000000000000..1e9c9b2247060
--- /dev/null
+++ b/sdks/go/pkg/beam/util/errorx/guarded_test.go
@@ -0,0 +1,45 @@
+// 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 errorx
+
+import (
+ "errors"
+ "testing"
+)
+
+func TestTrySetError(t *testing.T) {
+ var gu GuardedError
+ setErr := errors.New("attempted error")
+ success := gu.TrySetError(setErr)
+ if !success {
+ t.Fatal("got false when trying to set error, want true")
+ }
+ if got, want := gu.Error(), setErr; got != want {
+ t.Errorf("got error %v when checking message, want %v", got, want)
+ }
+}
+
+func TestTrySetError_bad(t *testing.T) {
+ setErr := errors.New("old error")
+ gu := &GuardedError{err: setErr}
+ success := gu.TrySetError(setErr)
+ if success {
+ t.Fatal("got true when trying to set error, want false")
+ }
+ if got, want := gu.Error(), setErr; got != want {
+ t.Errorf("got error %v when checking message, want %v", got, want)
+ }
+}
diff --git a/playground/backend/internal/utils/validators_utils.go b/sdks/go/pkg/beam/util/gcsx/example_test.go
similarity index 54%
rename from playground/backend/internal/utils/validators_utils.go
rename to sdks/go/pkg/beam/util/gcsx/example_test.go
index 5d9406e912df7..52664a054a878 100644
--- a/playground/backend/internal/utils/validators_utils.go
+++ b/sdks/go/pkg/beam/util/gcsx/example_test.go
@@ -13,28 +13,35 @@
// See the License for the specific language governing permissions and
// limitations under the License.
-package utils
+package gcsx_test
import (
- pb "beam.apache.org/playground/backend/internal/api/v1"
- "beam.apache.org/playground/backend/internal/validators"
- "fmt"
+ "context"
+ "time"
+
+ "cloud.google.com/go/storage"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/util/gcsx"
)
-// GetValidators returns slice of validators.Validator according to sdk
-func GetValidators(sdk pb.Sdk, filepath string) (*[]validators.Validator, error) {
- var val *[]validators.Validator
- switch sdk {
- case pb.Sdk_SDK_JAVA:
- val = validators.GetJavaValidators(filepath)
- case pb.Sdk_SDK_GO:
- val = validators.GetGoValidators(filepath)
- case pb.Sdk_SDK_PYTHON:
- val = validators.GetPyValidators(filepath)
- case pb.Sdk_SDK_SCIO:
- val = validators.GetScioValidators(filepath)
- default:
- return nil, fmt.Errorf("incorrect sdk: %s", sdk)
+func Example() {
+ ctx := context.Background()
+ c, err := gcsx.NewClient(ctx, storage.ScopeReadOnly)
+ if err != nil {
+ // do something
+ }
+
+ buckets, object, err := gcsx.ParseObject("gs://some-bucket/some-object")
+ if err != nil {
+ // do something
}
- return val, nil
+
+ ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
+ defer cancel()
+
+ bytes, err := gcsx.ReadObject(ctx, c, buckets, object)
+ if err != nil {
+ // do something
+ }
+
+ _ = bytes
}
diff --git a/sdks/go/pkg/beam/util/gcsx/gcs_test.go b/sdks/go/pkg/beam/util/gcsx/gcs_test.go
index 52664a054a878..90fb4b59f2fe8 100644
--- a/sdks/go/pkg/beam/util/gcsx/gcs_test.go
+++ b/sdks/go/pkg/beam/util/gcsx/gcs_test.go
@@ -13,35 +13,86 @@
// See the License for the specific language governing permissions and
// limitations under the License.
-package gcsx_test
+package gcsx
import (
- "context"
- "time"
+ "strings"
+ "testing"
- "cloud.google.com/go/storage"
- "github.com/apache/beam/sdks/v2/go/pkg/beam/util/gcsx"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
)
-func Example() {
- ctx := context.Background()
- c, err := gcsx.NewClient(ctx, storage.ScopeReadOnly)
- if err != nil {
- // do something
+func TestMakeObject(t *testing.T) {
+ if got, want := MakeObject("some-bucket", "some/path"), "gs://some-bucket/some/path"; got != want {
+ t.Fatalf("MakeObject() Got: %v Want: %v", got, want)
}
+}
- buckets, object, err := gcsx.ParseObject("gs://some-bucket/some-object")
- if err != nil {
- // do something
+func TestParseObject(t *testing.T) {
+ tests := []struct {
+ object string
+ bucket string
+ path string
+ err error
+ }{
+ {
+ object: "gs://some-bucket/some-object",
+ bucket: "some-bucket",
+ path: "some-object",
+ err: nil,
+ },
+ {
+ object: "gs://some-bucket",
+ bucket: "some-bucket",
+ path: "",
+ err: nil,
+ },
+ {
+ object: "gs://",
+ bucket: "",
+ path: "",
+ err: errors.Errorf("object gs:// must have bucket"),
+ },
+ {
+ object: "other://some-bucket/some-object",
+ bucket: "",
+ path: "",
+ err: errors.Errorf("object other://some-bucket/some-object must have 'gs' scheme"),
+ },
}
- ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
- defer cancel()
-
- bytes, err := gcsx.ReadObject(ctx, c, buckets, object)
- if err != nil {
- // do something
+ for _, test := range tests {
+ if bucket, path, err := ParseObject(test.object); bucket != test.bucket || path != test.path || (err != nil && test.err == nil) || (err == nil && test.err != nil) {
+ t.Errorf("ParseObject(%v) Got: %v, %v, %v Want: %v, %v, %v", test.object, bucket, path, err, test.bucket, test.path, test.err)
+ }
}
+}
- _ = bytes
+func TestJoin(t *testing.T) {
+ tests := []struct {
+ object string
+ elms []string
+ result string
+ }{
+ {
+ object: "gs://some-bucket/some-object",
+ elms: []string{"some/path", "more/pathing"},
+ result: "gs://some-bucket/some-object/some/path/more/pathing",
+ },
+ {
+ object: "gs://some-bucket/some-object",
+ elms: []string{"some/path"},
+ result: "gs://some-bucket/some-object/some/path",
+ },
+ {
+ object: "gs://some-bucket/some-object",
+ elms: []string{},
+ result: "gs://some-bucket/some-object",
+ },
+ }
+ for _, test := range tests {
+ if got, want := Join(test.object, test.elms...), test.result; got != want {
+ t.Errorf("Join(%v, %v) Got: %v Want: %v", test.object, strings.Join(test.elms, ", "), got, want)
+ }
+ }
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsFactory.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsFactory.java
index 4cb7eb4eac45f..d014f08ac2b8c 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsFactory.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsFactory.java
@@ -36,11 +36,13 @@
import com.fasterxml.jackson.databind.annotation.JsonSerialize;
import com.fasterxml.jackson.databind.deser.DefaultDeserializationContext;
import com.fasterxml.jackson.databind.deser.impl.MethodProperty;
+import com.fasterxml.jackson.databind.deser.impl.TypeWrappedDeserializer;
import com.fasterxml.jackson.databind.introspect.AnnotatedMember;
import com.fasterxml.jackson.databind.introspect.AnnotatedMethod;
import com.fasterxml.jackson.databind.introspect.AnnotationCollector;
import com.fasterxml.jackson.databind.introspect.BeanPropertyDefinition;
import com.fasterxml.jackson.databind.introspect.TypeResolutionContext;
+import com.fasterxml.jackson.databind.jsontype.TypeDeserializer;
import com.fasterxml.jackson.databind.node.TreeTraversingParser;
import com.fasterxml.jackson.databind.ser.DefaultSerializerProvider;
import com.fasterxml.jackson.databind.type.TypeBindings;
@@ -1730,21 +1732,23 @@ private static JsonDeserializer computeDeserializerForMethod(Method meth
BeanProperty prop = createBeanProperty(method);
AnnotatedMember annotatedMethod = prop.getMember();
+ DefaultDeserializationContext context = DESERIALIZATION_CONTEXT.get();
Object maybeDeserializerClass =
- DESERIALIZATION_CONTEXT
- .get()
- .getAnnotationIntrospector()
- .findDeserializer(annotatedMethod);
+ context.getAnnotationIntrospector().findDeserializer(annotatedMethod);
JsonDeserializer jsonDeserializer =
- DESERIALIZATION_CONTEXT
- .get()
- .deserializerInstance(annotatedMethod, maybeDeserializerClass);
+ context.deserializerInstance(annotatedMethod, maybeDeserializerClass);
if (jsonDeserializer == null) {
- jsonDeserializer =
- DESERIALIZATION_CONTEXT.get().findContextualValueDeserializer(prop.getType(), prop);
+ jsonDeserializer = context.findContextualValueDeserializer(prop.getType(), prop);
}
+
+ TypeDeserializer typeDeserializer =
+ context.getFactory().findTypeDeserializer(context.getConfig(), prop.getType());
+ if (typeDeserializer != null) {
+ jsonDeserializer = new TypeWrappedDeserializer(typeDeserializer, jsonDeserializer);
+ }
+
return jsonDeserializer;
} catch (JsonMappingException e) {
throw new RuntimeException(e);
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/options/PipelineOptionsFactoryTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/options/PipelineOptionsFactoryTest.java
index 94fd3f41faacd..ffdfbc8681a13 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/options/PipelineOptionsFactoryTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/options/PipelineOptionsFactoryTest.java
@@ -39,6 +39,8 @@
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonSubTypes;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.JsonProcessingException;
@@ -60,7 +62,6 @@
import java.util.Collection;
import java.util.List;
import java.util.Map;
-import java.util.Objects;
import java.util.Set;
import org.apache.beam.model.jobmanagement.v1.JobApi.PipelineOptionDescriptor;
import org.apache.beam.model.jobmanagement.v1.JobApi.PipelineOptionType;
@@ -1070,6 +1071,53 @@ public void testComplexTypes() {
assertEquals("value2", options.getObjectValue().get().value2);
}
+ @JsonTypeInfo(use = JsonTypeInfo.Id.NAME)
+ @JsonSubTypes({
+ @JsonSubTypes.Type(value = PolymorphicTypeOne.class, name = "one"),
+ @JsonSubTypes.Type(value = PolymorphicTypeTwo.class, name = "two")
+ })
+ public abstract static class PolymorphicType {
+ String key;
+
+ @JsonProperty("key")
+ public String getKey() {
+ return key;
+ }
+
+ public void setKey(String key) {
+ this.key = key;
+ }
+ }
+
+ public static class PolymorphicTypeOne extends PolymorphicType {}
+
+ public static class PolymorphicTypeTwo extends PolymorphicType {}
+
+ public interface PolymorphicTypes extends PipelineOptions {
+ PolymorphicType getObject();
+
+ void setObject(PolymorphicType value);
+
+ ValueProvider getObjectValue();
+
+ void setObjectValue(ValueProvider value);
+ }
+
+ @Test
+ public void testPolymorphicType() {
+ String[] args =
+ new String[] {
+ "--object={\"key\":\"value\",\"@type\":\"one\"}",
+ "--objectValue={\"key\":\"value\",\"@type\":\"two\"}"
+ };
+ PolymorphicTypes options = PipelineOptionsFactory.fromArgs(args).as(PolymorphicTypes.class);
+ assertEquals("value", options.getObject().key);
+ assertEquals(PolymorphicTypeOne.class, options.getObject().getClass());
+
+ assertEquals("value", options.getObjectValue().get().key);
+ assertEquals(PolymorphicTypeTwo.class, options.getObjectValue().get().getClass());
+ }
+
@Test
public void testMissingArgument() {
String[] args = new String[] {};
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
index 7a20afbd151e8..89c05b8f9f238 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
@@ -180,6 +180,7 @@ static class Factory
@@ -242,6 +243,7 @@ static class Factory onTimerContext;
private final OnWindowExpirationContext onWindowExpirationContext;
private final FinishBundleArgumentProvider finishBundleArgumentProvider;
+ private final Duration allowedLateness;
/**
* Used to guarantee a consistent view of this {@link FnApiDoFnRunner} while setting up for {@link
@@ -344,6 +346,7 @@ private interface TriFunction {
Map windowingStrategies,
Consumer addStartFunction,
Consumer addFinishFunction,
+ Consumer addResetFunction,
Consumer addTearDownFunction,
Function>> getPCollectionConsumer,
TriFunction>, Coder> addPCollectionConsumer,
@@ -457,6 +460,13 @@ private interface TriFunction {
}
timerFamilyInfos = timerFamilyInfosBuilder.build();
+ this.mainInputId = ParDoTranslation.getMainInputName(pTransform);
+ this.allowedLateness =
+ rehydratedComponents
+ .getPCollection(pTransform.getInputsOrThrow(mainInputId))
+ .getWindowingStrategy()
+ .getAllowedLateness();
+
} catch (IOException exn) {
throw new IllegalArgumentException("Malformed ParDoPayload", exn);
}
@@ -473,12 +483,11 @@ private interface TriFunction {
this.bundleFinalizer = bundleFinalizer;
this.onTimerContext = new OnTimerContext();
this.onWindowExpirationContext = new OnWindowExpirationContext<>();
+ this.timerBundleTracker =
+ new FnApiTimerBundleTracker(
+ keyCoder, windowCoder, this::getCurrentKey, () -> currentWindow);
+ addResetFunction.accept(timerBundleTracker::reset);
- try {
- this.mainInputId = ParDoTranslation.getMainInputName(pTransform);
- } catch (IOException e) {
- throw new RuntimeException(e);
- }
this.mainOutputConsumers =
(Collection>>)
(Collection) localNameToConsumer.get(mainOutputTag.getId());
@@ -756,9 +765,6 @@ private Object getCurrentKey() {
}
private void startBundle() {
- timerBundleTracker =
- new FnApiTimerBundleTracker(
- keyCoder, windowCoder, this::getCurrentKey, () -> currentWindow);
doFnInvoker.invokeStartBundle(startBundleArgumentProvider);
}
@@ -1694,14 +1700,9 @@ private void processTimer(
// The timerIdOrTimerFamilyId contains either a timerId from timer declaration or
// timerFamilyId
// from timer family declaration.
- String timerId =
- timerIdOrTimerFamilyId.startsWith(TimerFamilyDeclaration.PREFIX)
- ? ""
- : timerIdOrTimerFamilyId;
- String timerFamilyId =
- timerIdOrTimerFamilyId.startsWith(TimerFamilyDeclaration.PREFIX)
- ? timerIdOrTimerFamilyId
- : "";
+ boolean isFamily = timerIdOrTimerFamilyId.startsWith(TimerFamilyDeclaration.PREFIX);
+ String timerId = isFamily ? "" : timerIdOrTimerFamilyId;
+ String timerFamilyId = isFamily ? timerIdOrTimerFamilyId : "";
processTimerDirect(timerFamilyId, timerId, timeDomain, timer);
}
}
@@ -1778,7 +1779,6 @@ private class FnApiTimer implements org.apache.beam.sdk.state.Timer {
private final K userKey;
private final String dynamicTimerTag;
private final TimeDomain timeDomain;
- private final Duration allowedLateness;
private final Instant fireTimestamp;
private final Instant elementTimestampOrTimerHoldTimestamp;
private final BoundedWindow boundedWindow;
@@ -1817,18 +1817,6 @@ private class FnApiTimer implements org.apache.beam.sdk.state.Timer {
throw new IllegalArgumentException(
String.format("Unknown or unsupported time domain %s", timeDomain));
}
-
- try {
- this.allowedLateness =
- rehydratedComponents
- .getPCollection(
- pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform)))
- .getWindowingStrategy()
- .getAllowedLateness();
- } catch (IOException e) {
- throw new IllegalArgumentException(
- String.format("Unable to get allowed lateness for timer %s", timerIdOrFamily));
- }
}
@Override
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiTimerBundleTracker.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiTimerBundleTracker.java
index f8a1dc3f9a258..39f735d98efe4 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiTimerBundleTracker.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiTimerBundleTracker.java
@@ -153,6 +153,10 @@ public FnApiTimerBundleTracker(
});
}
+ public void reset() {
+ timerModifications.clear();
+ }
+
public void timerModified(String timerFamilyOrId, TimeDomain timeDomain, Timer timer) {
ByteString keyString = encodedCurrentKeySupplier.get();
ByteString windowString = encodedCurrentWindowSupplier.get();
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
index 13b8c24ed74af..9970e802ee340 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
@@ -509,6 +509,7 @@ public void testOrderOfSetupTeardownCalls() throws Exception {
PCollection.newBuilder()
.setWindowingStrategyId("window-strategy")
.setCoderId("2L-output-coder")
+ .setIsBounded(IsBounded.Enum.BOUNDED)
.build())
.putWindowingStrategies(
"window-strategy",
diff --git a/sdks/java/io/amazon-web-services2/build.gradle b/sdks/java/io/amazon-web-services2/build.gradle
index 1af0532fd13c8..dceb4c41bed91 100644
--- a/sdks/java/io/amazon-web-services2/build.gradle
+++ b/sdks/java/io/amazon-web-services2/build.gradle
@@ -59,10 +59,10 @@ dependencies {
testImplementation project(path: ":sdks:java:io:common", configuration: "testRuntimeMigration")
testImplementation project(path: ":sdks:java:io:kinesis", configuration: "testRuntimeMigration")
testImplementation "io.findify:s3mock_2.12:0.2.6"
+ testImplementation 'org.elasticmq:elasticmq-rest-sqs_2.12:1.3.5'
testImplementation library.java.mockito_core
testImplementation library.java.guava_testlib
testImplementation library.java.junit
- testImplementation 'org.elasticmq:elasticmq-rest-sqs_2.12:0.15.6' // later versions conflict with s3mock
testImplementation library.java.hamcrest
testImplementation "org.assertj:assertj-core:3.11.1"
testRuntimeOnly library.java.slf4j_jdk14
diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/sqs/SqsUnboundedReaderTest.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/sqs/SqsUnboundedReaderTest.java
index 2c2329974df2b..5b1ac0cbb28b2 100644
--- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/sqs/SqsUnboundedReaderTest.java
+++ b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/sqs/SqsUnboundedReaderTest.java
@@ -37,7 +37,6 @@
import org.apache.beam.sdk.io.aws2.sqs.EmbeddedSqsServer.TestCaseEnv;
import org.apache.beam.sdk.util.CoderUtils;
import org.junit.ClassRule;
-import org.junit.Ignore;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -110,7 +109,6 @@ public void testAckDeletedMessage() throws IOException {
}
@Test
- @Ignore("Behavior of SQSRestServer is broken: https://issues.apache.org/jira/browse/BEAM-13738")
public void testExtendDeletedMessage() throws IOException {
setupMessages(DATA);
Clock clock = mock(Clock.class);
diff --git a/sdks/java/io/amazon-web-services2/src/test/resources/application.conf b/sdks/java/io/amazon-web-services2/src/test/resources/application.conf
new file mode 100644
index 0000000000000..f1fcb1e9f3da4
--- /dev/null
+++ b/sdks/java/io/amazon-web-services2/src/test/resources/application.conf
@@ -0,0 +1,21 @@
+###############################################################################
+# 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.
+###############################################################################
+
+# Enable implicit HEAD requests using GET routes for S3Mock
+# https://doc.akka.io/docs/akka-http/current/migration-guide/migration-guide-10.2.x.html#transparent-head-requests-now-disabled-by-default
+akka.http.server.transparent-head-requests = on
diff --git a/sdks/java/io/google-cloud-platform/build.gradle b/sdks/java/io/google-cloud-platform/build.gradle
index be4e551cd6b9a..364321e0b27ed 100644
--- a/sdks/java/io/google-cloud-platform/build.gradle
+++ b/sdks/java/io/google-cloud-platform/build.gradle
@@ -52,6 +52,7 @@ dependencies {
// BEAM-13781: gax-grpc's gRPC version was older than Beam declared
exclude group: 'io.grpc', module: 'grpc-netty-shaded'
}
+ implementation library.java.gax_grpc_test
implementation library.java.gax_httpjson
permitUnusedDeclared library.java.gax_httpjson // BEAM-8755
implementation library.java.google_api_client
@@ -146,6 +147,7 @@ dependencies {
testImplementation library.java.powermock
testImplementation library.java.powermock_mockito
testImplementation library.java.joda_time
+ testImplementation library.java.google_cloud_spanner_test
testRuntimeOnly library.java.slf4j_jdk14
}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
index f6d433498fde1..97a9bfa1595c7 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
@@ -74,7 +74,6 @@
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.TableSchemaToJsonSchema;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.TableSpecToTableRef;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.TimePartitioningToJson;
-import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead.Method;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryResourceNaming.JobType;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.DatasetService;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.JobService;
@@ -91,6 +90,8 @@
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+import org.apache.beam.sdk.schemas.FieldAccessDescriptor;
+import org.apache.beam.sdk.schemas.ProjectionProducer;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
@@ -117,6 +118,7 @@
import org.apache.beam.sdk.values.ValueInSingleWindow;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Predicates;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
@@ -730,7 +732,8 @@ public Read withTemplateCompatibility() {
/** Implementation of {@link BigQueryIO#read(SerializableFunction)}. */
@AutoValue
- public abstract static class TypedRead extends PTransform> {
+ public abstract static class TypedRead extends PTransform>
+ implements ProjectionProducer>> {
/** Determines the method used to read data from BigQuery. */
public enum Method {
/** The default behavior if no method is explicitly set. Currently {@link #EXPORT}. */
@@ -1619,6 +1622,29 @@ public TypedRead withTestServices(BigQueryServices testServices) {
public TypedRead useAvroLogicalTypes() {
return toBuilder().setUseAvroLogicalTypes(true).build();
}
+
+ @Override
+ public boolean supportsProjectionPushdown() {
+ // We can't do projection pushdown when a query is set. The query may project certain fields
+ // itself, and we can't know without parsing the query.
+ return Method.DIRECT_READ.equals(getMethod()) && getQuery() == null;
+ }
+
+ @Override
+ public PTransform> actuateProjectionPushdown(
+ Map, FieldAccessDescriptor> outputFields) {
+ Preconditions.checkArgument(supportsProjectionPushdown());
+ FieldAccessDescriptor fieldAccessDescriptor = outputFields.get(new TupleTag<>("output"));
+ org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull(
+ fieldAccessDescriptor, "Expected pushdown on the main output (tagged 'output')");
+ Preconditions.checkArgument(
+ outputFields.size() == 1,
+ "Expected only to pushdown on the main output (tagged 'output'). Requested tags: %s",
+ outputFields.keySet());
+ ImmutableList fields =
+ ImmutableList.copyOf(fieldAccessDescriptor.fieldNamesAccessed());
+ return withSelectedFields(fields);
+ }
}
static String getExtractDestinationUri(String extractDestinationDir) {
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessor.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessor.java
index e34863ea97de0..a4223dc804eaf 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessor.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessor.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.io.gcp.spanner;
+import com.google.api.gax.grpc.testing.LocalChannelProvider;
import com.google.api.gax.retrying.RetrySettings;
import com.google.api.gax.rpc.FixedHeaderProvider;
import com.google.api.gax.rpc.ServerStreamingCallSettings;
@@ -107,34 +108,59 @@ public static SpannerAccessor getOrCreate(SpannerConfig spannerConfig) {
private static SpannerAccessor createAndConnect(SpannerConfig spannerConfig) {
SpannerOptions.Builder builder = SpannerOptions.newBuilder();
- ValueProvider commitDeadline = spannerConfig.getCommitDeadline();
- if (commitDeadline != null && commitDeadline.get().getMillis() > 0) {
+ // Set retryable codes for all API methods
+ if (spannerConfig.getRetryableCodes() != null) {
+ builder
+ .getSpannerStubSettingsBuilder()
+ .applyToAllUnaryMethods(
+ input -> {
+ input.setRetryableCodes(spannerConfig.getRetryableCodes());
+ return null;
+ });
+ builder
+ .getSpannerStubSettingsBuilder()
+ .executeStreamingSqlSettings()
+ .setRetryableCodes(spannerConfig.getRetryableCodes());
+ }
+ // Set commit retry settings
+ UnaryCallSettings.Builder commitSettings =
+ builder.getSpannerStubSettingsBuilder().commitSettings();
+ ValueProvider commitDeadline = spannerConfig.getCommitDeadline();
+ if (spannerConfig.getCommitRetrySettings() != null) {
+ commitSettings.setRetrySettings(spannerConfig.getCommitRetrySettings());
+ } else if (commitDeadline != null && commitDeadline.get().getMillis() > 0) {
// Set the GRPC deadline on the Commit API call.
- UnaryCallSettings.Builder commitSettings =
- builder.getSpannerStubSettingsBuilder().commitSettings();
- RetrySettings.Builder commitRetrySettings = commitSettings.getRetrySettings().toBuilder();
+ RetrySettings.Builder commitRetrySettingsBuilder =
+ commitSettings.getRetrySettings().toBuilder();
commitSettings.setRetrySettings(
- commitRetrySettings
+ commitRetrySettingsBuilder
.setTotalTimeout(org.threeten.bp.Duration.ofMillis(commitDeadline.get().getMillis()))
.setMaxRpcTimeout(org.threeten.bp.Duration.ofMillis(commitDeadline.get().getMillis()))
.setInitialRpcTimeout(
org.threeten.bp.Duration.ofMillis(commitDeadline.get().getMillis()))
.build());
}
- // Setting the timeout for streaming read to 2 hours. This is 1 hour by default
- // after BEAM 2.20.
+
+ // Set execute streaming sql retry settings
ServerStreamingCallSettings.Builder
executeStreamingSqlSettings =
builder.getSpannerStubSettingsBuilder().executeStreamingSqlSettings();
- RetrySettings.Builder executeSqlStreamingRetrySettings =
- executeStreamingSqlSettings.getRetrySettings().toBuilder();
- executeStreamingSqlSettings.setRetrySettings(
- executeSqlStreamingRetrySettings
- .setInitialRpcTimeout(org.threeten.bp.Duration.ofMinutes(120))
- .setMaxRpcTimeout(org.threeten.bp.Duration.ofMinutes(120))
- .setTotalTimeout(org.threeten.bp.Duration.ofMinutes(120))
- .build());
+ if (spannerConfig.getExecuteStreamingSqlRetrySettings() != null) {
+ executeStreamingSqlSettings.setRetrySettings(
+ spannerConfig.getExecuteStreamingSqlRetrySettings());
+ } else {
+ // Setting the timeout for streaming read to 2 hours. This is 1 hour by default
+ // after BEAM 2.20.
+ RetrySettings.Builder executeSqlStreamingRetrySettings =
+ executeStreamingSqlSettings.getRetrySettings().toBuilder();
+ executeStreamingSqlSettings.setRetrySettings(
+ executeSqlStreamingRetrySettings
+ .setInitialRpcTimeout(org.threeten.bp.Duration.ofMinutes(120))
+ .setMaxRpcTimeout(org.threeten.bp.Duration.ofMinutes(120))
+ .setTotalTimeout(org.threeten.bp.Duration.ofMinutes(120))
+ .build());
+ }
ValueProvider projectId = spannerConfig.getProjectId();
if (projectId != null) {
@@ -151,6 +177,10 @@ private static SpannerAccessor createAndConnect(SpannerConfig spannerConfig) {
ValueProvider emulatorHost = spannerConfig.getEmulatorHost();
if (emulatorHost != null) {
builder.setEmulatorHost(emulatorHost.get());
+ if (spannerConfig.getIsLocalChannelProvider() != null
+ && spannerConfig.getIsLocalChannelProvider().get()) {
+ builder.setChannelProvider(LocalChannelProvider.create(emulatorHost.get()));
+ }
builder.setCredentials(NoCredentials.getInstance());
}
String userAgentString = USER_AGENT_PREFIX + "/" + ReleaseInfo.getReleaseInfo().getVersion();
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
index 7fcc0b5676b92..a00b7896c35af 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
@@ -19,6 +19,8 @@
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.rpc.StatusCode.Code;
import com.google.auto.value.AutoValue;
import com.google.cloud.ServiceFactory;
import com.google.cloud.spanner.Options.RpcPriority;
@@ -28,6 +30,7 @@
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Duration;
@@ -56,10 +59,18 @@ public abstract class SpannerConfig implements Serializable {
public abstract @Nullable ValueProvider getEmulatorHost();
+ public abstract @Nullable ValueProvider getIsLocalChannelProvider();
+
public abstract @Nullable ValueProvider getCommitDeadline();
public abstract @Nullable ValueProvider getMaxCumulativeBackoff();
+ public abstract @Nullable RetrySettings getExecuteStreamingSqlRetrySettings();
+
+ public abstract @Nullable RetrySettings getCommitRetrySettings();
+
+ public abstract @Nullable ImmutableSet getRetryableCodes();
+
public abstract @Nullable ValueProvider getRpcPriority();
@VisibleForTesting
@@ -117,10 +128,19 @@ public abstract static class Builder {
abstract Builder setEmulatorHost(ValueProvider emulatorHost);
+ abstract Builder setIsLocalChannelProvider(ValueProvider isLocalChannelProvider);
+
abstract Builder setCommitDeadline(ValueProvider commitDeadline);
abstract Builder setMaxCumulativeBackoff(ValueProvider maxCumulativeBackoff);
+ abstract Builder setExecuteStreamingSqlRetrySettings(
+ RetrySettings executeStreamingSqlRetrySettings);
+
+ abstract Builder setCommitRetrySettings(RetrySettings commitRetrySettings);
+
+ abstract Builder setRetryableCodes(ImmutableSet retryableCodes);
+
abstract Builder setServiceFactory(ServiceFactory serviceFactory);
abstract Builder setRpcPriority(ValueProvider rpcPriority);
@@ -160,6 +180,10 @@ public SpannerConfig withEmulatorHost(ValueProvider emulatorHost) {
return toBuilder().setEmulatorHost(emulatorHost).build();
}
+ public SpannerConfig withIsLocalChannelProvider(ValueProvider isLocalChannelProvider) {
+ return toBuilder().setIsLocalChannelProvider(isLocalChannelProvider).build();
+ }
+
public SpannerConfig withCommitDeadline(Duration commitDeadline) {
return withCommitDeadline(ValueProvider.StaticValueProvider.of(commitDeadline));
}
@@ -176,6 +200,21 @@ public SpannerConfig withMaxCumulativeBackoff(ValueProvider maxCumulat
return toBuilder().setMaxCumulativeBackoff(maxCumulativeBackoff).build();
}
+ public SpannerConfig withExecuteStreamingSqlRetrySettings(
+ RetrySettings executeStreamingSqlRetrySettings) {
+ return toBuilder()
+ .setExecuteStreamingSqlRetrySettings(executeStreamingSqlRetrySettings)
+ .build();
+ }
+
+ public SpannerConfig withCommitRetrySettings(RetrySettings commitRetrySettings) {
+ return toBuilder().setCommitRetrySettings(commitRetrySettings).build();
+ }
+
+ public SpannerConfig withRetryableCodes(ImmutableSet retryableCodes) {
+ return toBuilder().setRetryableCodes(retryableCodes).build();
+ }
+
@VisibleForTesting
SpannerConfig withServiceFactory(ServiceFactory serviceFactory) {
return toBuilder().setServiceFactory(serviceFactory).build();
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
index 9fd929d097132..863d88ab54e68 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
@@ -23,6 +23,8 @@
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.rpc.StatusCode.Code;
import com.google.auto.value.AutoValue;
import com.google.cloud.ServiceFactory;
import com.google.cloud.Timestamp;
@@ -79,6 +81,7 @@
import org.apache.beam.sdk.metrics.Distribution;
import org.apache.beam.sdk.metrics.Metrics;
import org.apache.beam.sdk.options.ValueProvider;
+import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
@@ -115,6 +118,7 @@
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Stopwatch;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.primitives.UnsignedBytes;
import org.checkerframework.checker.nullness.qual.Nullable;
@@ -1349,6 +1353,8 @@ public abstract static class ReadChangeStream
abstract @Nullable String getMetadataDatabase();
+ abstract @Nullable String getMetadataTable();
+
abstract Timestamp getInclusiveStartAt();
abstract @Nullable Timestamp getInclusiveEndAt();
@@ -1370,6 +1376,8 @@ abstract static class Builder {
abstract Builder setMetadataDatabase(String metadataDatabase);
+ abstract Builder setMetadataTable(String metadataTable);
+
abstract Builder setInclusiveStartAt(Timestamp inclusiveStartAt);
abstract Builder setInclusiveEndAt(Timestamp inclusiveEndAt);
@@ -1434,6 +1442,10 @@ public ReadChangeStream withMetadataDatabase(String metadataDatabase) {
return toBuilder().setMetadataDatabase(metadataDatabase).build();
}
+ public ReadChangeStream withMetadataTable(String metadataTable) {
+ return toBuilder().setMetadataTable(metadataTable).build();
+ }
+
/** Specifies the time that the change stream should be read from. */
public ReadChangeStream withInclusiveStartAt(Timestamp timestamp) {
return toBuilder().setInclusiveStartAt(timestamp).build();
@@ -1497,7 +1509,8 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta
final String partitionMetadataDatabaseId =
MoreObjects.firstNonNull(getMetadataDatabase(), changeStreamDatabaseId.getDatabase());
final String partitionMetadataTableName =
- generatePartitionMetadataTableName(partitionMetadataDatabaseId);
+ MoreObjects.firstNonNull(
+ getMetadataTable(), generatePartitionMetadataTableName(partitionMetadataDatabaseId));
if (getTraceSampleProbability() != null) {
TraceConfig globalTraceConfig = Tracing.getTraceConfig();
@@ -1511,16 +1524,36 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta
.spanBuilder("SpannerIO.ReadChangeStream.expand")
.setRecordEvents(true)
.startScopedSpan()) {
- final SpannerConfig changeStreamSpannerConfig = getSpannerConfig();
+ SpannerConfig changeStreamSpannerConfig = getSpannerConfig();
+ // Set default retryable errors for ReadChangeStream
+ if (changeStreamSpannerConfig.getRetryableCodes() == null) {
+ ImmutableSet defaultRetryableCodes =
+ ImmutableSet.of(Code.UNAVAILABLE, Code.ABORTED);
+ changeStreamSpannerConfig =
+ changeStreamSpannerConfig
+ .toBuilder()
+ .setRetryableCodes(defaultRetryableCodes)
+ .build();
+ }
+ // Set default retry timeouts for ReadChangeStream
+ if (changeStreamSpannerConfig.getExecuteStreamingSqlRetrySettings() == null) {
+ changeStreamSpannerConfig =
+ changeStreamSpannerConfig
+ .toBuilder()
+ .setExecuteStreamingSqlRetrySettings(
+ RetrySettings.newBuilder()
+ .setTotalTimeout(org.threeten.bp.Duration.ofMinutes(5))
+ .setInitialRpcTimeout(org.threeten.bp.Duration.ofMinutes(1))
+ .setMaxRpcTimeout(org.threeten.bp.Duration.ofMinutes(1))
+ .build())
+ .build();
+ }
final SpannerConfig partitionMetadataSpannerConfig =
- SpannerConfig.create()
- .withProjectId(changeStreamSpannerConfig.getProjectId())
- .withHost(changeStreamSpannerConfig.getHost())
- .withInstanceId(partitionMetadataInstanceId)
- .withDatabaseId(partitionMetadataDatabaseId)
- .withCommitDeadline(changeStreamSpannerConfig.getCommitDeadline())
- .withEmulatorHost(changeStreamSpannerConfig.getEmulatorHost())
- .withMaxCumulativeBackoff(changeStreamSpannerConfig.getMaxCumulativeBackoff());
+ changeStreamSpannerConfig
+ .toBuilder()
+ .setInstanceId(StaticValueProvider.of(partitionMetadataInstanceId))
+ .setDatabaseId(StaticValueProvider.of(partitionMetadataDatabaseId))
+ .build();
final String changeStreamName = getChangeStreamName();
// FIXME: The backend only supports microsecond granularity. Remove when fixed.
final Timestamp startTimestamp = TimestampConverter.truncateNanos(getInclusiveStartAt());
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOStorageQueryTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOStorageQueryTest.java
index 86f538a4d0f3b..b23de6b3fc334 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOStorageQueryTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOStorageQueryTest.java
@@ -93,7 +93,10 @@
import org.junit.runners.JUnit4;
import org.junit.runners.model.Statement;
-/** Tests for {@link BigQueryIO#readTableRows()} using {@link Method#DIRECT_READ}. */
+/**
+ * Tests for {@link BigQueryIO#read(SerializableFunction)} using {@link Method#DIRECT_READ} and
+ * {@link BigQueryIO.TypedRead#fromQuery(String)}.
+ */
@RunWith(JUnit4.class)
public class BigQueryIOStorageQueryTest {
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOStorageReadTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOStorageReadTest.java
index fcf335c01b309..8f8ab867a9101 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOStorageReadTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOStorageReadTest.java
@@ -25,6 +25,7 @@
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
@@ -92,20 +93,29 @@
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead.Method;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.StorageClient;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryUtils.ConversionOptions;
import org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices;
import org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices.FakeBigQueryServerStream;
import org.apache.beam.sdk.io.gcp.testing.FakeDatasetService;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+import org.apache.beam.sdk.schemas.FieldAccessDescriptor;
+import org.apache.beam.sdk.schemas.SchemaCoder;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.display.DisplayData;
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.Row;
+import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.hamcrest.Matchers;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
@@ -2064,6 +2074,56 @@ public void testStreamSourceSplitAtFractionFailsWhenParentIsPastSplitPointArrow(
assertFalse(parent.advance());
}
+ @Test
+ public void testActuateProjectionPushdown() {
+ org.apache.beam.sdk.schemas.Schema schema =
+ org.apache.beam.sdk.schemas.Schema.builder()
+ .addStringField("foo")
+ .addStringField("bar")
+ .build();
+ TypedRead read =
+ BigQueryIO.read(
+ record ->
+ BigQueryUtils.toBeamRow(
+ record.getRecord(), schema, ConversionOptions.builder().build()))
+ .withMethod(Method.DIRECT_READ)
+ .withCoder(SchemaCoder.of(schema));
+
+ assertTrue(read.supportsProjectionPushdown());
+ PTransform> pushdownT =
+ read.actuateProjectionPushdown(
+ ImmutableMap.of(new TupleTag<>("output"), FieldAccessDescriptor.withFieldNames("foo")));
+
+ TypedRead pushdownRead = (TypedRead) pushdownT;
+ assertEquals(Method.DIRECT_READ, pushdownRead.getMethod());
+ assertThat(pushdownRead.getSelectedFields().get(), Matchers.containsInAnyOrder("foo"));
+ }
+
+ @Test
+ public void testReadFromQueryDoesNotSupportProjectionPushdown() {
+ org.apache.beam.sdk.schemas.Schema schema =
+ org.apache.beam.sdk.schemas.Schema.builder()
+ .addStringField("foo")
+ .addStringField("bar")
+ .build();
+ TypedRead read =
+ BigQueryIO.read(
+ record ->
+ BigQueryUtils.toBeamRow(
+ record.getRecord(), schema, ConversionOptions.builder().build()))
+ .fromQuery("SELECT bar FROM `dataset.table`")
+ .withMethod(Method.DIRECT_READ)
+ .withCoder(SchemaCoder.of(schema));
+
+ assertFalse(read.supportsProjectionPushdown());
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ read.actuateProjectionPushdown(
+ ImmutableMap.of(
+ new TupleTag<>("output"), FieldAccessDescriptor.withFieldNames("foo"))));
+ }
+
private static org.apache.arrow.vector.types.pojo.Field field(
String name,
boolean nullable,
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java
new file mode 100644
index 0000000000000..34d7cd4a91a7d
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java
@@ -0,0 +1,497 @@
+/*
+ * 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.spanner.changestreams;
+
+import static org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataAdminDao.COLUMN_CREATED_AT;
+import static org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataAdminDao.COLUMN_END_TIMESTAMP;
+import static org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataAdminDao.COLUMN_FINISHED_AT;
+import static org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataAdminDao.COLUMN_HEARTBEAT_MILLIS;
+import static org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataAdminDao.COLUMN_PARENT_TOKENS;
+import static org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataAdminDao.COLUMN_PARTITION_TOKEN;
+import static org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataAdminDao.COLUMN_RUNNING_AT;
+import static org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataAdminDao.COLUMN_SCHEDULED_AT;
+import static org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataAdminDao.COLUMN_START_TIMESTAMP;
+import static org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataAdminDao.COLUMN_STATE;
+import static org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataAdminDao.COLUMN_WATERMARK;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import com.google.api.gax.grpc.testing.MockServiceHelper;
+import com.google.api.gax.retrying.RetrySettings;
+import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.ErrorCode;
+import com.google.cloud.spanner.MockSpannerServiceImpl;
+import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
+import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.Statement;
+import com.google.protobuf.ListValue;
+import com.google.protobuf.NullValue;
+import com.google.protobuf.Value;
+import com.google.spanner.v1.ExecuteSqlRequest;
+import com.google.spanner.v1.ResultSet;
+import com.google.spanner.v1.ResultSetMetadata;
+import com.google.spanner.v1.StructType;
+import com.google.spanner.v1.StructType.Field;
+import com.google.spanner.v1.Type;
+import com.google.spanner.v1.TypeCode;
+import io.grpc.Status;
+import java.io.Serializable;
+import java.util.Collections;
+import org.apache.beam.sdk.Pipeline.PipelineExecutionException;
+import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig;
+import org.apache.beam.sdk.io.gcp.spanner.SpannerIO;
+import org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.DaoFactory;
+import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata.State;
+import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.hamcrest.Matchers;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class SpannerChangeStreamErrorTest implements Serializable {
+
+ public static final String SPANNER_HOST = "my-host";
+ private static final String TEST_PROJECT = "my-project";
+ private static final String TEST_INSTANCE = "my-instance";
+ private static final String TEST_DATABASE = "my-database";
+ private static final String TEST_TABLE = "my-metadata-table";
+ private static final String TEST_CHANGE_STREAM = "my-change-stream";
+
+ @Rule
+ public final transient TestPipeline pipeline =
+ TestPipeline.create().enableAbandonedNodeEnforcement(false);
+
+ @Rule public final transient ExpectedException thrown = ExpectedException.none();
+
+ private MockSpannerServiceImpl mockSpannerService;
+ private MockServiceHelper serviceHelper;
+
+ @Before
+ public void setUp() throws Exception {
+ mockSpannerService = new MockSpannerServiceImpl();
+ serviceHelper =
+ new MockServiceHelper(SPANNER_HOST, Collections.singletonList(mockSpannerService));
+ serviceHelper.start();
+ serviceHelper.reset();
+ }
+
+ @After
+ public void tearDown() throws NoSuchFieldException, IllegalAccessException {
+ serviceHelper.reset();
+ serviceHelper.stop();
+ mockSpannerService.reset();
+ resetDaoFactoryFields();
+ }
+
+ @Test
+ public void testResourceExhaustedDoesNotRetry() {
+ mockSpannerService.setExecuteStreamingSqlExecutionTime(
+ SimulatedExecutionTime.ofStickyException(Status.RESOURCE_EXHAUSTED.asRuntimeException()));
+
+ final Timestamp now = Timestamp.now();
+ final Timestamp after3Seconds =
+ Timestamp.ofTimeSecondsAndNanos(now.getSeconds() + 3, now.getNanos());
+ try {
+ pipeline.apply(
+ SpannerIO.readChangeStream()
+ .withSpannerConfig(getSpannerConfig())
+ .withChangeStreamName(TEST_CHANGE_STREAM)
+ .withMetadataDatabase(TEST_DATABASE)
+ .withInclusiveStartAt(now)
+ .withInclusiveEndAt(after3Seconds));
+ pipeline.run().waitUntilFinish();
+ } finally {
+ thrown.expect(PipelineExecutionException.class);
+ thrown.expectMessage(ErrorCode.RESOURCE_EXHAUSTED.name());
+ }
+ }
+
+ @Test
+ public void testUnavailableExceptionRetries() {
+ mockSpannerService.setExecuteStreamingSqlExecutionTime(
+ SimulatedExecutionTime.ofExceptions(
+ ImmutableSet.of(
+ Status.UNAVAILABLE.asRuntimeException(),
+ Status.RESOURCE_EXHAUSTED.asRuntimeException())));
+
+ final Timestamp now = Timestamp.now();
+ final Timestamp after3Seconds =
+ Timestamp.ofTimeSecondsAndNanos(now.getSeconds() + 3, now.getNanos());
+ try {
+ pipeline.apply(
+ SpannerIO.readChangeStream()
+ .withSpannerConfig(getSpannerConfig())
+ .withChangeStreamName(TEST_CHANGE_STREAM)
+ .withMetadataDatabase(TEST_DATABASE)
+ .withInclusiveStartAt(now)
+ .withInclusiveEndAt(after3Seconds));
+ pipeline.run().waitUntilFinish();
+ } finally {
+ assertThat(
+ mockSpannerService.countRequestsOfType(ExecuteSqlRequest.class), Matchers.greaterThan(1));
+ thrown.expect(PipelineExecutionException.class);
+ thrown.expectMessage(ErrorCode.RESOURCE_EXHAUSTED.name());
+ }
+ }
+
+ @Test
+ public void testAbortedExceptionRetries() {
+ mockSpannerService.setExecuteStreamingSqlExecutionTime(
+ SimulatedExecutionTime.ofStickyException(Status.ABORTED.asRuntimeException()));
+
+ final Timestamp now = Timestamp.now();
+ final Timestamp after3Seconds =
+ Timestamp.ofTimeSecondsAndNanos(now.getSeconds() + 3, now.getNanos());
+ try {
+ pipeline.apply(
+ SpannerIO.readChangeStream()
+ .withSpannerConfig(getSpannerConfig())
+ .withChangeStreamName(TEST_CHANGE_STREAM)
+ .withMetadataDatabase(TEST_DATABASE)
+ .withInclusiveStartAt(now)
+ .withInclusiveEndAt(after3Seconds));
+ pipeline.run().waitUntilFinish();
+ } finally {
+ assertThat(
+ mockSpannerService.countRequestsOfType(ExecuteSqlRequest.class), Matchers.greaterThan(1));
+ thrown.expect(PipelineExecutionException.class);
+ thrown.expectMessage(ErrorCode.ABORTED.name());
+ }
+ }
+
+ @Test
+ public void testUnknownExceptionDoesNotRetry() {
+ mockSpannerService.setExecuteStreamingSqlExecutionTime(
+ SimulatedExecutionTime.ofStickyException(Status.UNKNOWN.asRuntimeException()));
+
+ final Timestamp now = Timestamp.now();
+ final Timestamp after3Seconds =
+ Timestamp.ofTimeSecondsAndNanos(now.getSeconds() + 3, now.getNanos());
+ try {
+ pipeline.apply(
+ SpannerIO.readChangeStream()
+ .withSpannerConfig(getSpannerConfig())
+ .withChangeStreamName(TEST_CHANGE_STREAM)
+ .withMetadataDatabase(TEST_DATABASE)
+ .withInclusiveStartAt(now)
+ .withInclusiveEndAt(after3Seconds));
+ pipeline.run().waitUntilFinish();
+ } finally {
+ assertThat(
+ mockSpannerService.countRequestsOfType(ExecuteSqlRequest.class), Matchers.equalTo(1));
+ thrown.expect(PipelineExecutionException.class);
+ thrown.expectMessage(ErrorCode.UNKNOWN.name());
+ }
+ }
+
+ @Test
+ public void testInvalidRecordReceived() {
+ final Timestamp now = Timestamp.now();
+ final Timestamp after3Seconds =
+ Timestamp.ofTimeSecondsAndNanos(now.getSeconds() + 3, now.getNanos());
+
+ mockTableExists();
+ ResultSet getPartitionResultSet = mockGetParentPartition(now, after3Seconds);
+ mockGetWatermark(now);
+ mockGetPartitionsAfter(
+ Timestamp.ofTimeSecondsAndNanos(now.getSeconds(), now.getNanos() + 1000),
+ getPartitionResultSet);
+ mockGetPartitionsAfter(
+ Timestamp.ofTimeSecondsAndNanos(now.getSeconds(), now.getNanos() - 1000),
+ getPartitionResultSet);
+ mockInvalidChangeStreamRecordReceived(now, after3Seconds);
+
+ try {
+ pipeline.apply(
+ SpannerIO.readChangeStream()
+ .withSpannerConfig(getSpannerConfig())
+ .withChangeStreamName(TEST_CHANGE_STREAM)
+ .withMetadataDatabase(TEST_DATABASE)
+ .withMetadataTable(TEST_TABLE)
+ .withInclusiveStartAt(now)
+ .withInclusiveEndAt(after3Seconds));
+ pipeline.run().waitUntilFinish();
+ } finally {
+ thrown.expect(PipelineExecutionException.class);
+ thrown.expectMessage("Field not found");
+ }
+ }
+
+ private void mockInvalidChangeStreamRecordReceived(Timestamp now, Timestamp after3Seconds) {
+ Statement changeStreamQueryStatement =
+ Statement.newBuilder(
+ "SELECT * FROM READ_my-change-stream( start_timestamp => @startTimestamp, end_timestamp => @endTimestamp, partition_token => @partitionToken, read_options => null, heartbeat_milliseconds => @heartbeatMillis)")
+ .bind("startTimestamp")
+ .to(now)
+ .bind("endTimestamp")
+ .to(after3Seconds)
+ .bind("partitionToken")
+ .to((String) null)
+ .bind("heartbeatMillis")
+ .to(500)
+ .build();
+ ResultSetMetadata readChangeStreamResultSetMetadata =
+ ResultSetMetadata.newBuilder()
+ .setRowType(
+ StructType.newBuilder()
+ .addFields(
+ Field.newBuilder()
+ .setName("COL1")
+ .setType(
+ Type.newBuilder()
+ .setCode(TypeCode.ARRAY)
+ .setArrayElementType(
+ Type.newBuilder()
+ .setCode(TypeCode.STRUCT)
+ .setStructType(
+ StructType.newBuilder()
+ .addFields(
+ Field.newBuilder()
+ .setName("field_name")
+ .setType(
+ Type.newBuilder()
+ .setCode(TypeCode.STRUCT)
+ .setStructType(
+ StructType.newBuilder()
+ .addFields(
+ Field.newBuilder()
+ .setType(
+ Type
+ .newBuilder()
+ .setCode(
+ TypeCode
+ .STRING)))))))))))
+ .build();
+ ResultSet readChangeStreamResultSet =
+ ResultSet.newBuilder()
+ .addRows(
+ ListValue.newBuilder()
+ .addValues(
+ Value.newBuilder()
+ .setListValue(
+ ListValue.newBuilder()
+ .addValues(
+ Value.newBuilder()
+ .setListValue(
+ ListValue.newBuilder()
+ .addValues(
+ Value.newBuilder()
+ .setListValue(
+ ListValue.newBuilder()
+ .addValues(
+ Value.newBuilder()
+ .setStringValue(
+ "bad_value")))))))))
+ .setMetadata(readChangeStreamResultSetMetadata)
+ .build();
+ mockSpannerService.putStatementResult(
+ StatementResult.query(changeStreamQueryStatement, readChangeStreamResultSet));
+ }
+
+ private void mockGetPartitionsAfter(Timestamp timestamp, ResultSet getPartitionResultSet) {
+ Statement getPartitionsAfterStatement =
+ Statement.newBuilder(
+ "SELECT * FROM my-metadata-table WHERE CreatedAt > @timestamp ORDER BY CreatedAt ASC, StartTimestamp ASC")
+ .bind("timestamp")
+ .to(Timestamp.ofTimeSecondsAndNanos(timestamp.getSeconds(), timestamp.getNanos()))
+ .build();
+ mockSpannerService.putStatementResult(
+ StatementResult.query(getPartitionsAfterStatement, getPartitionResultSet));
+ }
+
+ private void mockGetWatermark(Timestamp now) {
+ Statement watermarkStatement =
+ Statement.newBuilder(
+ "SELECT Watermark FROM my-metadata-table WHERE State != @state ORDER BY Watermark ASC LIMIT 1")
+ .bind("state")
+ .to(State.FINISHED.name())
+ .build();
+ ResultSetMetadata watermarkResultSetMetadata =
+ ResultSetMetadata.newBuilder()
+ .setRowType(
+ StructType.newBuilder()
+ .addFields(
+ Field.newBuilder()
+ .setName("Watermark")
+ .setType(Type.newBuilder().setCode(TypeCode.TIMESTAMP).build())
+ .build())
+ .build())
+ .build();
+ ResultSet watermarkResultSet =
+ ResultSet.newBuilder()
+ .addRows(
+ ListValue.newBuilder()
+ .addValues(Value.newBuilder().setStringValue(now.toString()).build())
+ .build())
+ .setMetadata(watermarkResultSetMetadata)
+ .build();
+ mockSpannerService.putStatementResult(
+ StatementResult.query(watermarkStatement, watermarkResultSet));
+ }
+
+ private ResultSet mockGetParentPartition(Timestamp now, Timestamp after3Seconds) {
+ Statement getPartitionStatement =
+ Statement.newBuilder("SELECT * FROM my-metadata-table WHERE PartitionToken = @partition")
+ .bind("partition")
+ .to("Parent0")
+ .build();
+ ResultSet getPartitionResultSet =
+ ResultSet.newBuilder()
+ .addRows(
+ ListValue.newBuilder()
+ .addValues(Value.newBuilder().setStringValue("Parent0"))
+ .addValues(Value.newBuilder().setListValue(ListValue.newBuilder().build()))
+ .addValues(Value.newBuilder().setStringValue(now.toString()))
+ .addValues(Value.newBuilder().setStringValue(after3Seconds.toString()))
+ .addValues(Value.newBuilder().setStringValue("500"))
+ .addValues(Value.newBuilder().setStringValue(State.CREATED.name()))
+ .addValues(Value.newBuilder().setStringValue(now.toString()))
+ .addValues(Value.newBuilder().setStringValue(now.toString()))
+ .addValues(Value.newBuilder().setNullValue(NullValue.NULL_VALUE).build())
+ .addValues(Value.newBuilder().setNullValue(NullValue.NULL_VALUE).build())
+ .addValues(Value.newBuilder().setNullValue(NullValue.NULL_VALUE).build())
+ .build())
+ .setMetadata(PARTITION_METADATA_RESULT_SET_METADATA)
+ .build();
+ mockSpannerService.putStatementResult(
+ StatementResult.query(getPartitionStatement, getPartitionResultSet));
+ return getPartitionResultSet;
+ }
+
+ private void mockTableExists() {
+ Statement tableExistsStatement =
+ Statement.of(
+ "SELECT t.table_name FROM information_schema.tables AS t WHERE t.table_catalog = '' AND t.table_schema = '' AND t.table_name = 'my-metadata-table'");
+ ResultSetMetadata tableExistsResultSetMetadata =
+ ResultSetMetadata.newBuilder()
+ .setRowType(
+ StructType.newBuilder()
+ .addFields(
+ Field.newBuilder()
+ .setName("table_name")
+ .setType(Type.newBuilder().setCode(TypeCode.STRING).build())
+ .build())
+ .build())
+ .build();
+ ResultSet tableExistsResultSet =
+ ResultSet.newBuilder()
+ .addRows(
+ ListValue.newBuilder()
+ .addValues(Value.newBuilder().setStringValue(TEST_TABLE).build())
+ .build())
+ .setMetadata(tableExistsResultSetMetadata)
+ .build();
+ mockSpannerService.putStatementResult(
+ StatementResult.query(tableExistsStatement, tableExistsResultSet));
+ }
+
+ private SpannerConfig getSpannerConfig() {
+ RetrySettings quickRetrySettings =
+ RetrySettings.newBuilder()
+ .setInitialRetryDelay(org.threeten.bp.Duration.ofMillis(250))
+ .setMaxRetryDelay(org.threeten.bp.Duration.ofSeconds(1))
+ .setRetryDelayMultiplier(5)
+ .setTotalTimeout(org.threeten.bp.Duration.ofSeconds(1))
+ .build();
+ return SpannerConfig.create()
+ .withEmulatorHost(StaticValueProvider.of(SPANNER_HOST))
+ .withIsLocalChannelProvider(StaticValueProvider.of(true))
+ .withCommitRetrySettings(quickRetrySettings)
+ .withExecuteStreamingSqlRetrySettings(quickRetrySettings)
+ .withProjectId(TEST_PROJECT)
+ .withInstanceId(TEST_INSTANCE)
+ .withDatabaseId(TEST_DATABASE);
+ }
+
+ private static void resetDaoFactoryFields() throws NoSuchFieldException, IllegalAccessException {
+ java.lang.reflect.Field partitionMetadataAdminDaoField =
+ DaoFactory.class.getDeclaredField("partitionMetadataAdminDao");
+ partitionMetadataAdminDaoField.setAccessible(true);
+ partitionMetadataAdminDaoField.set(null, null);
+ java.lang.reflect.Field partitionMetadataDaoInstanceField =
+ DaoFactory.class.getDeclaredField("partitionMetadataDaoInstance");
+ partitionMetadataDaoInstanceField.setAccessible(true);
+ partitionMetadataDaoInstanceField.set(null, null);
+ java.lang.reflect.Field changeStreamDaoInstanceField =
+ DaoFactory.class.getDeclaredField("changeStreamDaoInstance");
+ changeStreamDaoInstanceField.setAccessible(true);
+ changeStreamDaoInstanceField.set(null, null);
+ }
+
+ private static final ResultSetMetadata PARTITION_METADATA_RESULT_SET_METADATA =
+ ResultSetMetadata.newBuilder()
+ .setRowType(
+ StructType.newBuilder()
+ .addFields(
+ Field.newBuilder()
+ .setName(COLUMN_PARTITION_TOKEN)
+ .setType(Type.newBuilder().setCode(TypeCode.STRING))
+ .build())
+ .addFields(
+ Field.newBuilder()
+ .setName(COLUMN_PARENT_TOKENS)
+ .setType(
+ Type.newBuilder()
+ .setCode(TypeCode.ARRAY)
+ .setArrayElementType(Type.newBuilder().setCode(TypeCode.STRING)))
+ .build())
+ .addFields(
+ Field.newBuilder()
+ .setName(COLUMN_START_TIMESTAMP)
+ .setType(Type.newBuilder().setCode(TypeCode.TIMESTAMP)))
+ .addFields(
+ Field.newBuilder()
+ .setName(COLUMN_END_TIMESTAMP)
+ .setType(Type.newBuilder().setCode(TypeCode.TIMESTAMP)))
+ .addFields(
+ Field.newBuilder()
+ .setName(COLUMN_HEARTBEAT_MILLIS)
+ .setType(Type.newBuilder().setCode(TypeCode.INT64)))
+ .addFields(
+ Field.newBuilder()
+ .setName(COLUMN_STATE)
+ .setType(Type.newBuilder().setCode(TypeCode.STRING)))
+ .addFields(
+ Field.newBuilder()
+ .setName(COLUMN_WATERMARK)
+ .setType(Type.newBuilder().setCode(TypeCode.TIMESTAMP)))
+ .addFields(
+ Field.newBuilder()
+ .setName(COLUMN_CREATED_AT)
+ .setType(Type.newBuilder().setCode(TypeCode.TIMESTAMP)))
+ .addFields(
+ Field.newBuilder()
+ .setName(COLUMN_SCHEDULED_AT)
+ .setType(Type.newBuilder().setCode(TypeCode.TIMESTAMP)))
+ .addFields(
+ Field.newBuilder()
+ .setName(COLUMN_RUNNING_AT)
+ .setType(Type.newBuilder().setCode(TypeCode.TIMESTAMP)))
+ .addFields(
+ Field.newBuilder()
+ .setName(COLUMN_FINISHED_AT)
+ .setType(Type.newBuilder().setCode(TypeCode.TIMESTAMP)))
+ .build())
+ .build();
+}
diff --git a/sdks/java/io/neo4j/OWNERS b/sdks/java/io/neo4j/OWNERS
new file mode 100644
index 0000000000000..0ff6e82359fb6
--- /dev/null
+++ b/sdks/java/io/neo4j/OWNERS
@@ -0,0 +1,5 @@
+# See the OWNERS docs at https://s.apache.org/beam-owners
+
+reviewers:
+ - mcasters
+
diff --git a/sdks/java/io/neo4j/build.gradle b/sdks/java/io/neo4j/build.gradle
new file mode 100644
index 0000000000000..9d5adfc32b1d0
--- /dev/null
+++ b/sdks/java/io/neo4j/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.neo4j')
+provideIntegrationTestingDependencies()
+enableJavaPerformanceTesting()
+
+description = "Apache Beam :: SDKs :: Java :: IO :: Neo4j"
+ext.summary = "IO to read from and write to Neo4j graphs"
+
+dependencies {
+ implementation project(path: ":sdks:java:core", configuration: "shadow")
+ implementation "org.neo4j.driver:neo4j-java-driver:4.4.3"
+ implementation library.java.slf4j_api
+ implementation library.java.vendored_guava_26_0_jre
+ testImplementation library.java.junit
+ testImplementation library.java.hamcrest
+ testImplementation project(path: ":sdks:java:io:common", configuration: "testRuntimeMigration")
+ testImplementation project(path: ":sdks:java:testing:test-utils", configuration: "testRuntimeMigration")
+ testImplementation "org.testcontainers:neo4j:1.16.2"
+ testRuntimeOnly library.java.slf4j_jdk14
+ testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow")
+}
diff --git a/sdks/java/io/neo4j/src/main/java/org/apache/beam/sdk/io/neo4j/Neo4jIO.java b/sdks/java/io/neo4j/src/main/java/org/apache/beam/sdk/io/neo4j/Neo4jIO.java
new file mode 100644
index 0000000000000..9b011aff410d2
--- /dev/null
+++ b/sdks/java/io/neo4j/src/main/java/org/apache/beam/sdk/io/neo4j/Neo4jIO.java
@@ -0,0 +1,1221 @@
+/*
+ * 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.neo4j;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.StringUtils;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.harness.JvmInitializer;
+import org.apache.beam.sdk.options.ValueProvider;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
+import org.checkerframework.checker.initialization.qual.Initialized;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.neo4j.driver.AuthToken;
+import org.neo4j.driver.AuthTokens;
+import org.neo4j.driver.Config;
+import org.neo4j.driver.Driver;
+import org.neo4j.driver.GraphDatabase;
+import org.neo4j.driver.Record;
+import org.neo4j.driver.Result;
+import org.neo4j.driver.Session;
+import org.neo4j.driver.SessionConfig;
+import org.neo4j.driver.TransactionConfig;
+import org.neo4j.driver.TransactionWork;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * This is a Beam IO to read from, and write data to, Neo4j.
+ *
+ *
+ *
+ *
+ *
+ *
Driver configuration
+ *
+ * To read from or write to Neo4j you have to provide a {@link DriverConfiguration} using
+ *
+ *
+ *
+ *
Neo4j Aura
+ *
+ * If you have trouble connecting to a Neo4j Aura database please try to disable a few security
+ * algorithms in your JVM. This makes sure that the right one is picked to connect:
+ *
+ *
+ *
+ *
{@code
+ * Security.setProperty(
+ * "jdk.tls.disabledAlgorithms",
+ * "SSLv3, RC4, DES, MD5withRSA, DH keySize < 1024, EC keySize < 224, 3DES_EDE_CBC, anon, NULL");
+ * }
+ *
+ *
+ *
+ *
+ *
+ *
To execute this code on GCP Dataflow you can create a class which extends {@link
+ * JvmInitializer} and implement the {@link JvmInitializer#onStartup()} method. You need to annotate
+ * this new class with {@link com.google.auto.service.AutoService}
+ *
+ *
{@code
+ * @AutoService(value = JvmInitializer.class)
+ * }
+ *
+ *
+ *
+ *
Reading from Neo4j
+ *
+ * {@link Neo4jIO#readAll()} source returns a bounded collection of {@code OuptutT} as a {@code
+ * PCollection}. OutputT is the type returned by the provided {@link RowMapper}. It accepts
+ * parameters as input in the form of {@code ParameterT} as a {@code PCollection}
+ *
+ * The following example reads ages to return the IDs of Person nodes. It runs a Cypher query for
+ * each provided age.
+ *
+ *
The mapping {@link SerializableFunction} maps input values to each execution of the Cypher
+ * statement. In the function simply return a map containing the parameters you want to set.
+ *
+ *
The {@link RowMapper} converts output Neo4j {@link Record} values to the output of the source.
+ *
+ *
{@code
+ * pipeline
+ * .apply(Create.of(40, 50, 60))
+ * .apply(Neo4jIO.readAll()
+ * .withDriverConfiguration(Neo4jIO.DriverConfiguration.create("neo4j://localhost:7687", "neo4j", "password"))
+ * .withCypher("MATCH(n:Person) WHERE n.age = $age RETURN n.id")
+ * .withReadTransaction()
+ * .withCoder(StringUtf8Coder.of())
+ * .withParametersFunction( age -> Collections.singletonMap( "age", age ))
+ * .withRowMapper( record -> return record.get(0).asString() )
+ * );
+ * }
+ *
+ * Writing to Neo4j
+ *
+ * The Neo4j {@link WriteUnwind} transform supports writing data to a graph. It writes a {@link
+ * PCollection} to the graph by collecting a batch of elements after which all elements in the batch
+ * are written together to Neo4j.
+ *
+ *
Like the source, to configure this sink, you have to provide a {@link DriverConfiguration}.
+ *
+ *
In the following example we'll merge a collection of {@link org.apache.beam.sdk.values.Row}
+ * into Person nodes. Since this is a Sink it has no output and as such no RowMapper is needed. The
+ * rows are being used as a container for the parameters of the Cypher statement. The used Cypher in
+ * question needs to be an UNWIND statement. Like in the read case, the parameters {@link
+ * SerializableFunction} converts parameter values to a {@link Map}. The difference here is that the
+ * resulting Map is stored in a {@link List} (containing maps) which in turn is stored in another
+ * Map under the name provided by the {@link WriteUnwind#withUnwindMapName(String)} method. All of
+ * this is handled automatically. You do need to provide the unwind map name so that you can
+ * reference that in the UNWIND statement.
+ *
+ *
+ *
+ *
For example:
+ *
+ *
{@code
+ * pipeline
+ * .apply(...)
+ * .apply(Neo4jIO.writeUnwind()
+ * .withDriverConfiguration(Neo4jIO.DriverConfiguration.create("neo4j://localhost:7687", "neo4j", "password"))
+ * .withUnwindMapName("rows")
+ * .withCypher("UNWIND $rows AS row MERGE(n:Person { id : row.id } ) SET n.firstName = row.first, n.lastName = row.last")
+ * .withParametersFunction( row -> ImmutableMap.of(
+ * "id", row.getString("id),
+ * "first", row.getString("firstName")
+ * "last", row.getString("lastName")))
+ * );
+ * }
+ */
+@Experimental(Experimental.Kind.SOURCE_SINK)
+public class Neo4jIO {
+
+ private static final Logger LOG = LoggerFactory.getLogger(Neo4jIO.class);
+
+ /**
+ * Read all rows using a Neo4j Cypher query.
+ *
+ * @param Type of the data representing query parameters.
+ * @param Type of the data to be read.
+ */
+ public static ReadAll readAll() {
+ return new AutoValue_Neo4jIO_ReadAll.Builder().build();
+ }
+
+ /**
+ * Write all rows using a Neo4j Cypher UNWIND cypher statement. This sets a default batch size of
+ * 5000.
+ *
+ * @param Type of the data representing query parameters.
+ */
+ public static WriteUnwind writeUnwind() {
+ return new AutoValue_Neo4jIO_WriteUnwind.Builder()
+ .setBatchSize(ValueProvider.StaticValueProvider.of(5000L))
+ .build();
+ }
+
+ private static PCollection getOutputPCollection(
+ PCollection input,
+ DoFn writeFn,
+ @Nullable Coder coder) {
+ PCollection output = input.apply(ParDo.of(writeFn));
+ if (coder != null) {
+ output.setCoder(coder);
+ try {
+ TypeDescriptor typeDesc = coder.getEncodedTypeDescriptor();
+ SchemaRegistry registry = input.getPipeline().getSchemaRegistry();
+ Schema schema = registry.getSchema(typeDesc);
+ output.setSchema(
+ schema,
+ typeDesc,
+ registry.getToRowFunction(typeDesc),
+ registry.getFromRowFunction(typeDesc));
+ } catch (NoSuchSchemaException e) {
+ // ignore
+ }
+ }
+ return output;
+ }
+
+ /**
+ * An interface used by {@link ReadAll} for converting each row of a Neo4j {@link Result} record
+ * {@link Record} into an element of the resulting {@link PCollection}.
+ */
+ @FunctionalInterface
+ public interface RowMapper extends Serializable {
+ T mapRow(Record record) throws Exception;
+ }
+
+ /**
+ * A convenience method to clarify the way {@link ValueProvider} works to the static code checker
+ * framework for {@link Nullable} values.
+ *
+ * @param valueProvider
+ * @param
+ * @return The provided value or null if none was specified.
+ */
+ private static T getProvidedValue(@Nullable ValueProvider valueProvider) {
+ if (valueProvider == null) {
+ return (T) null;
+ }
+ return valueProvider.get();
+ }
+
+ /** This describes all the information needed to create a Neo4j {@link Session}. */
+ @AutoValue
+ public abstract static class DriverConfiguration implements Serializable {
+ public static DriverConfiguration create() {
+ return new AutoValue_Neo4jIO_DriverConfiguration.Builder()
+ .build()
+ .withDefaultConfig(true)
+ .withConfig(Config.defaultConfig());
+ }
+
+ public static DriverConfiguration create(String url, String username, String password) {
+ checkArgument(url != null, "url can not be null");
+ checkArgument(username != null, "username can not be null");
+ checkArgument(password != null, "password can not be null");
+ return new AutoValue_Neo4jIO_DriverConfiguration.Builder()
+ .build()
+ .withDefaultConfig(true)
+ .withConfig(Config.defaultConfig())
+ .withUrl(url)
+ .withUsername(username)
+ .withPassword(password);
+ }
+
+ abstract @Nullable ValueProvider getUrl();
+
+ abstract @Nullable ValueProvider> getUrls();
+
+ abstract @Nullable ValueProvider getUsername();
+
+ abstract @Nullable ValueProvider getPassword();
+
+ abstract @Nullable Config getConfig();
+
+ abstract @Nullable ValueProvider getHasDefaultConfig();
+
+ abstract Builder builder();
+
+ public DriverConfiguration withUrl(String url) {
+ return withUrl(ValueProvider.StaticValueProvider.of(url));
+ }
+
+ public DriverConfiguration withUrl(ValueProvider url) {
+ Preconditions.checkArgument(
+ url != null, "a neo4j connection URL can not be empty or null", url);
+ Preconditions.checkArgument(
+ StringUtils.isNotEmpty(url.get()),
+ "a neo4j connection URL can not be empty or null",
+ url);
+ return builder().setUrl(url).build();
+ }
+
+ public DriverConfiguration withUrls(List urls) {
+ return withUrls(ValueProvider.StaticValueProvider.of(urls));
+ }
+
+ public DriverConfiguration withUrls(ValueProvider> urls) {
+ Preconditions.checkArgument(
+ urls != null, "a list of neo4j connection URLs can not be empty or null", urls);
+ Preconditions.checkArgument(
+ urls.get() != null && !urls.get().isEmpty(),
+ "a neo4j connection URL can not be empty or null",
+ urls);
+ return builder().setUrls(urls).build();
+ }
+
+ public DriverConfiguration withConfig(Config config) {
+ return builder().setConfig(config).build();
+ }
+
+ public DriverConfiguration withUsername(String username) {
+ return withUsername(ValueProvider.StaticValueProvider.of(username));
+ }
+
+ public DriverConfiguration withUsername(ValueProvider username) {
+ Preconditions.checkArgument(username != null, "neo4j username can not be null", username);
+ Preconditions.checkArgument(
+ username.get() != null, "neo4j username can not be null", username);
+ return builder().setUsername(username).build();
+ }
+
+ public DriverConfiguration withPassword(String password) {
+ return withPassword(ValueProvider.StaticValueProvider.of(password));
+ }
+
+ public DriverConfiguration withPassword(ValueProvider password) {
+ Preconditions.checkArgument(password != null, "neo4j password can not be null", password);
+ Preconditions.checkArgument(
+ password.get() != null, "neo4j password can not be null", password);
+ return builder().setPassword(password).build();
+ }
+
+ public DriverConfiguration withDefaultConfig(boolean useDefault) {
+ return withDefaultConfig(ValueProvider.StaticValueProvider.of(useDefault));
+ }
+
+ public DriverConfiguration withDefaultConfig(ValueProvider useDefault) {
+ Preconditions.checkArgument(
+ useDefault != null, "withDefaultConfig parameter useDefault can not be null", useDefault);
+ Preconditions.checkArgument(
+ useDefault.get() != null,
+ "withDefaultConfig parameter useDefault can not be null",
+ useDefault);
+ return builder().setHasDefaultConfig(useDefault).build();
+ }
+
+ void populateDisplayData(DisplayData.Builder builder) {
+ builder.addIfNotNull(DisplayData.item("neo4j-url", getUrl()));
+ builder.addIfNotNull(DisplayData.item("neo4j-username", getUsername()));
+ builder.addIfNotNull(
+ DisplayData.item(
+ "neo4j-password", getPassword() != null ? "" : ""));
+ }
+
+ Driver buildDriver() {
+ // Create the Neo4j Driver
+ // This uses the provided Neo4j configuration along with URLs, username and password
+ //
+ Config config = getConfig();
+ if (config == null) {
+ throw new RuntimeException("please provide a neo4j config");
+ }
+ // We're trying to work around a subtle serialisation bug in the Neo4j Java driver.
+ // The fix is work in progress. For now, we harden our code to avoid
+ // wild goose chases.
+ //
+ Boolean hasDefaultConfig = getProvidedValue(getHasDefaultConfig());
+ if (hasDefaultConfig != null && hasDefaultConfig) {
+ config = Config.defaultConfig();
+ }
+
+ // Get the list of the URI to connect with
+ //
+ List uris = new ArrayList<>();
+ String url = getProvidedValue(getUrl());
+ if (url != null) {
+ try {
+ uris.add(new URI(url));
+ } catch (URISyntaxException e) {
+ throw new RuntimeException("Error creating URI from URL '" + url + "'", e);
+ }
+ }
+ List providedUrls = getProvidedValue(getUrls());
+ if (providedUrls != null) {
+ for (String providedUrl : providedUrls) {
+ try {
+ uris.add(new URI(providedUrl));
+ } catch (URISyntaxException e) {
+ throw new RuntimeException(
+ "Error creating URI '"
+ + providedUrl
+ + "' from a list of "
+ + providedUrls.size()
+ + " URLs",
+ e);
+ }
+ }
+ }
+
+ // A specific routing driver can be used to connect to specific clustered configurations.
+ // Often we don't need it because the Java driver automatically can figure this out
+ // automatically. To keep things simple we use the routing driver in case we have more
+ // than one URL specified. This is an exceptional case.
+ //
+ Driver driver;
+ AuthToken authTokens =
+ getAuthToken(getProvidedValue(getUsername()), getProvidedValue(getPassword()));
+ if (uris.size() > 1) {
+ driver = GraphDatabase.routingDriver(uris, authTokens, config);
+ } else {
+ // Just take the first URI that was provided
+ driver = GraphDatabase.driver(uris.get(0), authTokens, config);
+ }
+
+ return driver;
+ }
+
+ /**
+ * Certain embedded scenarios and so on actually allow for having no authentication at all.
+ *
+ * @param username The username if one is needed
+ * @param password The password if one is needed
+ * @return The AuthToken
+ */
+ protected AuthToken getAuthToken(String username, String password) {
+ if (username != null && password != null) {
+ return AuthTokens.basic(username, password);
+ } else {
+ return AuthTokens.none();
+ }
+ }
+
+ /**
+ * The Builder class below is not visible. We use it to service the "with" methods below the
+ * Builder class.
+ */
+ @AutoValue.Builder
+ abstract static class Builder {
+ abstract Builder setUrl(ValueProvider url);
+
+ abstract Builder setUrls(ValueProvider> url);
+
+ abstract Builder setUsername(ValueProvider username);
+
+ abstract Builder setPassword(ValueProvider password);
+
+ abstract Builder setConfig(Config config);
+
+ abstract Builder setHasDefaultConfig(ValueProvider useDefault);
+
+ abstract DriverConfiguration build();
+ }
+ }
+
+ /** This is the class which handles the work behind the {@link #readAll} method. */
+ @AutoValue
+ public abstract static class ReadAll
+ extends PTransform, PCollection> {
+
+ abstract @Nullable SerializableFunction getDriverProviderFn();
+
+ abstract @Nullable SessionConfig getSessionConfig();
+
+ abstract @Nullable TransactionConfig getTransactionConfig();
+
+ abstract @Nullable ValueProvider getCypher();
+
+ abstract @Nullable ValueProvider getWriteTransaction();
+
+ abstract @Nullable RowMapper getRowMapper();
+
+ abstract @Nullable SerializableFunction>
+ getParametersFunction();
+
+ abstract @Nullable Coder getCoder();
+
+ abstract @Nullable ValueProvider getLogCypher();
+
+ abstract Builder toBuilder();
+
+ public ReadAll withDriverConfiguration(DriverConfiguration config) {
+ return toBuilder()
+ .setDriverProviderFn(new DriverProviderFromDriverConfiguration(config))
+ .build();
+ }
+
+ public ReadAll withCypher(String cypher) {
+ checkArgument(
+ cypher != null, "Neo4jIO.readAll().withCypher(query) called with null cypher query");
+ return withCypher(ValueProvider.StaticValueProvider.of(cypher));
+ }
+
+ public ReadAll withCypher(ValueProvider cypher) {
+ checkArgument(cypher != null, "Neo4jIO.readAll().withCypher(cypher) called with null cypher");
+ return toBuilder().setCypher(cypher).build();
+ }
+
+ public ReadAll withSessionConfig(SessionConfig sessionConfig) {
+ checkArgument(
+ sessionConfig != null,
+ "Neo4jIO.readAll().withSessionConfig(sessionConfig) called with null sessionConfig");
+ return toBuilder().setSessionConfig(sessionConfig).build();
+ }
+
+ public ReadAll withTransactionConfig(TransactionConfig transactionConfig) {
+ checkArgument(
+ transactionConfig != null,
+ "Neo4jIO.readAll().withTransactionConfig(transactionConfig) called with null transactionConfig");
+ return toBuilder().setTransactionConfig(transactionConfig).build();
+ }
+
+ public ReadAll withRowMapper(RowMapper rowMapper) {
+ checkArgument(
+ rowMapper != null,
+ "Neo4jIO.readAll().withRowMapper(rowMapper) called with null rowMapper");
+ return toBuilder().setRowMapper(rowMapper).build();
+ }
+
+ public ReadAll withParametersFunction(
+ SerializableFunction> parametersFunction) {
+ checkArgument(
+ parametersFunction != null,
+ "Neo4jIO.readAll().withParametersFunction(parametersFunction) called with null parametersFunction");
+ return toBuilder().setParametersFunction(parametersFunction).build();
+ }
+
+ public ReadAll withCoder(Coder coder) {
+ checkArgument(coder != null, "Neo4jIO.readAll().withCoder(coder) called with null coder");
+ return toBuilder().setCoder(coder).build();
+ }
+
+ public ReadAll withReadTransaction() {
+ return toBuilder()
+ .setWriteTransaction(ValueProvider.StaticValueProvider.of(Boolean.FALSE))
+ .build();
+ }
+
+ public ReadAll withWriteTransaction() {
+ return toBuilder()
+ .setWriteTransaction(ValueProvider.StaticValueProvider.of(Boolean.TRUE))
+ .build();
+ }
+
+ public ReadAll withCypherLogging() {
+ return toBuilder().setLogCypher(ValueProvider.StaticValueProvider.of(Boolean.TRUE)).build();
+ }
+
+ @Override
+ public PCollection expand(PCollection input) {
+
+ final SerializableFunction driverProviderFn = getDriverProviderFn();
+ final RowMapper rowMapper = getRowMapper();
+ SerializableFunction> parametersFunction =
+ getParametersFunction();
+
+ final String cypher = getProvidedValue(getCypher());
+ checkArgument(cypher != null, "please provide a cypher statement to execute");
+
+ SessionConfig sessionConfig = getSessionConfig();
+ if (sessionConfig == null) {
+ // Create a default session configuration as recommended by Neo4j
+ //
+ sessionConfig = SessionConfig.defaultConfig();
+ }
+
+ TransactionConfig transactionConfig = getTransactionConfig();
+ if (transactionConfig == null) {
+ transactionConfig = TransactionConfig.empty();
+ }
+
+ Boolean writeTransaction = getProvidedValue(getWriteTransaction());
+ if (writeTransaction == null) {
+ writeTransaction = Boolean.FALSE;
+ }
+
+ Boolean logCypher = getProvidedValue(getLogCypher());
+ if (logCypher == null) {
+ logCypher = Boolean.FALSE;
+ }
+
+ if (driverProviderFn == null) {
+ throw new RuntimeException("please provide a driver provider");
+ }
+ if (rowMapper == null) {
+ throw new RuntimeException("please provide a row mapper");
+ }
+ if (parametersFunction == null) {
+ parametersFunction = t -> Collections.emptyMap();
+ }
+
+ ReadFn readFn =
+ new ReadFn<>(
+ driverProviderFn,
+ sessionConfig,
+ transactionConfig,
+ cypher,
+ rowMapper,
+ parametersFunction,
+ writeTransaction,
+ logCypher);
+
+ return getOutputPCollection(input, readFn, getCoder());
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ super.populateDisplayData(builder);
+ String cypher = getProvidedValue(getCypher());
+ if (cypher == null) {
+ cypher = "";
+ }
+ builder.add(DisplayData.item("cypher", cypher));
+ SerializableFunction driverProviderFn = getDriverProviderFn();
+ if (driverProviderFn != null) {
+ if (driverProviderFn instanceof HasDisplayData) {
+ ((HasDisplayData) driverProviderFn).populateDisplayData(builder);
+ }
+ }
+ }
+
+ @AutoValue.Builder
+ abstract static class Builder {
+ abstract Builder setDriverProviderFn(
+ SerializableFunction driverProviderFn);
+
+ abstract Builder setCypher(ValueProvider cypher);
+
+ abstract Builder setSessionConfig(SessionConfig sessionConfig);
+
+ abstract Builder setTransactionConfig(
+ TransactionConfig transactionConfig);
+
+ abstract Builder setWriteTransaction(
+ ValueProvider writeTransaction);
+
+ abstract Builder setRowMapper(RowMapper rowMapper);
+
+ abstract Builder setParametersFunction(
+ SerializableFunction> parametersFunction);
+
+ abstract Builder setCoder(Coder coder);
+
+ abstract Builder setLogCypher(ValueProvider logCypher);
+
+ abstract ReadAll build();
+ }
+ }
+
+ /** A {@link DoFn} to execute a Cypher query to read from Neo4j. */
+ private static class ReadWriteFn extends DoFn {
+ protected static class DriverSession {
+ public @Nullable Driver driver;
+ public @Nullable Session session;
+ public @NonNull AtomicBoolean closed;
+
+ protected DriverSession(Driver driver, Session session) {
+ this.driver = driver;
+ this.session = session;
+ this.closed = new AtomicBoolean(false);
+ }
+
+ private DriverSession() {
+ this.driver = null;
+ this.session = null;
+ this.closed = new AtomicBoolean(true);
+ }
+
+ protected static @NonNull DriverSession emptyClosed() {
+ return new DriverSession();
+ }
+ }
+
+ protected final @NonNull SerializableFunction driverProviderFn;
+ protected final @NonNull SessionConfig sessionConfig;
+ protected final @NonNull TransactionConfig transactionConfig;
+
+ protected transient @NonNull DriverSession driverSession;
+
+ protected ReadWriteFn(
+ @NonNull SerializableFunction driverProviderFn,
+ @NonNull SessionConfig sessionConfig,
+ @NonNull TransactionConfig transactionConfig) {
+ this.driverProviderFn = driverProviderFn;
+ this.sessionConfig = sessionConfig;
+ this.transactionConfig = transactionConfig;
+ this.driverSession = DriverSession.emptyClosed();
+ }
+
+ /**
+ * Delay the creation of driver and session until we actually have data to do something with.
+ */
+ @Setup
+ public void setup() {}
+
+ protected @NonNull Driver createDriver() {
+ Driver driver = driverProviderFn.apply(null);
+ if (driver == null) {
+ throw new RuntimeException("null driver given by driver provider");
+ }
+ return driver;
+ }
+
+ protected @Initialized @NonNull DriverSession buildDriverSession() {
+ @NonNull Driver driver = createDriver();
+ @NonNull Session session = driver.session(sessionConfig);
+ return new DriverSession(driver, session);
+ }
+
+ @StartBundle
+ public void startBundle() {
+ if (driverSession == null || driverSession.closed.get()) {
+ driverSession = buildDriverSession();
+ }
+ }
+
+ @FinishBundle
+ public void finishBundle() {
+ cleanUpDriverSession();
+ }
+
+ @Override
+ protected void finalize() {
+ cleanUpDriverSession();
+ }
+
+ protected void cleanUpDriverSession() {
+ if (!driverSession.closed.get()) {
+ try {
+ if (driverSession.session != null) {
+ driverSession.session.close();
+ }
+ if (driverSession.driver != null) {
+ driverSession.driver.close();
+ }
+ } finally {
+ driverSession.closed.set(true);
+ }
+ }
+ }
+
+ protected String getParametersString(Map parametersMap) {
+ StringBuilder parametersString = new StringBuilder();
+ parametersMap
+ .keySet()
+ .forEach(
+ key -> {
+ if (parametersString.length() > 0) {
+ parametersString.append(',');
+ }
+ parametersString.append(key).append('=');
+ Object value = parametersMap.get(key);
+ if (value == null) {
+ parametersString.append("");
+ } else {
+ parametersString.append(value);
+ }
+ });
+ return parametersString.toString();
+ }
+ }
+
+ /** A {@link DoFn} to execute a Cypher query to read from Neo4j. */
+ private static class ReadFn extends ReadWriteFn {
+ protected final @NonNull String cypher;
+ protected final @NonNull RowMapper rowMapper;
+ protected final @Nullable SerializableFunction>
+ parametersFunction;
+
+ private final boolean writeTransaction;
+ private final boolean logCypher;
+
+ private ReadFn(
+ @NonNull SerializableFunction driverProviderFn,
+ @NonNull SessionConfig sessionConfig,
+ @NonNull TransactionConfig transactionConfig,
+ @NonNull String cypher,
+ @NonNull RowMapper rowMapper,
+ @Nullable SerializableFunction> parametersFunction,
+ boolean writeTransaction,
+ boolean logCypher) {
+ super(driverProviderFn, sessionConfig, transactionConfig);
+ this.cypher = cypher;
+ this.rowMapper = rowMapper;
+ this.parametersFunction = parametersFunction;
+ this.writeTransaction = writeTransaction;
+ this.logCypher = logCypher;
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext context) {
+ // Map the input data to the parameters map...
+ //
+ ParameterT parameters = context.element();
+ final Map parametersMap;
+ if (parametersFunction != null) {
+ parametersMap = parametersFunction.apply(parameters);
+ } else {
+ parametersMap = Collections.emptyMap();
+ }
+ executeReadCypherStatement(context, parametersMap);
+ }
+
+ private void executeReadCypherStatement(
+ final ProcessContext processContext, Map parametersMap) {
+ // The actual "reading" work needs to happen in a transaction.
+ // We could actually read and write here depending on the type of transaction
+ // we picked. As long as the Cypher statement returns values it's fine.
+ //
+ TransactionWork transactionWork =
+ transaction -> {
+ Result result = transaction.run(cypher, parametersMap);
+ while (result.hasNext()) {
+ Record record = result.next();
+ try {
+ OutputT outputT = rowMapper.mapRow(record);
+ processContext.output(outputT);
+ } catch (Exception e) {
+ throw new RuntimeException("error mapping Neo4j record to row", e);
+ }
+ }
+
+ // We deliver no specific Neo4j transaction output beyond what goes to the context
+ // output
+ //
+ return null;
+ };
+
+ if (logCypher) {
+ String parametersString = getParametersString(parametersMap);
+
+ String readWrite = writeTransaction ? "write" : "read";
+ LOG.info(
+ "Starting a "
+ + readWrite
+ + " transaction for cypher: "
+ + cypher
+ + ", parameters: "
+ + parametersString);
+ }
+
+ // There are 2 ways to do a transaction on Neo4j: read or write
+ // It's important that the right type is selected, especially in clustered configurations.
+ //
+ if (driverSession.session == null) {
+ throw new RuntimeException("neo4j session was not initialized correctly");
+ } else {
+ if (writeTransaction) {
+ driverSession.session.writeTransaction(transactionWork, transactionConfig);
+ } else {
+ driverSession.session.readTransaction(transactionWork, transactionConfig);
+ }
+ }
+ }
+ }
+
+ /**
+ * Wraps a {@link DriverConfiguration} to provide a {@link Driver}.
+ *
+ * At most a single {@link Driver} instance will be constructed during pipeline execution for
+ * each unique {@link DriverConfiguration} within the pipeline.
+ */
+ public static class DriverProviderFromDriverConfiguration
+ implements SerializableFunction, HasDisplayData {
+ private final DriverConfiguration config;
+
+ private DriverProviderFromDriverConfiguration(DriverConfiguration config) {
+ this.config = config;
+ }
+
+ public static SerializableFunction of(DriverConfiguration config) {
+ return new DriverProviderFromDriverConfiguration(config);
+ }
+
+ @Override
+ public Driver apply(Void input) {
+ return config.buildDriver();
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ config.populateDisplayData(builder);
+ }
+ }
+
+ /** This is the class which handles the work behind the {@link #writeUnwind()} method. */
+ @AutoValue
+ public abstract static class WriteUnwind
+ extends PTransform, PDone> {
+
+ abstract @Nullable SerializableFunction getDriverProviderFn();
+
+ abstract @Nullable ValueProvider getSessionConfig();
+
+ abstract @Nullable ValueProvider getCypher();
+
+ abstract @Nullable ValueProvider getUnwindMapName();
+
+ abstract @Nullable ValueProvider getTransactionConfig();
+
+ abstract @Nullable SerializableFunction>
+ getParametersFunction();
+
+ abstract @Nullable ValueProvider getBatchSize();
+
+ abstract @Nullable ValueProvider getLogCypher();
+
+ abstract Builder toBuilder();
+
+ public WriteUnwind withDriverConfiguration(DriverConfiguration config) {
+ return toBuilder()
+ .setDriverProviderFn(new DriverProviderFromDriverConfiguration(config))
+ .build();
+ }
+
+ public WriteUnwind withCypher(String cypher) {
+ checkArgument(
+ cypher != null, "Neo4jIO.writeUnwind().withCypher(query) called with null cypher query");
+ return withCypher(ValueProvider.StaticValueProvider.of(cypher));
+ }
+
+ public WriteUnwind withCypher(ValueProvider cypher) {
+ checkArgument(
+ cypher != null, "Neo4jIO.writeUnwind().withCypher(cypher) called with null cypher");
+ return toBuilder().setCypher(cypher).build();
+ }
+
+ public WriteUnwind withUnwindMapName(String mapName) {
+ checkArgument(
+ mapName != null,
+ "Neo4jIO.writeUnwind().withUnwindMapName(query) called with null mapName");
+ return withUnwindMapName(ValueProvider.StaticValueProvider.of(mapName));
+ }
+
+ public WriteUnwind withUnwindMapName(ValueProvider mapName) {
+ checkArgument(
+ mapName != null,
+ "Neo4jIO.writeUnwind().withUnwindMapName(cypher) called with null mapName");
+ return toBuilder().setUnwindMapName(mapName).build();
+ }
+
+ public WriteUnwind withTransactionConfig(TransactionConfig transactionConfig) {
+ checkArgument(
+ transactionConfig != null,
+ "Neo4jIO.writeUnwind().withTransactionConfig(config) called with null transactionConfig");
+ return withTransactionConfig(ValueProvider.StaticValueProvider.of(transactionConfig));
+ }
+
+ public WriteUnwind withTransactionConfig(
+ ValueProvider transactionConfig) {
+ checkArgument(
+ transactionConfig != null,
+ "Neo4jIO.writeUnwind().withTransactionConfig(config) called with null transactionConfig");
+ return toBuilder().setTransactionConfig(transactionConfig).build();
+ }
+
+ public WriteUnwind withSessionConfig(SessionConfig sessionConfig) {
+ checkArgument(
+ sessionConfig != null,
+ "Neo4jIO.writeUnwind().withSessionConfig(sessionConfig) called with null sessionConfig");
+ return withSessionConfig(ValueProvider.StaticValueProvider.of(sessionConfig));
+ }
+
+ public WriteUnwind withSessionConfig(ValueProvider sessionConfig) {
+ checkArgument(
+ sessionConfig != null,
+ "Neo4jIO.writeUnwind().withSessionConfig(sessionConfig) called with null sessionConfig");
+ return toBuilder().setSessionConfig(sessionConfig).build();
+ }
+
+ // Batch size
+ public WriteUnwind withBatchSize(long batchSize) {
+ checkArgument(
+ batchSize > 0, "Neo4jIO.writeUnwind().withFetchSize(query) called with batchSize<=0");
+ return withBatchSize(ValueProvider.StaticValueProvider.of(batchSize));
+ }
+
+ public WriteUnwind withBatchSize(ValueProvider batchSize) {
+ checkArgument(
+ batchSize != null && batchSize.get() >= 0,
+ "Neo4jIO.readAll().withBatchSize(query) called with batchSize<=0");
+ return toBuilder().setBatchSize(batchSize).build();
+ }
+
+ public WriteUnwind withParametersFunction(
+ SerializableFunction> parametersFunction) {
+ checkArgument(
+ parametersFunction != null,
+ "Neo4jIO.readAll().withParametersFunction(parametersFunction) called with null parametersFunction");
+ return toBuilder().setParametersFunction(parametersFunction).build();
+ }
+
+ public WriteUnwind withCypherLogging() {
+ return toBuilder().setLogCypher(ValueProvider.StaticValueProvider.of(Boolean.TRUE)).build();
+ }
+
+ @Override
+ public PDone expand(PCollection input) {
+
+ final SerializableFunction