feat(neo4j): Allow datahub to connect to specific neo4j database (#9179)

Co-authored-by: david-leifker <[email protected]>
datahub-project · Nov 8, 2023 · 399e032 · 399e032
1 parent e73e926
commit 399e032
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 55 deletions.
diff --git a/docker/docker-compose.override.yml b/docker/docker-compose.override.yml
@@ -7,8 +7,12 @@ services:
     environment:
     - DATAHUB_SERVER_TYPE=${DATAHUB_SERVER_TYPE:-quickstart}
     - DATAHUB_TELEMETRY_ENABLED=${DATAHUB_TELEMETRY_ENABLED:-true}
+    - GRAPH_SERVICE_IMPL=${GRAPH_SERVICE_IMPL:-elasticsearch}
     volumes:
     - ${HOME}/.datahub/plugins:/etc/datahub/plugins
+  datahub-upgrade:
+    environment:
+      - GRAPH_SERVICE_IMPL=${GRAPH_SERVICE_IMPL:-elasticsearch}
   mysql-setup:
     container_name: mysql-setup
     hostname: mysql-setup

diff --git a/docker/quickstart/docker-compose-m1.quickstart.yml b/docker/quickstart/docker-compose-m1.quickstart.yml
@@ -81,32 +81,32 @@ services:
     - DATAHUB_SERVER_TYPE=${DATAHUB_SERVER_TYPE:-quickstart}
     - DATAHUB_TELEMETRY_ENABLED=${DATAHUB_TELEMETRY_ENABLED:-true}
     - DATAHUB_UPGRADE_HISTORY_KAFKA_CONSUMER_GROUP_ID=generic-duhe-consumer-job-client-gms
-    - EBEAN_DATASOURCE_USERNAME=datahub
-    - EBEAN_DATASOURCE_PASSWORD=datahub
+    - EBEAN_DATASOURCE_DRIVER=com.mysql.jdbc.Driver
     - EBEAN_DATASOURCE_HOST=mysql:3306
+    - EBEAN_DATASOURCE_PASSWORD=datahub
     - EBEAN_DATASOURCE_URL=jdbc:mysql://mysql:3306/datahub?verifyServerCertificate=false&useSSL=true&useUnicode=yes&characterEncoding=UTF-8&enabledTLSProtocols=TLSv1.2
-    - EBEAN_DATASOURCE_DRIVER=com.mysql.jdbc.Driver
-    - KAFKA_BOOTSTRAP_SERVER=broker:29092
-    - KAFKA_SCHEMAREGISTRY_URL=http://schema-registry:8081
+    - EBEAN_DATASOURCE_USERNAME=datahub
     - ELASTICSEARCH_HOST=elasticsearch
-    - ELASTICSEARCH_PORT=9200
-    - ES_BULK_REFRESH_POLICY=WAIT_UNTIL
-    - ELASTICSEARCH_INDEX_BUILDER_SETTINGS_REINDEX=true
     - ELASTICSEARCH_INDEX_BUILDER_MAPPINGS_REINDEX=true
-    - NEO4J_HOST=http://neo4j:7474
-    - NEO4J_URI=bolt://neo4j
-    - NEO4J_USERNAME=neo4j
-    - NEO4J_PASSWORD=datahub
-    - JAVA_OPTS=-Xms1g -Xmx1g
-    - GRAPH_SERVICE_DIFF_MODE_ENABLED=true
-    - GRAPH_SERVICE_IMPL=neo4j
+    - ELASTICSEARCH_INDEX_BUILDER_SETTINGS_REINDEX=true
+    - ELASTICSEARCH_PORT=9200
     - ENTITY_REGISTRY_CONFIG_PATH=/datahub/datahub-gms/resources/entity-registry.yml
     - ENTITY_SERVICE_ENABLE_RETENTION=true
+    - ES_BULK_REFRESH_POLICY=WAIT_UNTIL
+    - GRAPH_SERVICE_DIFF_MODE_ENABLED=true
+    - GRAPH_SERVICE_IMPL=${GRAPH_SERVICE_IMPL:-elasticsearch}
+    - JAVA_OPTS=-Xms1g -Xmx1g
+    - KAFKA_BOOTSTRAP_SERVER=broker:29092
+    - KAFKA_SCHEMAREGISTRY_URL=http://schema-registry:8081
     - MAE_CONSUMER_ENABLED=true
     - MCE_CONSUMER_ENABLED=true
+    - METADATA_SERVICE_AUTH_ENABLED=false
+    - NEO4J_HOST=http://neo4j:7474
+    - NEO4J_PASSWORD=datahub
+    - NEO4J_URI=bolt://neo4j
+    - NEO4J_USERNAME=neo4j
     - PE_CONSUMER_ENABLED=true
     - UI_INGESTION_ENABLED=true
-    - METADATA_SERVICE_AUTH_ENABLED=false
     healthcheck:
       interval: 1s
       retries: 3
@@ -134,23 +134,23 @@ services:
       neo4j:
         condition: service_healthy
     environment:
-    - EBEAN_DATASOURCE_USERNAME=datahub
-    - EBEAN_DATASOURCE_PASSWORD=datahub
+    - BACKFILL_BROWSE_PATHS_V2=true
+    - DATAHUB_GMS_HOST=datahub-gms
+    - DATAHUB_GMS_PORT=8080
+    - EBEAN_DATASOURCE_DRIVER=com.mysql.jdbc.Driver
     - EBEAN_DATASOURCE_HOST=mysql:3306
+    - EBEAN_DATASOURCE_PASSWORD=datahub
     - EBEAN_DATASOURCE_URL=jdbc:mysql://mysql:3306/datahub?verifyServerCertificate=false&useSSL=true&useUnicode=yes&characterEncoding=UTF-8
-    - EBEAN_DATASOURCE_DRIVER=com.mysql.jdbc.Driver
-    - KAFKA_BOOTSTRAP_SERVER=broker:29092
-    - KAFKA_SCHEMAREGISTRY_URL=http://schema-registry:8081
+    - EBEAN_DATASOURCE_USERNAME=datahub
+    - ELASTICSEARCH_BUILD_INDICES_CLONE_INDICES=false
     - ELASTICSEARCH_HOST=elasticsearch
-    - ELASTICSEARCH_PORT=9200
     - ELASTICSEARCH_INDEX_BUILDER_MAPPINGS_REINDEX=true
     - ELASTICSEARCH_INDEX_BUILDER_SETTINGS_REINDEX=true
-    - ELASTICSEARCH_BUILD_INDICES_CLONE_INDICES=false
-    - GRAPH_SERVICE_IMPL=elasticsearch
-    - DATAHUB_GMS_HOST=datahub-gms
-    - DATAHUB_GMS_PORT=8080
+    - ELASTICSEARCH_PORT=9200
     - ENTITY_REGISTRY_CONFIG_PATH=/datahub/datahub-gms/resources/entity-registry.yml
-    - BACKFILL_BROWSE_PATHS_V2=true
+    - GRAPH_SERVICE_IMPL=${GRAPH_SERVICE_IMPL:-elasticsearch}
+    - KAFKA_BOOTSTRAP_SERVER=broker:29092
+    - KAFKA_SCHEMAREGISTRY_URL=http://schema-registry:8081
     - REPROCESS_DEFAULT_BROWSE_PATHS_V2=false
     hostname: datahub-upgrade
     image: ${DATAHUB_UPGRADE_IMAGE:-acryldata/datahub-upgrade}:${DATAHUB_VERSION:-head}

diff --git a/docker/quickstart/docker-compose.quickstart.yml b/docker/quickstart/docker-compose.quickstart.yml
@@ -81,32 +81,32 @@ services:
     - DATAHUB_SERVER_TYPE=${DATAHUB_SERVER_TYPE:-quickstart}
     - DATAHUB_TELEMETRY_ENABLED=${DATAHUB_TELEMETRY_ENABLED:-true}
     - DATAHUB_UPGRADE_HISTORY_KAFKA_CONSUMER_GROUP_ID=generic-duhe-consumer-job-client-gms
-    - EBEAN_DATASOURCE_USERNAME=datahub
-    - EBEAN_DATASOURCE_PASSWORD=datahub
+    - EBEAN_DATASOURCE_DRIVER=com.mysql.jdbc.Driver
     - EBEAN_DATASOURCE_HOST=mysql:3306
+    - EBEAN_DATASOURCE_PASSWORD=datahub
     - EBEAN_DATASOURCE_URL=jdbc:mysql://mysql:3306/datahub?verifyServerCertificate=false&useSSL=true&useUnicode=yes&characterEncoding=UTF-8&enabledTLSProtocols=TLSv1.2
-    - EBEAN_DATASOURCE_DRIVER=com.mysql.jdbc.Driver
-    - KAFKA_BOOTSTRAP_SERVER=broker:29092
-    - KAFKA_SCHEMAREGISTRY_URL=http://schema-registry:8081
+    - EBEAN_DATASOURCE_USERNAME=datahub
     - ELASTICSEARCH_HOST=elasticsearch
-    - ELASTICSEARCH_PORT=9200
-    - ES_BULK_REFRESH_POLICY=WAIT_UNTIL
-    - ELASTICSEARCH_INDEX_BUILDER_SETTINGS_REINDEX=true
     - ELASTICSEARCH_INDEX_BUILDER_MAPPINGS_REINDEX=true
-    - NEO4J_HOST=http://neo4j:7474
-    - NEO4J_URI=bolt://neo4j
-    - NEO4J_USERNAME=neo4j
-    - NEO4J_PASSWORD=datahub
-    - JAVA_OPTS=-Xms1g -Xmx1g
-    - GRAPH_SERVICE_DIFF_MODE_ENABLED=true
-    - GRAPH_SERVICE_IMPL=neo4j
+    - ELASTICSEARCH_INDEX_BUILDER_SETTINGS_REINDEX=true
+    - ELASTICSEARCH_PORT=9200
     - ENTITY_REGISTRY_CONFIG_PATH=/datahub/datahub-gms/resources/entity-registry.yml
     - ENTITY_SERVICE_ENABLE_RETENTION=true
+    - ES_BULK_REFRESH_POLICY=WAIT_UNTIL
+    - GRAPH_SERVICE_DIFF_MODE_ENABLED=true
+    - GRAPH_SERVICE_IMPL=${GRAPH_SERVICE_IMPL:-elasticsearch}
+    - JAVA_OPTS=-Xms1g -Xmx1g
+    - KAFKA_BOOTSTRAP_SERVER=broker:29092
+    - KAFKA_SCHEMAREGISTRY_URL=http://schema-registry:8081
     - MAE_CONSUMER_ENABLED=true
     - MCE_CONSUMER_ENABLED=true
+    - METADATA_SERVICE_AUTH_ENABLED=false
+    - NEO4J_HOST=http://neo4j:7474
+    - NEO4J_PASSWORD=datahub
+    - NEO4J_URI=bolt://neo4j
+    - NEO4J_USERNAME=neo4j
     - PE_CONSUMER_ENABLED=true
     - UI_INGESTION_ENABLED=true
-    - METADATA_SERVICE_AUTH_ENABLED=false
     healthcheck:
       interval: 1s
       retries: 3
@@ -134,23 +134,23 @@ services:
       neo4j:
         condition: service_healthy
     environment:
-    - EBEAN_DATASOURCE_USERNAME=datahub
-    - EBEAN_DATASOURCE_PASSWORD=datahub
+    - BACKFILL_BROWSE_PATHS_V2=true
+    - DATAHUB_GMS_HOST=datahub-gms
+    - DATAHUB_GMS_PORT=8080
+    - EBEAN_DATASOURCE_DRIVER=com.mysql.jdbc.Driver
     - EBEAN_DATASOURCE_HOST=mysql:3306
+    - EBEAN_DATASOURCE_PASSWORD=datahub
     - EBEAN_DATASOURCE_URL=jdbc:mysql://mysql:3306/datahub?verifyServerCertificate=false&useSSL=true&useUnicode=yes&characterEncoding=UTF-8
-    - EBEAN_DATASOURCE_DRIVER=com.mysql.jdbc.Driver
-    - KAFKA_BOOTSTRAP_SERVER=broker:29092
-    - KAFKA_SCHEMAREGISTRY_URL=http://schema-registry:8081
+    - EBEAN_DATASOURCE_USERNAME=datahub
+    - ELASTICSEARCH_BUILD_INDICES_CLONE_INDICES=false
     - ELASTICSEARCH_HOST=elasticsearch
-    - ELASTICSEARCH_PORT=9200
     - ELASTICSEARCH_INDEX_BUILDER_MAPPINGS_REINDEX=true
     - ELASTICSEARCH_INDEX_BUILDER_SETTINGS_REINDEX=true
-    - ELASTICSEARCH_BUILD_INDICES_CLONE_INDICES=false
-    - GRAPH_SERVICE_IMPL=elasticsearch
-    - DATAHUB_GMS_HOST=datahub-gms
-    - DATAHUB_GMS_PORT=8080
+    - ELASTICSEARCH_PORT=9200
     - ENTITY_REGISTRY_CONFIG_PATH=/datahub/datahub-gms/resources/entity-registry.yml
-    - BACKFILL_BROWSE_PATHS_V2=true
+    - GRAPH_SERVICE_IMPL=${GRAPH_SERVICE_IMPL:-elasticsearch}
+    - KAFKA_BOOTSTRAP_SERVER=broker:29092
+    - KAFKA_SCHEMAREGISTRY_URL=http://schema-registry:8081
     - REPROCESS_DEFAULT_BROWSE_PATHS_V2=false
     hostname: datahub-upgrade
     image: ${DATAHUB_UPGRADE_IMAGE:-acryldata/datahub-upgrade}:${DATAHUB_VERSION:-head}

diff --git a/metadata-service/configuration/src/main/resources/application.yml b/metadata-service/configuration/src/main/resources/application.yml
@@ -251,6 +251,7 @@ neo4j:
   username: ${NEO4J_USERNAME:neo4j}
   password: ${NEO4J_PASSWORD:datahub}
   uri: ${NEO4J_URI:bolt://localhost}
+  database: ${NEO4J_DATABASE:graph.db}
   maxConnectionPoolSize: ${NEO4J_MAX_CONNECTION_POOL_SIZE:100}
   maxConnectionAcquisitionTimeout: ${NEO4J_MAX_CONNECTION_ACQUISITION_TIMEOUT_IN_SECONDS:60}
   maxConnectionLifetimeInSeconds: ${NEO4j_MAX_CONNECTION_LIFETIME_IN_SECONDS:3600}

diff --git a/...ice/factories/src/main/java/com/linkedin/gms/factory/common/Neo4jGraphServiceFactory.java b/...ice/factories/src/main/java/com/linkedin/gms/factory/common/Neo4jGraphServiceFactory.java
@@ -6,8 +6,10 @@
 import com.linkedin.metadata.models.registry.EntityRegistry;
 import javax.annotation.Nonnull;
 import org.neo4j.driver.Driver;
+import org.neo4j.driver.SessionConfig;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Import;
@@ -24,10 +26,13 @@ public class Neo4jGraphServiceFactory {
   @Qualifier("entityRegistry")
   private EntityRegistry entityRegistry;
 
+  @Value("${neo4j.database}")
+  private String neo4jDatabase;
+
   @Bean(name = "neo4jGraphService")
   @Nonnull
   protected Neo4jGraphService getInstance() {
     LineageRegistry lineageRegistry = new LineageRegistry(entityRegistry);
-    return new Neo4jGraphService(lineageRegistry, neo4jDriver);
+    return new Neo4jGraphService(lineageRegistry, neo4jDriver, SessionConfig.forDatabase(neo4jDatabase));
   }
 }