Correctly inject TaskStore instance

Co-authored-by: Ioannis Panagiotas <[email protected]>

Author: vnickolov

progress-tracking/src/main/java/org/neo4j/gds/core/utils/progress/TaskRegistryExtension.java

@@ -54,16 +54,15 @@ public Lifecycle newInstance(ExtensionContext context, TaskRegistryExtension.Dep
         var registry = dependencies.globalProceduresRegistry();
         var enabled = dependencies.config().get(ProgressFeatureSettings.progress_tracking_enabled);
         String databaseName = dependencies.graphDatabaseService().databaseName();
+
         if (enabled) {
+            var taskStoreProvider = new TaskStoreProvider();
             // Use the centrally managed task stores
-            var taskStore = TaskStoreHolder.getTaskStore(databaseName);
-            var taskRegistryFactoryProvider = new TaskRegistryFactoryProvider(taskStore);
-
-            registry.registerComponent(TaskStore.class, ctx -> taskStore, true);
-            registry.registerComponent(TaskRegistryFactory.class, taskRegistryFactoryProvider, true);
+            registry.registerComponent(TaskStore.class, taskStoreProvider, true);
+            registry.registerComponent(TaskRegistryFactory.class, new TaskRegistryFactoryProvider(taskStoreProvider), true);
 
             // hey this is just for tests? TaskRegistryExtensionMultiDBTest breaks if it is missing
-            context.dependencySatisfier().satisfyDependency(taskStore);
+            context.dependencySatisfier().satisfyDependency(TaskStoreHolder.getTaskStore(databaseName));
         } else {
             registry.registerComponent(TaskRegistryFactory.class, ctx -> EmptyTaskRegistryFactory.INSTANCE, true);
             registry.registerComponent(TaskStore.class, ctx -> EmptyTaskStore.INSTANCE, true);
@@ -87,7 +86,7 @@ public void shutdown() {
         };
     }
 
-    interface Dependencies {
+    public interface Dependencies {
         Config config();
 
         LogService logService();

progress-tracking/src/main/java/org/neo4j/gds/core/utils/progress/TaskRegistryFactoryProvider.java

@@ -25,13 +25,17 @@
 import org.neo4j.kernel.api.procedure.Context;
 
 class TaskRegistryFactoryProvider implements ThrowingFunction<Context, TaskRegistryFactory, ProcedureException> {
-    private final TaskStore taskStore;
 
-    TaskRegistryFactoryProvider(TaskStore taskStore) {this.taskStore = taskStore;}
+    private final TaskStoreProvider taskStoreProvider;
+
+    TaskRegistryFactoryProvider(TaskStoreProvider taskStoreProvider) {
+
+        this.taskStoreProvider = taskStoreProvider;
+    }
 
     @Override
     public TaskRegistryFactory apply(Context context) throws ProcedureException {
         var username = Neo4jProxy.username(context.securityContext().subject());
-        return new LocalTaskRegistryFactory(username, taskStore);
+        return new LocalTaskRegistryFactory(username, taskStoreProvider.apply(context));
     }
 }

progress-tracking/src/main/java/org/neo4j/gds/core/utils/progress/TaskStoreProvider.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.neo4j.gds.core.utils.progress;
+
+import org.neo4j.function.ThrowingFunction;
+import org.neo4j.internal.kernel.api.exceptions.ProcedureException;
+import org.neo4j.kernel.api.procedure.Context;
+
+class TaskStoreProvider implements ThrowingFunction<Context, TaskStore, ProcedureException> {
+
+    @Override
+    public TaskStore apply(Context context) throws ProcedureException {
+        return TaskStoreHolder.getTaskStore(context.graphDatabaseAPI().databaseName());
+    }
+}