From 788b07ccbf1dcbf2bff12fa3212f657dfb3bf9df Mon Sep 17 00:00:00 2001
From: Lasse Westh-Nielsen <lassewesth@gmail.com>
Date: Fri, 2 Aug 2024 09:37:44 +0200
Subject: [PATCH] create private modules to host system monitor

---
 .../extension/OpenGraphDataScienceExtension.java   | 14 ++++++++++----
 .../ExporterBuildersProviderService.java           |  4 ++--
 .../GraphDataScienceExtensionBuilder.java          | 13 +++++++------
 .../integration/GraphDataScienceProvider.java      | 10 +++++-----
 .../GraphDataScienceProviderFactory.java           | 14 +++++++-------
 5 files changed, 31 insertions(+), 24 deletions(-)

diff --git a/procedures/extension/src/main/java/org/neo4j/gds/extension/OpenGraphDataScienceExtension.java b/procedures/extension/src/main/java/org/neo4j/gds/extension/OpenGraphDataScienceExtension.java
index ac9b6c1dd0..ce5bf86728 100644
--- a/procedures/extension/src/main/java/org/neo4j/gds/extension/OpenGraphDataScienceExtension.java
+++ b/procedures/extension/src/main/java/org/neo4j/gds/extension/OpenGraphDataScienceExtension.java
@@ -57,18 +57,24 @@ public Lifecycle newInstance(ExtensionContext extensionContext, Dependencies dep
         var modelCatalog = new OpenModelCatalogProvider().get(null);
 
         var graphSageModelRepository = new DisableGraphSageModelRepository(); // no model storing in OpenGDS
-        return GraphDataScienceExtensionBuilder.create(
+
+        var neo4jConfiguration = dependencies.config();
+
+        var graphDataScienceExtensionBuilderAndAssociatedProducts = GraphDataScienceExtensionBuilder.create(
             log,
-            dependencies.config(),
+            neo4jConfiguration,
             dependencies.globalProcedures(),
             Optional.empty(), // no extra checks in OpenGDS
             Optional.empty(), // no extra checks in OpenGDS
             exporterBuildersProviderService,
             metricsFacade,
             modelCatalog,
-            dependencies.config(),
             graphSageModelRepository
-        ).build();
+        );
+
+        var graphDataScienceExtensionBuilder = graphDataScienceExtensionBuilderAndAssociatedProducts.getLeft();
+
+        return graphDataScienceExtensionBuilder.build();
     }
 
     public interface Dependencies {
diff --git a/procedures/facade/src/main/java/org/neo4j/gds/procedures/ExporterBuildersProviderService.java b/procedures/facade/src/main/java/org/neo4j/gds/procedures/ExporterBuildersProviderService.java
index daeff136f6..92714a3b23 100644
--- a/procedures/facade/src/main/java/org/neo4j/gds/procedures/ExporterBuildersProviderService.java
+++ b/procedures/facade/src/main/java/org/neo4j/gds/procedures/ExporterBuildersProviderService.java
@@ -19,10 +19,10 @@
  */
 package org.neo4j.gds.procedures;
 
-import org.neo4j.configuration.Config;
 import org.neo4j.gds.core.write.ExportBuildersProvider;
 import org.neo4j.graphdb.GraphDatabaseService;
+import org.neo4j.graphdb.config.Configuration;
 
 public interface ExporterBuildersProviderService {
-    ExportBuildersProvider identifyExportBuildersProvider(GraphDatabaseService graphDatabaseService, Config config);
+    ExportBuildersProvider identifyExportBuildersProvider(GraphDatabaseService graphDatabaseService, Configuration neo4jConfiguration);
 }
diff --git a/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/GraphDataScienceExtensionBuilder.java b/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/GraphDataScienceExtensionBuilder.java
index 06aa51122e..18adce4688 100644
--- a/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/GraphDataScienceExtensionBuilder.java
+++ b/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/GraphDataScienceExtensionBuilder.java
@@ -19,7 +19,7 @@
  */
 package org.neo4j.gds.procedures.integration;
 
-import org.neo4j.configuration.Config;
+import org.apache.commons.lang3.tuple.Pair;
 import org.neo4j.gds.applications.algorithms.embeddings.GraphSageModelRepository;
 import org.neo4j.gds.applications.algorithms.machinery.AlgorithmProcessingTemplate;
 import org.neo4j.gds.applications.graphstorecatalog.CatalogBusinessFacade;
@@ -80,7 +80,7 @@ private GraphDataScienceExtensionBuilder(
         ModelCatalog modelCatalog,
         TaskStoreService taskStoreService,
         TaskRegistryFactoryService taskRegistryFactoryService,
-        Boolean useMaxMemoryEstimation,
+        boolean useMaxMemoryEstimation,
         UserLogServices userLogServices,
         Lifecycle gcListener
     ) {
@@ -100,7 +100,7 @@ private GraphDataScienceExtensionBuilder(
      * We want to build a GDS, we receive a few customisations and are able to read configuration,
      * and all the rest of the machinery goes here
      */
-    public static GraphDataScienceExtensionBuilder create(
+    public static Pair<GraphDataScienceExtensionBuilder, TaskStoreService> create(
         Log log,
         Configuration neo4jConfiguration,
         GlobalProcedures globalProcedures,
@@ -109,7 +109,6 @@ public static GraphDataScienceExtensionBuilder create(
         ExporterBuildersProviderService exporterBuildersProviderService,
         MetricsFacade metricsFacade,
         ModelCatalog modelCatalog,
-        Config config,
         GraphSageModelRepository graphSageModelRepository
     ) {
         // Read some configuration used to select behaviour
@@ -153,11 +152,11 @@ public static GraphDataScienceExtensionBuilder create(
             memoryGauge,
             metricsFacade,
             modelCatalog,
-            config,
+            neo4jConfiguration,
             graphSageModelRepository
         );
 
-        return new GraphDataScienceExtensionBuilder(
+        var graphDataScienceExtensionBuilder = new GraphDataScienceExtensionBuilder(
             log,
             componentRegistration,
             graphDataScienceProviderFactory,
@@ -169,6 +168,8 @@ public static GraphDataScienceExtensionBuilder create(
             userLogServices,
             gcListener
         );
+
+        return Pair.of(graphDataScienceExtensionBuilder, taskStoreService);
     }
 
     /**
diff --git a/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/GraphDataScienceProvider.java b/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/GraphDataScienceProvider.java
index 50ca7e7621..8aac310bc6 100644
--- a/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/GraphDataScienceProvider.java
+++ b/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/GraphDataScienceProvider.java
@@ -19,7 +19,6 @@
  */
 package org.neo4j.gds.procedures.integration;
 
-import org.neo4j.configuration.Config;
 import org.neo4j.function.ThrowingFunction;
 import org.neo4j.gds.applications.algorithms.embeddings.GraphSageModelRepository;
 import org.neo4j.gds.applications.algorithms.machinery.AlgorithmProcessingTemplate;
@@ -50,6 +49,7 @@
 import org.neo4j.gds.procedures.UserAccessor;
 import org.neo4j.gds.procedures.UserLogServices;
 import org.neo4j.graphdb.GraphDatabaseService;
+import org.neo4j.graphdb.config.Configuration;
 import org.neo4j.internal.kernel.api.exceptions.ProcedureException;
 import org.neo4j.kernel.api.procedure.Context;
 
@@ -83,7 +83,7 @@ public class GraphDataScienceProvider implements ThrowingFunction<Context, Graph
     private final TaskRegistryFactoryService taskRegistryFactoryService;
     private final TaskStoreService taskStoreService;
     private final UserLogServices userLogServices;
-    private final Config config;
+    private final Configuration neo4jConfiguration;
     private final ModelCatalog modelCatalog;
     private final GraphSageModelRepository graphSageModelRepository;
 
@@ -104,7 +104,7 @@ public class GraphDataScienceProvider implements ThrowingFunction<Context, Graph
         TaskRegistryFactoryService taskRegistryFactoryService,
         TaskStoreService taskStoreService,
         UserLogServices userLogServices,
-        Config config,
+        Configuration neo4jConfiguration,
         ModelCatalog modelCatalog,
         GraphSageModelRepository graphSageModelRepository
     ) {
@@ -124,7 +124,7 @@ public class GraphDataScienceProvider implements ThrowingFunction<Context, Graph
         this.taskRegistryFactoryService = taskRegistryFactoryService;
         this.taskStoreService = taskStoreService;
         this.userLogServices = userLogServices;
-        this.config = config;
+        this.neo4jConfiguration = neo4jConfiguration;
         this.modelCatalog = modelCatalog;
         this.graphSageModelRepository = graphSageModelRepository;
     }
@@ -198,7 +198,7 @@ public GraphDataScienceProcedures apply(Context context) throws ProcedureExcepti
     private WriteContext constructWriteContext(Context context, GraphDatabaseService graphDatabaseService) {
         var exportBuildersProvider = exporterBuildersProviderService.identifyExportBuildersProvider(
             graphDatabaseService,
-            config
+            neo4jConfiguration
         );
 
         var exporterContext = new ExporterContext.ProcedureContextWrapper(context);
diff --git a/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/GraphDataScienceProviderFactory.java b/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/GraphDataScienceProviderFactory.java
index c04ba6517c..178a1bf050 100644
--- a/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/GraphDataScienceProviderFactory.java
+++ b/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/GraphDataScienceProviderFactory.java
@@ -19,7 +19,6 @@
  */
 package org.neo4j.gds.procedures.integration;
 
-import org.neo4j.configuration.Config;
 import org.neo4j.gds.applications.algorithms.embeddings.GraphSageModelRepository;
 import org.neo4j.gds.applications.algorithms.machinery.AlgorithmProcessingTemplate;
 import org.neo4j.gds.applications.algorithms.machinery.DefaultMemoryGuard;
@@ -37,6 +36,7 @@
 import org.neo4j.gds.procedures.ExporterBuildersProviderService;
 import org.neo4j.gds.procedures.TaskRegistryFactoryService;
 import org.neo4j.gds.procedures.UserLogServices;
+import org.neo4j.graphdb.config.Configuration;
 
 import java.util.Optional;
 import java.util.function.Function;
@@ -64,7 +64,7 @@ final class GraphDataScienceProviderFactory {
     private final MemoryGauge memoryGauge;
     private final MetricsFacade metricsFacade;
     private final ModelCatalog modelCatalog;
-    private final Config config;
+    private final Configuration neo4jConfiguration;
     private final GraphSageModelRepository graphSageModelRepository;
 
     private GraphDataScienceProviderFactory(
@@ -75,7 +75,7 @@ private GraphDataScienceProviderFactory(
         MemoryGauge memoryGauge,
         MetricsFacade metricsFacade,
         ModelCatalog modelCatalog,
-        Config config,
+        Configuration neo4jConfiguration,
         GraphSageModelRepository graphSageModelRepository
     ) {
         this.log = log;
@@ -85,7 +85,7 @@ private GraphDataScienceProviderFactory(
         this.memoryGauge = memoryGauge;
         this.metricsFacade = metricsFacade;
         this.modelCatalog = modelCatalog;
-        this.config = config;
+        this.neo4jConfiguration = neo4jConfiguration;
         this.graphSageModelRepository = graphSageModelRepository;
     }
 
@@ -120,7 +120,7 @@ GraphDataScienceProvider createGraphDataScienceProvider(
             taskRegistryFactoryService,
             taskStoreService,
             userLogServices,
-            config,
+            neo4jConfiguration,
             modelCatalog,
             graphSageModelRepository
         );
@@ -134,7 +134,7 @@ static GraphDataScienceProviderFactory create(
         MemoryGauge memoryGauge,
         MetricsFacade metricsFacade,
         ModelCatalog modelCatalog,
-        Config config,
+        Configuration neo4jConfiguration,
         GraphSageModelRepository graphSageModelRepository
     ) {
         return new GraphDataScienceProviderFactory(
@@ -145,7 +145,7 @@ static GraphDataScienceProviderFactory create(
             memoryGauge,
             metricsFacade,
             modelCatalog,
-            config,
+            neo4jConfiguration,
             graphSageModelRepository
         );
     }