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..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 @@ -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,6 +44,8 @@ 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; @@ -49,6 +53,7 @@ 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 +65,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 +189,64 @@ 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,209 @@ 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 +1482,131 @@ 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; + } + } + } }