Skip to content

Commit

Permalink
feat(neo4j): Extend neo4j writer dao to write to a given db (#94)
Browse files Browse the repository at this point in the history
  • Loading branch information
Kerem Sahin authored Apr 13, 2021
1 parent 8a7846a commit 1a66e49
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 11 deletions.
33 changes: 26 additions & 7 deletions dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/internal/Neo4jGraphWriterDAO.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,15 @@
import java.util.Set;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.time.StopWatch;
import org.neo4j.driver.Driver;
import org.neo4j.driver.Record;
import org.neo4j.driver.Session;
import org.neo4j.driver.SessionConfig;
import org.neo4j.driver.exceptions.Neo4jException;

import static com.linkedin.metadata.dao.Neo4jUtil.*;
Expand Down Expand Up @@ -123,17 +125,33 @@ public void onRelationshipsRemoved(int relationshipCount, long updateTimeMs, int

private static final int MAX_TRANSACTION_RETRY = 3;
private final Driver _driver;
private SessionConfig _sessionConfig;
private static Map<String, String> _urnToEntityMap = null;
private DelegateMetricListener _metricListener = new DelegateMetricListener();

public Neo4jGraphWriterDAO(@Nonnull Driver driver) {
this._driver = driver;
buildUrnToEntityMap(getAllEntities());
this(driver, SessionConfig.defaultConfig());
}

/**
* WARNING: Do NOT use this! This is not tested yet.
* Multi-DB support comes with Neo4j 4+.
* Although DAO works with Neo4j 4+, we can't bump Neo4j test harness to 4+ to test this because it needs Java 11
* And Java 11 build is blocked by ES7 migration.
*/
public Neo4jGraphWriterDAO(@Nonnull Driver driver, @Nonnull String databaseName) {
this(driver, SessionConfig.forDatabase(databaseName));
}

public Neo4jGraphWriterDAO(@Nonnull Driver driver, @Nonnull SessionConfig sessionConfig) {
this(driver, sessionConfig, getAllEntities());
}

/* Should only be sed for testing */
public Neo4jGraphWriterDAO(@Nonnull Driver driver, @Nonnull Set<Class<? extends RecordTemplate>> allEntities) {
/* Should only be used for testing */
public Neo4jGraphWriterDAO(@Nonnull Driver driver, @Nonnull SessionConfig sessionConfig,
@Nonnull Set<Class<? extends RecordTemplate>> allEntities) {
this._driver = driver;
this._sessionConfig = sessionConfig;
buildUrnToEntityMap(allEntities);
}

Expand Down Expand Up @@ -222,7 +240,7 @@ private ExecutionResult executeStatements(@Nonnull List<Statement> statements) {
final StopWatch stopWatch = new StopWatch();
stopWatch.start();
Exception lastException;
try (final Session session = _driver.session()) {
try (final Session session = _driver.session(_sessionConfig)) {
do {
try {
session.writeTransaction(tx -> {
Expand Down Expand Up @@ -255,7 +273,7 @@ private ExecutionResult executeStatements(@Nonnull List<Statement> statements) {
*/
@Nonnull
private List<Record> runQuery(@Nonnull Statement statement) {
try (final Session session = _driver.session()) {
try (final Session session = _driver.session(_sessionConfig)) {
return session.run(statement.getCommandText(), statement.getParams()).list();
}
}
Expand Down Expand Up @@ -483,8 +501,9 @@ private <RELATIONSHIP extends RecordTemplate> Statement removeEdge(@Nonnull RELA
return buildStatement(statement, params);
}

// visible for testing
@Nonnull
private Statement buildStatement(@Nonnull String queryTemplate, @Nonnull Map<String, Object> params) {
Statement buildStatement(@Nonnull String queryTemplate, @Nonnull Map<String, Object> params) {
for (Map.Entry<String, Object> entry : params.entrySet()) {
String k = entry.getKey();
Object v = entry.getValue();
Expand Down
3 changes: 2 additions & 1 deletion dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/Neo4jQueryDAOTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import org.neo4j.driver.Driver;
import org.neo4j.driver.GraphDatabase;
import org.neo4j.driver.Record;
import org.neo4j.driver.SessionConfig;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;
Expand All @@ -49,7 +50,7 @@ public void init() {

final Driver driver = GraphDatabase.driver(_serverBuilder.boltURI());
_dao = new Neo4jQueryDAO(driver);
_writer = new Neo4jGraphWriterDAO(driver, TestUtils.getAllTestEntities());
_writer = new Neo4jGraphWriterDAO(driver, SessionConfig.defaultConfig(), TestUtils.getAllTestEntities());
}

@AfterMethod
Expand Down
24 changes: 21 additions & 3 deletions ...l/neo4j-dao/src/test/java/com/linkedin/metadata/dao/internal/Neo4jGraphWriterDAOTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import com.linkedin.metadata.dao.BaseQueryDAO;
import com.linkedin.metadata.dao.Neo4jQueryDAO;
import com.linkedin.metadata.dao.Neo4jTestServerBuilder;
import com.linkedin.metadata.dao.utils.Statement;
import com.linkedin.metadata.query.Criterion;
import com.linkedin.metadata.query.CriterionArray;
import com.linkedin.metadata.query.Filter;
Expand All @@ -13,14 +14,17 @@
import com.linkedin.testing.TestUtils;
import com.linkedin.testing.urn.BarUrn;
import com.linkedin.testing.urn.FooUrn;

import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import javax.annotation.Nonnull;
import org.neo4j.driver.Driver;
import org.neo4j.driver.GraphDatabase;
import org.neo4j.driver.SessionConfig;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;
Expand All @@ -34,6 +38,7 @@
public class Neo4jGraphWriterDAOTest {

private Neo4jTestServerBuilder _serverBuilder;
private Driver _driver;
private Neo4jGraphWriterDAO _dao;
private BaseQueryDAO _queryDao;
private TestMetricListener _testMetricListener;
Expand Down Expand Up @@ -78,9 +83,9 @@ public void init() {
_serverBuilder = new Neo4jTestServerBuilder();
_serverBuilder.newServer();
_testMetricListener = new TestMetricListener();
final Driver driver = GraphDatabase.driver(_serverBuilder.boltURI());
_dao = new Neo4jGraphWriterDAO(driver, TestUtils.getAllTestEntities());
_queryDao = new Neo4jQueryDAO(driver);
_driver = GraphDatabase.driver(_serverBuilder.boltURI());
_dao = new Neo4jGraphWriterDAO(_driver, SessionConfig.defaultConfig(), TestUtils.getAllTestEntities());
_queryDao = new Neo4jQueryDAO(_driver);
_dao.addMetricListener(_testMetricListener);
}

Expand Down Expand Up @@ -405,6 +410,19 @@ public void testGetNodeTypeFromUrn() {
assertEquals(_dao.getNodeType(makeFooUrn(1)), getTypeOrEmptyString(EntityFoo.class));
}

@Test
public void testBuildStatement() {
final FooUrn urn = makeFooUrn(0);
final String queryTemplate = "dummy query template";
final Map<String, Object> queryParams = new HashMap<>();
queryParams.put("urn", urn);

final Statement queryStatement = _dao.buildStatement(queryTemplate, queryParams);

assertEquals(queryStatement.getCommandText(), queryTemplate);
assertEquals(queryStatement.getParams().get("urn"), urn.toString());
}

private void assertEntityFoo(@Nonnull Map<String, Object> node, @Nonnull EntityFoo entity) {
assertEquals(node.get("urn"), entity.getUrn().toString());
assertEquals(node.get("value"), entity.getValue());
Expand Down

0 comments on commit 1a66e49

Please sign in to comment.