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

Co-authored-by: david-leifker <114954101+david-leifker@users.noreply.github.com>
2025-12-11 01:56:58 +00:00 · 2023-11-08 22:32:13 +05:30 · 2023-11-08 22:32:13 +05:30 · 399e032dfa
commit 399e032dfa
parent e73e926999
5 changed files with 73 additions and 63 deletions
--- 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
--- 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_HOST=mysql:3306
-    - 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_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_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
-    - EBEAN_DATASOURCE_HOST=mysql:3306
-    - 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
-    - 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
+    - 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_USERNAME=datahub
+    - ELASTICSEARCH_BUILD_INDICES_CLONE_INDICES=false
+    - ELASTICSEARCH_HOST=elasticsearch
+    - ELASTICSEARCH_INDEX_BUILDER_MAPPINGS_REINDEX=true
+    - ELASTICSEARCH_INDEX_BUILDER_SETTINGS_REINDEX=true
+    - 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}
--- 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_HOST=mysql:3306
-    - 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_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_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
-    - EBEAN_DATASOURCE_HOST=mysql:3306
-    - 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
-    - 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
+    - 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_USERNAME=datahub
+    - ELASTICSEARCH_BUILD_INDICES_CLONE_INDICES=false
+    - ELASTICSEARCH_HOST=elasticsearch
+    - ELASTICSEARCH_INDEX_BUILDER_MAPPINGS_REINDEX=true
+    - ELASTICSEARCH_INDEX_BUILDER_SETTINGS_REINDEX=true
+    - 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}
--- 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}
--- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/Neo4jGraphServiceFactory.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/Neo4jGraphServiceFactory.java
@ -6,8 +6,10 @@ import com.linkedin.metadata.graph.neo4j.Neo4jGraphService;
 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));
  }
 }