feat(sourcedb-to-spanner): Implement Neo4j batch data extraction pipeline. #4029#4030
feat(sourcedb-to-spanner): Implement Neo4j batch data extraction pipeline. #4029#4030yogeshtewari wants to merge 2 commits into
Conversation
…line - Register NEO4J as a supported source database dialect. - Add pipeline options for --neo4jUri, --neo4jUser, --neo4jPassword, and --neo4jReadChunkSize. - Add beam-sdks-java-io-neo4j dependency to sourcedb-to-spanner pom.xml. - Implement Neo4jSrcToSpSourceConnector and route dialect options to Neo4j. - Implement virtual schema discovery (Neo4jIoWrapper) mapping elements into GraphNode and GraphEdge virtual tables. - Implement Neo4jTableReader to partition ID ranges and run parallel Cypher query extractions on worker pools. - Implement Neo4jRowMapper, registering Jackson JavaTimeModule to support serializing date/time attributes. - Create unit tests for Neo4jIoWrapper and Neo4jIoWrapperFactory. Closes GoogleCloudPlatform#4029
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces comprehensive support for Neo4j as a source database within the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for Neo4j as a source database dialect in the sourcedb-to-spanner pipeline, introducing new configuration options, connector classes, schema references, and a custom reader to partition and query Neo4j nodes and edges. The review feedback highlights three key issues: a potential NotSerializableException in Neo4jRowMapper due to a non-serializable ObjectMapper field, an issue in Neo4jTableReader where catching all exceptions and returning 0L can silently hide critical connection errors, and a potential NullPointerException in Neo4jIoWrapperFactory from auto-unboxing a nullable Integer chunk size.
| public class Neo4jRowMapper implements Serializable { | ||
|
|
||
| private final SourceSchemaReference sourceSchemaReference; | ||
| private final SourceTableSchema sourceTableSchema; | ||
| private final ObjectMapper objectMapper; | ||
|
|
||
| public Neo4jRowMapper( | ||
| SourceSchemaReference sourceSchemaReference, SourceTableSchema sourceTableSchema) { | ||
| this.sourceSchemaReference = sourceSchemaReference; | ||
| this.sourceTableSchema = sourceTableSchema; | ||
| this.objectMapper = new ObjectMapper().registerModule(new JavaTimeModule()); | ||
| } |
There was a problem hiding this comment.
The ObjectMapper class is not serializable. Since Neo4jRowMapper implements Serializable, having a non-transient, non-static ObjectMapper instance field will cause a NotSerializableException if the mapper is ever serialized. Additionally, instantiating a new ObjectMapper and registering modules is an expensive operation.
To resolve this, declare ObjectMapper as a private static final constant. This ensures thread-safe, performant reuse and avoids serialization issues.
| public class Neo4jRowMapper implements Serializable { | |
| private final SourceSchemaReference sourceSchemaReference; | |
| private final SourceTableSchema sourceTableSchema; | |
| private final ObjectMapper objectMapper; | |
| public Neo4jRowMapper( | |
| SourceSchemaReference sourceSchemaReference, SourceTableSchema sourceTableSchema) { | |
| this.sourceSchemaReference = sourceSchemaReference; | |
| this.sourceTableSchema = sourceTableSchema; | |
| this.objectMapper = new ObjectMapper().registerModule(new JavaTimeModule()); | |
| } | |
| public class Neo4jRowMapper implements Serializable { | |
| private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper().registerModule(new JavaTimeModule()); | |
| private final SourceSchemaReference sourceSchemaReference; | |
| private final SourceTableSchema sourceTableSchema; | |
| public Neo4jRowMapper( | |
| SourceSchemaReference sourceSchemaReference, SourceTableSchema sourceTableSchema) { | |
| this.sourceSchemaReference = sourceSchemaReference; | |
| this.sourceTableSchema = sourceTableSchema; | |
| } |
| private long fetchMaxId(String query) { | ||
| try (Driver driver = GraphDatabase.driver(uri, AuthTokens.basic(user, password)); | ||
| Session session = driver.session()) { | ||
| Result result = session.run(query); | ||
| if (result.hasNext()) { | ||
| Record record = result.next(); | ||
| Value val = record.get(0); | ||
| if (!val.isNull()) { | ||
| return val.asLong(); | ||
| } | ||
| } | ||
| } catch (Exception e) { | ||
| LOG.warn("Error querying max ID from Neo4j (defaulting to 0):", e); | ||
| } | ||
| return 0L; | ||
| } |
There was a problem hiding this comment.
Catching Exception and silently returning 0L with a warning hides critical connection, authentication, or query errors during pipeline submission. This can lead to empty migrations or late failures on workers.
Instead, throw a RuntimeException to fail fast. If the database is empty, max(id(n)) returns null (which is a NullValue in Neo4j), and val.isNull() correctly handles this by returning 0L without throwing an exception. Additionally, log the error message before throwing the exception to ensure the failure is recorded in Dataflow environments.
| private long fetchMaxId(String query) { | |
| try (Driver driver = GraphDatabase.driver(uri, AuthTokens.basic(user, password)); | |
| Session session = driver.session()) { | |
| Result result = session.run(query); | |
| if (result.hasNext()) { | |
| Record record = result.next(); | |
| Value val = record.get(0); | |
| if (!val.isNull()) { | |
| return val.asLong(); | |
| } | |
| } | |
| } catch (Exception e) { | |
| LOG.warn("Error querying max ID from Neo4j (defaulting to 0):", e); | |
| } | |
| return 0L; | |
| } | |
| private long fetchMaxId(String query) { | |
| try (Driver driver = GraphDatabase.driver(uri, AuthTokens.basic(user, password)); | |
| Session session = driver.session()) { | |
| Result result = session.run(query); | |
| if (result.hasNext()) { | |
| Record record = result.next(); | |
| Value val = record.get(0); | |
| if (!val.isNull()) { | |
| return val.asLong(); | |
| } | |
| } | |
| } catch (Exception e) { | |
| LOG.error("Failed to fetch max ID from Neo4j using query: {}", query, e); | |
| throw new RuntimeException("Failed to fetch max ID from Neo4j using query: " + query, e); | |
| } | |
| return 0L; | |
| } |
References
- When throwing exceptions in environments where the caller or global exception handler might not log them (such as certain Dataflow templates), log the error message before throwing the exception to ensure the failure is recorded.
| public static Neo4jIoWrapperFactory fromPipelineOptions(SourceDbToSpannerOptions options) { | ||
| Preconditions.checkArgument( | ||
| options.getSourceDbDialect().equals(SourceDbToSpannerOptions.NEO4J_SOURCE_DIALECT), | ||
| "Unexpected Dialect " + options.getSourceDbDialect() + " for Neo4j Source"); | ||
| Preconditions.checkNotNull(options.getNeo4jUri(), "Neo4j URI must be specified"); | ||
| Preconditions.checkNotNull(options.getNeo4jUser(), "Neo4j Username must be specified"); | ||
| Preconditions.checkNotNull(options.getNeo4jPassword(), "Neo4j Password must be specified"); | ||
|
|
||
| return new Neo4jIoWrapperFactory( | ||
| options.getNeo4jUri(), | ||
| options.getNeo4jUser(), | ||
| options.getNeo4jPassword(), | ||
| options.getNeo4jReadChunkSize()); | ||
| } |
There was a problem hiding this comment.
The options.getNeo4jReadChunkSize() method returns an Integer object, which can be null. Passing it directly to the Neo4jIoWrapperFactory constructor (which expects a primitive int) will cause a NullPointerException due to auto-unboxing.
Add a defensive null check and fall back to the default chunk size of 100000 if it is null.
public static Neo4jIoWrapperFactory fromPipelineOptions(SourceDbToSpannerOptions options) {
Preconditions.checkArgument(
options.getSourceDbDialect().equals(SourceDbToSpannerOptions.NEO4J_SOURCE_DIALECT),
"Unexpected Dialect " + options.getSourceDbDialect() + " for Neo4j Source");
Preconditions.checkNotNull(options.getNeo4jUri(), "Neo4j URI must be specified");
Preconditions.checkNotNull(options.getNeo4jUser(), "Neo4j Username must be specified");
Preconditions.checkNotNull(options.getNeo4jPassword(), "Neo4j Password must be specified");
Integer chunkSize = options.getNeo4jReadChunkSize();
int readChunkSize = (chunkSize != null) ? chunkSize : 100000;
return new Neo4jIoWrapperFactory(
options.getNeo4jUri(),
options.getNeo4jUser(),
options.getNeo4jPassword(),
readChunkSize);
}References
- Ensure that mandatory parameters are validated for nullability at the entry point or in upstream components to prevent NullPointerExceptions (especially during auto-unboxing) in downstream transforms, avoiding the need for redundant null checks.
Codecov Report❌ Patch coverage is ❌ Your patch check has failed because the patch coverage (41.00%) is below the target coverage (80.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #4030 +/- ##
============================================
+ Coverage 55.69% 63.34% +7.64%
+ Complexity 6787 2521 -4266
============================================
Files 1124 543 -581
Lines 68447 30919 -37528
Branches 7721 3436 -4285
============================================
- Hits 38123 19586 -18537
+ Misses 27849 10319 -17530
+ Partials 2475 1014 -1461
🚀 New features to boost your workflow:
|
bharadwaj-aditya
left a comment
bharadwaj-aditya
left a comment
There was a problem hiding this comment.
Couple of major questions here
- Schema seems to be hardcoded. I would assume different labels will go to different spanner tables. How is that handled ?
- It does not look like Neo4jIO is being used here. We should rely on the library if it has all the support we need.
|
|
||
| void setGcsOutputDirectory(String value); | ||
|
|
||
| @TemplateParameter.Text( |
There was a problem hiding this comment.
These should move into the ConfigFileReader. Please don't add source specific parameters to the job
| public Neo4jIoWrapper( | ||
| String neo4jUri, | ||
| String neo4jUser, | ||
| String neo4jPassword, |
There was a problem hiding this comment.
will this work with secret manager as well ?
|
|
||
| SourceSchemaReference schemaReference = | ||
| SourceSchemaReference.ofNeo4j( | ||
| Neo4jSchemaReference.builder().setDatabaseName("neo4j").build()); |
There was a problem hiding this comment.
There can be multiple database / graph names in a neo4j instance. This should come from config.
| .setTableName("GraphNode") | ||
| .setPrimaryKeyColumns(ImmutableList.of("id")) | ||
| .addSourceColumnNameToSourceColumnType( | ||
| "id", new SourceColumnType("STRING", new Long[0], null)) |
There was a problem hiding this comment.
The schema seems hard coded here. There can be multiple label in a database. Ideally this should be fetched by something like db.labels() or some equivalent.
| // 2. Construct Virtual Edge Table Schema | ||
| SourceTableSchema edgeTableSchema = | ||
| SourceTableSchema.builder(typeMapper) | ||
| .setTableName("GraphEdge") |
There was a problem hiding this comment.
please use constants for such parameters where possible.
| public static ImmutableMap<String, UnifiedTypeMapping> getMapping() { | ||
| return ImmutableMap.<String, UnifiedTypeMapping>builder() | ||
| .put("STRING", UnifiedMappingProvider.getMapping(UnifiedMappingProvider.Type.STRING)) | ||
| .put("JSON", UnifiedMappingProvider.getMapping(UnifiedMappingProvider.Type.JSON)) |
There was a problem hiding this comment.
what about other types ?
| import java.io.Serializable; | ||
|
|
||
| /** Describes a bounded node/edge ID query partition task. */ | ||
| public class Neo4jPartitionTask implements Serializable { |
There was a problem hiding this comment.
nit: This seems like a partition definition. Why is this names task ?
| } | ||
| } | ||
| } catch (Exception e) { | ||
| LOG.warn("Error querying max ID from Neo4j (defaulting to 0):", e); |
There was a problem hiding this comment.
we should be throwing an exception here. If we want to ignore some specific exception, please add more specific handling
|
|
||
| List<Neo4jPartitionTask> partitionTasks = new ArrayList<>(); | ||
| for (long i = 0; i <= maxId; i += chunkSize) { | ||
| partitionTasks.add(new Neo4jPartitionTask(elementType, i, i + chunkSize)); |
There was a problem hiding this comment.
Is this the most efficient way to fetch data from neo4j ? It seems like we're traversing all the nodes to find them.
There was a problem hiding this comment.
Can you not use neo4j IO here ?
| } | ||
|
|
||
| public String getName() { | ||
| return "Database." + databaseName(); |
There was a problem hiding this comment.
is it fine ot hardcode this ?
2 participants
Closes #4029