From 080260e1f8c29f3fa4b181cfd64a2bbf4770b239 Mon Sep 17 00:00:00 2001 From: Suvrat1629 Date: Mon, 25 Aug 2025 00:40:11 +0530 Subject: [PATCH 1/2] output write statistics --- .../org/apache/beam/sdk/io/neo4j/Neo4jIO.java | 395 ++++++++++++++++++ 1 file changed, 395 insertions(+) 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 index 1dd95b44c2fc..144f010ae9f5 100644 --- 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 @@ -31,6 +31,8 @@ import java.util.concurrent.atomic.AtomicBoolean; import org.apache.beam.repackaged.core.org.apache.commons.lang3.StringUtils; import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.DefaultCoder; +import org.apache.beam.sdk.extensions.avro.coders.AvroCoder; import org.apache.beam.sdk.harness.JvmInitializer; import org.apache.beam.sdk.options.ValueProvider; import org.apache.beam.sdk.schemas.NoSuchSchemaException; @@ -42,13 +44,17 @@ 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.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; 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.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_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.joda.time.Instant; import org.neo4j.driver.AuthToken; import org.neo4j.driver.AuthTokens; import org.neo4j.driver.Config; @@ -60,6 +66,7 @@ import org.neo4j.driver.SessionConfig; import org.neo4j.driver.TransactionConfig; import org.neo4j.driver.TransactionWork; +import org.neo4j.driver.summary.SummaryCounters; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -183,6 +190,63 @@ public class Neo4jIO { private static final Logger LOG = LoggerFactory.getLogger(Neo4jIO.class); + /** + * Represents write statistics from a Neo4j UNWIND operation, used by WriteUnwindWithResults. + * Contains metrics such as the number of nodes created, relationships created, and properties set. + */ + @DefaultCoder(AvroCoder.class) + public static class Neo4jWriteStats implements Serializable { + private final int nodesCreated; + private final int nodesDeleted; + private final int relationshipsCreated; + private final int relationshipsDeleted; + private final int propertiesSet; + private final int labelsAdded; + private final int labelsRemoved; + + /** + * Constructs a Neo4jWriteStats instance from a Neo4j SummaryCounters object. + * + * @param counters The SummaryCounters object containing write statistics. + */ + public Neo4jWriteStats(SummaryCounters counters) { + this.nodesCreated = counters.nodesCreated(); + this.nodesDeleted = counters.nodesDeleted(); + this.relationshipsCreated = counters.relationshipsCreated(); + this.relationshipsDeleted = counters.relationshipsDeleted(); + this.propertiesSet = counters.propertiesSet(); + this.labelsAdded = counters.labelsAdded(); + this.labelsRemoved = counters.labelsRemoved(); + } + + public int getNodesCreated() { + return nodesCreated; + } + + public int getNodesDeleted() { + return nodesDeleted; + } + + public int getRelationshipsCreated() { + return relationshipsCreated; + } + + public int getRelationshipsDeleted() { + return relationshipsDeleted; + } + + public int getPropertiesSet() { + return propertiesSet; + } + + public int getLabelsAdded() { + return labelsAdded; + } + + public int getLabelsRemoved() { + return labelsRemoved; + } + } /** * Read all rows using a Neo4j Cypher query. * @@ -205,6 +269,18 @@ public static WriteUnwind writeUnwind() { .build(); } + /** + * Write all rows using a Neo4j Cypher UNWIND statement and return write statistics. + * This sets a default batch size of 5000 and returns a PCollection of write statistics. + * + * @param Type of the data representing query parameters. + */ + public static WriteUnwindWithResults writeUnwindWithResults() { + return new AutoValue_Neo4jIO_WriteUnwindWithResults.Builder() + .setBatchSize(ValueProvider.StaticValueProvider.of(5000L)) + .build(); + } + private static PCollection getOutputPCollection( PCollection input, DoFn writeFn, @@ -1088,6 +1164,201 @@ abstract Builder setParametersFunction( } } + /** This is the class which handles the work behind the {@link #writeUnwindWithResults()} method. */ + @AutoValue + public abstract static class WriteUnwindWithResults + extends PTransform, PCollection> { + + 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 WriteUnwindWithResults withDriverConfiguration(DriverConfiguration config) { + return toBuilder() + .setDriverProviderFn(new DriverProviderFromDriverConfiguration(config)) + .build(); + } + + public WriteUnwindWithResults withCypher(String cypher) { + checkArgument( + cypher != null, "Neo4jIO.writeUnwindWithResults().withCypher(cypher) called with null cypher"); + return withCypher(ValueProvider.StaticValueProvider.of(cypher)); + } + + public WriteUnwindWithResults withCypher(ValueProvider cypher) { + checkArgument( + cypher != null, "Neo4jIO.writeUnwindWithResults().withCypher(cypher) called with null cypher"); + return toBuilder().setCypher(cypher).build(); + } + + public WriteUnwindWithResults withUnwindMapName(String mapName) { + checkArgument( + mapName != null, + "Neo4jIO.writeUnwindWithResults().withUnwindMapName(mapName) called with null mapName"); + return withUnwindMapName(ValueProvider.StaticValueProvider.of(mapName)); + } + + public WriteUnwindWithResults withUnwindMapName(ValueProvider mapName) { + checkArgument( + mapName != null, + "Neo4jIO.writeUnwindWithResults().withUnwindMapName(mapName) called with null mapName"); + return toBuilder().setUnwindMapName(mapName).build(); + } + + public WriteUnwindWithResults withTransactionConfig(TransactionConfig transactionConfig) { + checkArgument( + transactionConfig != null, + "Neo4jIO.writeUnwindWithResults().withTransactionConfig(transactionConfig) called with null transactionConfig"); + return withTransactionConfig(ValueProvider.StaticValueProvider.of(transactionConfig)); + } + + public WriteUnwindWithResults withTransactionConfig( + ValueProvider transactionConfig) { + checkArgument( + transactionConfig != null, + "Neo4jIO.writeUnwindWithResults().withTransactionConfig(transactionConfig) called with null transactionConfig"); + return toBuilder().setTransactionConfig(transactionConfig).build(); + } + + public WriteUnwindWithResults withSessionConfig(SessionConfig sessionConfig) { + checkArgument( + sessionConfig != null, + "Neo4jIO.writeUnwindWithResults().withSessionConfig(sessionConfig) called with null sessionConfig"); + return withSessionConfig(ValueProvider.StaticValueProvider.of(sessionConfig)); + } + + public WriteUnwindWithResults withSessionConfig(ValueProvider sessionConfig) { + checkArgument( + sessionConfig != null, + "Neo4jIO.writeUnwindWithResults().withSessionConfig(sessionConfig) called with null sessionConfig"); + return toBuilder().setSessionConfig(sessionConfig).build(); + } + + public WriteUnwindWithResults withBatchSize(long batchSize) { + checkArgument( + batchSize > 0, "Neo4jIO.writeUnwindWithResults().withBatchSize(batchSize) called with batchSize<=0"); + return withBatchSize(ValueProvider.StaticValueProvider.of(batchSize)); + } + + public WriteUnwindWithResults withBatchSize(ValueProvider batchSize) { + checkArgument( + batchSize != null && batchSize.get() >= 0, + "Neo4jIO.writeUnwindWithResults().withBatchSize(batchSize) called with batchSize<=0"); + return toBuilder().setBatchSize(batchSize).build(); + } + + public WriteUnwindWithResults withParametersFunction( + SerializableFunction> parametersFunction) { + checkArgument( + parametersFunction != null, + "Neo4jIO.writeUnwindWithResults().withParametersFunction(parametersFunction) called with null parametersFunction"); + return toBuilder().setParametersFunction(parametersFunction).build(); + } + + public WriteUnwindWithResults withCypherLogging() { + return toBuilder().setLogCypher(ValueProvider.StaticValueProvider.of(Boolean.TRUE)).build(); + } + + @Override + public PCollection expand(PCollection input) { + final SerializableFunction driverProviderFn = getDriverProviderFn(); + final SerializableFunction> parametersFunction = + getParametersFunction(); + SessionConfig sessionConfig = getProvidedValue(getSessionConfig()); + if (sessionConfig == null) { + sessionConfig = SessionConfig.defaultConfig(); + } + TransactionConfig transactionConfig = getProvidedValue(getTransactionConfig()); + if (transactionConfig == null) { + transactionConfig = TransactionConfig.empty(); + } + final String cypher = getProvidedValue(getCypher()); + checkArgument(cypher != null, "please provide an unwind cypher statement to execute"); + final String unwindMapName = getProvidedValue(getUnwindMapName()); + checkArgument(unwindMapName != null, "please provide an unwind map name"); + + Long batchSize = getProvidedValue(getBatchSize()); + if (batchSize == null || batchSize <= 0) { + batchSize = 5000L; + } + + Boolean logCypher = getProvidedValue(getLogCypher()); + if (logCypher == null) { + logCypher = Boolean.FALSE; + } + + if (driverProviderFn == null) { + throw new RuntimeException("please provide a driver provider"); + } + if (parametersFunction == null) { + throw new RuntimeException("please provide a parameters function"); + } + + WriteUnwindWithResultsFn writeFn = + new WriteUnwindWithResultsFn<>( + driverProviderFn, + sessionConfig, + transactionConfig, + cypher, + parametersFunction, + batchSize, + logCypher, + unwindMapName); + + return getOutputPCollection(input, writeFn, AvroCoder.of(Neo4jWriteStats.class)); + } + + @Override + public void populateDisplayData(DisplayData.Builder builder) { + super.populateDisplayData(builder); + builder.add(DisplayData.item("cypher", getCypher())); + final 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 setSessionConfig(ValueProvider sessionConfig); + + abstract Builder setTransactionConfig( + ValueProvider transactionConfig); + + abstract Builder setCypher(ValueProvider cypher); + + abstract Builder setUnwindMapName(ValueProvider unwindMapName); + + abstract Builder setParametersFunction( + SerializableFunction> parametersFunction); + + abstract Builder setBatchSize(ValueProvider batchSize); + + abstract Builder setLogCypher(ValueProvider logCypher); + + abstract WriteUnwindWithResults build(); + } + } + /** A {@link DoFn} to execute a Cypher query to read from Neo4j. */ private static class WriteUnwindFn extends ReadWriteFn { @@ -1203,4 +1474,128 @@ public void finishBundle() { executeCypherUnwindStatement(); } } + + private static class WriteUnwindWithResultsFn extends ReadWriteFn { + private final @NonNull String cypher; + private final @Nullable SerializableFunction> parametersFunction; + private final boolean logCypher; + private final long batchSize; + private final @NonNull String unwindMapName; + + private long elementsInput; + private boolean loggingDone; + private List> unwindList; + + private WriteUnwindWithResultsFn( + @NonNull SerializableFunction driverProviderFn, + @NonNull SessionConfig sessionConfig, + @NonNull TransactionConfig transactionConfig, + @NonNull String cypher, + @Nullable SerializableFunction> parametersFunction, + long batchSize, + boolean logCypher, + String unwindMapName) { + super(driverProviderFn, sessionConfig, transactionConfig); + this.cypher = cypher; + this.parametersFunction = parametersFunction; + this.logCypher = logCypher; + this.batchSize = batchSize; + this.unwindMapName = unwindMapName; + + unwindList = new ArrayList<>(); + elementsInput = 0; + loggingDone = false; + } + + @ProcessElement + public void processElement(@Element ParameterT parameters, ProcessContext context) { + if (parametersFunction != null) { + unwindList.add(parametersFunction.apply(parameters)); + } else { + unwindList.add(Collections.emptyMap()); + } + elementsInput++; + + if (elementsInput >= batchSize) { + if (elementsInput > 0) { + final Map parametersMap = new HashMap<>(); + parametersMap.put(unwindMapName, unwindList); + + TransactionWork transactionWork = + transaction -> { + Result result = transaction.run(cypher, parametersMap); + Neo4jWriteStats stats = new Neo4jWriteStats(result.consume().counters()); + context.outputWithTimestamp(stats, context.timestamp()); + return null; + }; + + if (logCypher && !loggingDone) { + String parametersString = getParametersString(parametersMap); + LOG.info( + "Starting a write transaction for unwind statement cypher: " + + cypher + + ", parameters: " + + parametersString); + loggingDone = true; + } + + if (driverSession.session == null) { + throw new RuntimeException("neo4j session was not initialized correctly"); + } else { + try { + driverSession.session.writeTransaction(transactionWork, transactionConfig); + } catch (Exception e) { + throw new RuntimeException( + "Error writing " + unwindList.size() + " rows to Neo4j with Cypher: " + cypher, e); + } + } + + unwindList.clear(); + elementsInput = 0; + } + } + } + + @FinishBundle + public void finishBundle(FinishBundleContext context) { + if (elementsInput > 0) { + final Map parametersMap = new HashMap<>(); + parametersMap.put(unwindMapName, unwindList); + + TransactionWork transactionWork = + transaction -> { + Result result = transaction.run(cypher, parametersMap); + Neo4jWriteStats stats = new Neo4jWriteStats(result.consume().counters()); + Instant timestamp = Instant.now(); + BoundedWindow window = GlobalWindow.INSTANCE; + context.output(stats, timestamp, window); + return null; + }; + + if (logCypher && !loggingDone) { + String parametersString = getParametersString(parametersMap); + LOG.info( + "Starting a write transaction for unwind statement cypher: " + + cypher + + ", parameters: " + + parametersString); + loggingDone = true; + } + + if (driverSession.session == null) { + throw new RuntimeException("neo4j session was not initialized correctly"); + } else { + try { + driverSession.session.writeTransaction(transactionWork, transactionConfig); + } catch (Exception e) { + throw new RuntimeException( + "Error writing " + unwindList.size() + " rows to Neo4j with Cypher: " + cypher, e); + } + } + + unwindList.clear(); + elementsInput = 0; + } + } + } } From 476fec07cb23b3533220235a144dc469d35cf019 Mon Sep 17 00:00:00 2001 From: Suvrat1629 Date: Mon, 25 Aug 2025 01:13:06 +0530 Subject: [PATCH 2/2] spotless fix --- .../org/apache/beam/sdk/io/neo4j/Neo4jIO.java | 173 ++++++++++-------- 1 file changed, 92 insertions(+), 81 deletions(-) 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 index 144f010ae9f5..c2c1f16d5d84 100644 --- 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 @@ -49,7 +49,6 @@ 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.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.checkerframework.checker.initialization.qual.Initialized; import org.checkerframework.checker.nullness.qual.NonNull; @@ -192,7 +191,8 @@ public class Neo4jIO { /** * Represents write statistics from a Neo4j UNWIND operation, used by WriteUnwindWithResults. - * Contains metrics such as the number of nodes created, relationships created, and properties set. + * Contains metrics such as the number of nodes created, relationships created, and properties + * set. */ @DefaultCoder(AvroCoder.class) public static class Neo4jWriteStats implements Serializable { @@ -270,15 +270,15 @@ public static WriteUnwind writeUnwind() { } /** - * Write all rows using a Neo4j Cypher UNWIND statement and return write statistics. - * This sets a default batch size of 5000 and returns a PCollection of write statistics. + * Write all rows using a Neo4j Cypher UNWIND statement and return write statistics. This sets a + * default batch size of 5000 and returns a PCollection of write statistics. * * @param Type of the data representing query parameters. */ public static WriteUnwindWithResults writeUnwindWithResults() { return new AutoValue_Neo4jIO_WriteUnwindWithResults.Builder() - .setBatchSize(ValueProvider.StaticValueProvider.of(5000L)) - .build(); + .setBatchSize(ValueProvider.StaticValueProvider.of(5000L)) + .build(); } private static PCollection getOutputPCollection( @@ -1164,10 +1164,12 @@ abstract Builder setParametersFunction( } } - /** This is the class which handles the work behind the {@link #writeUnwindWithResults()} method. */ + /** + * This is the class which handles the work behind the {@link #writeUnwindWithResults()} method. + */ @AutoValue public abstract static class WriteUnwindWithResults - extends PTransform, PCollection> { + extends PTransform, PCollection> { abstract @Nullable SerializableFunction getDriverProviderFn(); @@ -1179,7 +1181,8 @@ public abstract static class WriteUnwindWithResults abstract @Nullable ValueProvider getTransactionConfig(); - abstract @Nullable SerializableFunction> getParametersFunction(); + abstract @Nullable SerializableFunction> + getParametersFunction(); abstract @Nullable ValueProvider getBatchSize(); @@ -1189,83 +1192,88 @@ public abstract static class WriteUnwindWithResults public WriteUnwindWithResults withDriverConfiguration(DriverConfiguration config) { return toBuilder() - .setDriverProviderFn(new DriverProviderFromDriverConfiguration(config)) - .build(); + .setDriverProviderFn(new DriverProviderFromDriverConfiguration(config)) + .build(); } public WriteUnwindWithResults withCypher(String cypher) { checkArgument( - cypher != null, "Neo4jIO.writeUnwindWithResults().withCypher(cypher) called with null cypher"); + cypher != null, + "Neo4jIO.writeUnwindWithResults().withCypher(cypher) called with null cypher"); return withCypher(ValueProvider.StaticValueProvider.of(cypher)); } public WriteUnwindWithResults withCypher(ValueProvider cypher) { checkArgument( - cypher != null, "Neo4jIO.writeUnwindWithResults().withCypher(cypher) called with null cypher"); + cypher != null, + "Neo4jIO.writeUnwindWithResults().withCypher(cypher) called with null cypher"); return toBuilder().setCypher(cypher).build(); } public WriteUnwindWithResults withUnwindMapName(String mapName) { checkArgument( - mapName != null, - "Neo4jIO.writeUnwindWithResults().withUnwindMapName(mapName) called with null mapName"); + mapName != null, + "Neo4jIO.writeUnwindWithResults().withUnwindMapName(mapName) called with null mapName"); return withUnwindMapName(ValueProvider.StaticValueProvider.of(mapName)); } public WriteUnwindWithResults withUnwindMapName(ValueProvider mapName) { checkArgument( - mapName != null, - "Neo4jIO.writeUnwindWithResults().withUnwindMapName(mapName) called with null mapName"); + mapName != null, + "Neo4jIO.writeUnwindWithResults().withUnwindMapName(mapName) called with null mapName"); return toBuilder().setUnwindMapName(mapName).build(); } - public WriteUnwindWithResults withTransactionConfig(TransactionConfig transactionConfig) { + public WriteUnwindWithResults withTransactionConfig( + TransactionConfig transactionConfig) { checkArgument( - transactionConfig != null, - "Neo4jIO.writeUnwindWithResults().withTransactionConfig(transactionConfig) called with null transactionConfig"); + transactionConfig != null, + "Neo4jIO.writeUnwindWithResults().withTransactionConfig(transactionConfig) called with null transactionConfig"); return withTransactionConfig(ValueProvider.StaticValueProvider.of(transactionConfig)); } public WriteUnwindWithResults withTransactionConfig( - ValueProvider transactionConfig) { + ValueProvider transactionConfig) { checkArgument( - transactionConfig != null, - "Neo4jIO.writeUnwindWithResults().withTransactionConfig(transactionConfig) called with null transactionConfig"); + transactionConfig != null, + "Neo4jIO.writeUnwindWithResults().withTransactionConfig(transactionConfig) called with null transactionConfig"); return toBuilder().setTransactionConfig(transactionConfig).build(); } public WriteUnwindWithResults withSessionConfig(SessionConfig sessionConfig) { checkArgument( - sessionConfig != null, - "Neo4jIO.writeUnwindWithResults().withSessionConfig(sessionConfig) called with null sessionConfig"); + sessionConfig != null, + "Neo4jIO.writeUnwindWithResults().withSessionConfig(sessionConfig) called with null sessionConfig"); return withSessionConfig(ValueProvider.StaticValueProvider.of(sessionConfig)); } - public WriteUnwindWithResults withSessionConfig(ValueProvider sessionConfig) { + public WriteUnwindWithResults withSessionConfig( + ValueProvider sessionConfig) { checkArgument( - sessionConfig != null, - "Neo4jIO.writeUnwindWithResults().withSessionConfig(sessionConfig) called with null sessionConfig"); + sessionConfig != null, + "Neo4jIO.writeUnwindWithResults().withSessionConfig(sessionConfig) called with null sessionConfig"); return toBuilder().setSessionConfig(sessionConfig).build(); } public WriteUnwindWithResults withBatchSize(long batchSize) { checkArgument( - batchSize > 0, "Neo4jIO.writeUnwindWithResults().withBatchSize(batchSize) called with batchSize<=0"); + batchSize > 0, + "Neo4jIO.writeUnwindWithResults().withBatchSize(batchSize) called with batchSize<=0"); return withBatchSize(ValueProvider.StaticValueProvider.of(batchSize)); } public WriteUnwindWithResults withBatchSize(ValueProvider batchSize) { checkArgument( - batchSize != null && batchSize.get() >= 0, - "Neo4jIO.writeUnwindWithResults().withBatchSize(batchSize) called with batchSize<=0"); + batchSize != null && batchSize.get() >= 0, + "Neo4jIO.writeUnwindWithResults().withBatchSize(batchSize) called with batchSize<=0"); return toBuilder().setBatchSize(batchSize).build(); } public WriteUnwindWithResults withParametersFunction( - SerializableFunction> parametersFunction) { + SerializableFunction> parametersFunction) { checkArgument( - parametersFunction != null, - "Neo4jIO.writeUnwindWithResults().withParametersFunction(parametersFunction) called with null parametersFunction"); + parametersFunction != null, + "Neo4jIO.writeUnwindWithResults().withParametersFunction(parametersFunction) called with null parametersFunction"); return toBuilder().setParametersFunction(parametersFunction).build(); } @@ -1277,7 +1285,7 @@ public WriteUnwindWithResults withCypherLogging() { public PCollection expand(PCollection input) { final SerializableFunction driverProviderFn = getDriverProviderFn(); final SerializableFunction> parametersFunction = - getParametersFunction(); + getParametersFunction(); SessionConfig sessionConfig = getProvidedValue(getSessionConfig()); if (sessionConfig == null) { sessionConfig = SessionConfig.defaultConfig(); @@ -1309,15 +1317,15 @@ public PCollection expand(PCollection input) { } WriteUnwindWithResultsFn writeFn = - new WriteUnwindWithResultsFn<>( - driverProviderFn, - sessionConfig, - transactionConfig, - cypher, - parametersFunction, - batchSize, - logCypher, - unwindMapName); + new WriteUnwindWithResultsFn<>( + driverProviderFn, + sessionConfig, + transactionConfig, + cypher, + parametersFunction, + batchSize, + logCypher, + unwindMapName); return getOutputPCollection(input, writeFn, AvroCoder.of(Neo4jWriteStats.class)); } @@ -1337,19 +1345,19 @@ public void populateDisplayData(DisplayData.Builder builder) { @AutoValue.Builder abstract static class Builder { abstract Builder setDriverProviderFn( - SerializableFunction driverProviderFn); + SerializableFunction driverProviderFn); abstract Builder setSessionConfig(ValueProvider sessionConfig); abstract Builder setTransactionConfig( - ValueProvider transactionConfig); + ValueProvider transactionConfig); abstract Builder setCypher(ValueProvider cypher); abstract Builder setUnwindMapName(ValueProvider unwindMapName); abstract Builder setParametersFunction( - SerializableFunction> parametersFunction); + SerializableFunction> parametersFunction); abstract Builder setBatchSize(ValueProvider batchSize); @@ -1475,9 +1483,11 @@ public void finishBundle() { } } - private static class WriteUnwindWithResultsFn extends ReadWriteFn { + private static class WriteUnwindWithResultsFn + extends ReadWriteFn { private final @NonNull String cypher; - private final @Nullable SerializableFunction> parametersFunction; + private final @Nullable SerializableFunction> + parametersFunction; private final boolean logCypher; private final long batchSize; private final @NonNull String unwindMapName; @@ -1487,14 +1497,14 @@ private static class WriteUnwindWithResultsFn extends ReadWriteFn> unwindList; private WriteUnwindWithResultsFn( - @NonNull SerializableFunction driverProviderFn, - @NonNull SessionConfig sessionConfig, - @NonNull TransactionConfig transactionConfig, - @NonNull String cypher, - @Nullable SerializableFunction> parametersFunction, - long batchSize, - boolean logCypher, - String unwindMapName) { + @NonNull SerializableFunction driverProviderFn, + @NonNull SessionConfig sessionConfig, + @NonNull TransactionConfig transactionConfig, + @NonNull String cypher, + @Nullable SerializableFunction> parametersFunction, + long batchSize, + boolean logCypher, + String unwindMapName) { super(driverProviderFn, sessionConfig, transactionConfig); this.cypher = cypher; this.parametersFunction = parametersFunction; @@ -1522,20 +1532,20 @@ public void processElement(@Element ParameterT parameters, ProcessContext contex parametersMap.put(unwindMapName, unwindList); TransactionWork transactionWork = - transaction -> { - Result result = transaction.run(cypher, parametersMap); - Neo4jWriteStats stats = new Neo4jWriteStats(result.consume().counters()); - context.outputWithTimestamp(stats, context.timestamp()); - return null; - }; + transaction -> { + Result result = transaction.run(cypher, parametersMap); + Neo4jWriteStats stats = new Neo4jWriteStats(result.consume().counters()); + context.outputWithTimestamp(stats, context.timestamp()); + return null; + }; if (logCypher && !loggingDone) { String parametersString = getParametersString(parametersMap); LOG.info( - "Starting a write transaction for unwind statement cypher: " - + cypher - + ", parameters: " - + parametersString); + "Starting a write transaction for unwind statement cypher: " + + cypher + + ", parameters: " + + parametersString); loggingDone = true; } @@ -1546,7 +1556,8 @@ public void processElement(@Element ParameterT parameters, ProcessContext contex driverSession.session.writeTransaction(transactionWork, transactionConfig); } catch (Exception e) { throw new RuntimeException( - "Error writing " + unwindList.size() + " rows to Neo4j with Cypher: " + cypher, e); + "Error writing " + unwindList.size() + " rows to Neo4j with Cypher: " + cypher, + e); } } @@ -1563,22 +1574,22 @@ public void finishBundle(FinishBundleContext context) { parametersMap.put(unwindMapName, unwindList); TransactionWork transactionWork = - transaction -> { - Result result = transaction.run(cypher, parametersMap); - Neo4jWriteStats stats = new Neo4jWriteStats(result.consume().counters()); - Instant timestamp = Instant.now(); - BoundedWindow window = GlobalWindow.INSTANCE; - context.output(stats, timestamp, window); - return null; - }; + transaction -> { + Result result = transaction.run(cypher, parametersMap); + Neo4jWriteStats stats = new Neo4jWriteStats(result.consume().counters()); + Instant timestamp = Instant.now(); + BoundedWindow window = GlobalWindow.INSTANCE; + context.output(stats, timestamp, window); + return null; + }; if (logCypher && !loggingDone) { String parametersString = getParametersString(parametersMap); LOG.info( - "Starting a write transaction for unwind statement cypher: " - + cypher - + ", parameters: " - + parametersString); + "Starting a write transaction for unwind statement cypher: " + + cypher + + ", parameters: " + + parametersString); loggingDone = true; } @@ -1589,7 +1600,7 @@ public void finishBundle(FinishBundleContext context) { driverSession.session.writeTransaction(transactionWork, transactionConfig); } catch (Exception e) { throw new RuntimeException( - "Error writing " + unwindList.size() + " rows to Neo4j with Cypher: " + cypher, e); + "Error writing " + unwindList.size() + " rows to Neo4j with Cypher: " + cypher, e); } }