Moved Neo4jTransactionalContextFactory creation out of DataSourceModule

community/cypher/cypher/src/test/scala/org/neo4j/cypher/internal/spi/v3_1/TransactionBoundQueryContextTest.scala

@@ -72,7 +72,7 @@ class TransactionBoundQueryContextTest extends CypherFunSuite {
     // GIVEN
     when(outerTx.failure()).thenThrow(new AssertionError("Shouldn't be called"))
     val tc = new Neo4jTransactionalContext(graph, outerTx, KernelTransaction.Type.`implicit`, AccessMode.Static.FULL,
-      statement, null, locker, null, null, null)
+      statement, null, locker, null, null, null, null)
     val transactionalContext = TransactionalContextWrapperv3_1(tc)
     val context = new TransactionBoundQueryContext(transactionalContext)(indexSearchMonitor)
     // WHEN
@@ -88,7 +88,7 @@ class TransactionBoundQueryContextTest extends CypherFunSuite {
     // GIVEN
     when(outerTx.success()).thenThrow(new AssertionError("Shouldn't be called"))
     val tc = new Neo4jTransactionalContext(graph, outerTx, KernelTransaction.Type.`implicit`, AccessMode.Static.FULL,
-      statement, null, locker, null, null, null)
+      statement, null, locker, null, null, null, null)
     val transactionalContext = TransactionalContextWrapperv3_1(tc)
     val context = new TransactionBoundQueryContext(transactionalContext)(indexSearchMonitor)
     // WHEN

community/kernel/src/main/java/org/neo4j/kernel/impl/factory/DataSourceModule.java

@@ -21,7 +21,6 @@
 
 import java.io.File;
 import java.time.Clock;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
 import org.neo4j.graphdb.DependencyResolver;
@@ -67,7 +66,6 @@
 import org.neo4j.kernel.impl.proc.ProcedureGDSFactory;
 import org.neo4j.kernel.impl.proc.Procedures;
 import org.neo4j.kernel.impl.proc.TypeMappers.SimpleConverter;
-import org.neo4j.kernel.impl.query.QueryEngineProvider;
 import org.neo4j.kernel.impl.query.QueryExecutionEngine;
 import org.neo4j.kernel.impl.store.StoreId;
 import org.neo4j.kernel.impl.transaction.log.PhysicalLogFile;
@@ -118,8 +116,8 @@ public class DataSourceModule
 
     public final AutoIndexing autoIndexing;
 
-    public DataSourceModule( final GraphDatabaseFacadeFactory.Dependencies dependencies,
-            final PlatformModule platformModule, EditionModule editionModule )
+    public DataSourceModule( final PlatformModule platformModule, EditionModule editionModule,
+            Supplier<QueryExecutionEngine> queryExecutionEngineSupplier )
     {
         final Dependencies deps = platformModule.dependencies;
         Config config = platformModule.config;
@@ -131,6 +129,7 @@ public DataSourceModule( final GraphDatabaseFacadeFactory.Dependencies dependenc
         RelationshipTypeTokenHolder relationshipTypeTokenHolder = editionModule.relationshipTypeTokenHolder;
         File storeDir = platformModule.storeDir;
         DiagnosticsManager diagnosticsManager = platformModule.diagnosticsManager;
+        this.queryExecutor = queryExecutionEngineSupplier;
 
         threadToTransactionBridge = deps.satisfyDependency( life.add( new ThreadToStatementContextBridge() ) );
 
@@ -169,8 +168,6 @@ public DataSourceModule( final GraphDatabaseFacadeFactory.Dependencies dependenc
 
         autoIndexing = new InternalAutoIndexing( platformModule.config, editionModule.propertyKeyTokenHolder );
 
-        AtomicReference<QueryExecutionEngine> queryExecutor = new AtomicReference<>( QueryEngineProvider.noEngine() );
-        this.queryExecutor = queryExecutor::get;
         Procedures procedures = setupProcedures( platformModule, editionModule );
 
         deps.satisfyDependency( new NonTransactionalDbmsOperations.Factory( procedures ) );
@@ -228,29 +225,6 @@ public DataSourceModule( final GraphDatabaseFacadeFactory.Dependencies dependenc
         // Kernel event handlers should be the very last, i.e. very first to receive shutdown events
         life.add( kernelEventHandlers );
 
-        dataSourceManager.addListener( new DataSourceManager.Listener()
-        {
-            private QueryExecutionEngine engine;
-
-            @Override
-            public void registered( NeoStoreDataSource dataSource )
-            {
-                if ( engine == null )
-                {
-                    engine = QueryEngineProvider.initialize( deps, platformModule.graphDatabaseFacade,
-                            dependencies.executionEngines() );
-                }
-
-                queryExecutor.set( engine );
-            }
-
-            @Override
-            public void unregistered( NeoStoreDataSource dataSource )
-            {
-                queryExecutor.set( QueryEngineProvider.noEngine() );
-            }
-        } );
-
         this.storeId = neoStoreDataSource::getStoreId;
         this.kernelAPI = neoStoreDataSource::getKernel;
     }

community/kernel/src/main/java/org/neo4j/kernel/impl/factory/GraphDatabaseFacade.java

@@ -23,7 +23,6 @@
 import java.net.URL;
 import java.util.Collections;
 import java.util.Map;
-import java.util.function.Supplier;
 import java.util.concurrent.TimeUnit;
 
 import org.neo4j.collection.primitive.PrimitiveLongCollections;
@@ -99,7 +98,6 @@
 
 import static java.lang.String.format;
 import static org.neo4j.collection.primitive.PrimitiveLongCollections.map;
-import static org.neo4j.function.Suppliers.lazySingleton;
 import static org.neo4j.helpers.collection.Iterators.emptyIterator;
 import static org.neo4j.kernel.impl.api.operations.KeyReadOperations.NO_SUCH_LABEL;
 import static org.neo4j.kernel.impl.api.operations.KeyReadOperations.NO_SUCH_PROPERTY_KEY;
@@ -118,7 +116,7 @@ public class GraphDatabaseFacade implements GraphDatabaseAPI
     private NodeProxy.NodeActions nodeActions;
     private RelationshipProxy.RelationshipActions relActions;
     private SPI spi;
-    private Supplier<TransactionalContextFactory> contextFactorySupplier;
+    private TransactionalContextFactory contextFactory;
     private long defaultTransactionTimeout;
 
     /**
@@ -197,7 +195,7 @@ public GraphDatabaseFacade()
     /**
      * Create a new Core API facade, backed by the given SPI.
      */
-    public void init( SPI spi, Config config )
+    public void initSPI( SPI spi, Config config )
     {
         defaultTransactionTimeout = config.get( GraphDatabaseSettings.transaction_timeout );
         IndexProviderImpl idxProvider = new IndexProviderImpl( this, spi::currentStatement );
@@ -219,7 +217,17 @@ public void init( SPI spi, Config config )
                         .getOrCreateRelationshipIndex( InternalAutoIndexing.RELATIONSHIP_AUTO_INDEX, null ) ),
                 spi.autoIndexing().relationships() );
         this.indexManager = new IndexManagerImpl( spi::currentStatement, idxProvider, nodeAutoIndexer, relAutoIndexer );
-        this.contextFactorySupplier = lazySingleton( () -> new Neo4jTransactionalContextFactory( spi, locker ) );
+    }
+
+    /**
+     * This needs to be called after data source creation, and the execution engine dependency is satisfied by the
+     * cypher module GraphData
+     *
+     * @see GraphDatabaseFacadeFactory#initFacade(File, Map, GraphDatabaseFacadeFactory.Dependencies, GraphDatabaseFacade)
+     */
+    public void initTransactionalContextFactoryFromSPI()
+    {
+        contextFactory = new Neo4jTransactionalContextFactory( spi, locker );
     }
 
     @Override
@@ -380,7 +388,6 @@ public Result execute( String query, Map<String,Object> parameters ) throws Quer
         // ensure we have a tx and create a context (the tx is gonna get closed by the Cypher result)
         InternalTransaction transaction = beginTransaction( KernelTransaction.Type.implicit, AccessMode.Static.FULL );
 
-        TransactionalContextFactory contextFactory = contextFactorySupplier.get();
         TransactionalContext context =
                 contextFactory.newContext( QueryEngineProvider.describe(), transaction, query, parameters );
         return spi.executeQuery( query, parameters, context );
@@ -399,7 +406,6 @@ public Result execute( String query, Map<String,Object> parameters, long timeout
     public Result execute( InternalTransaction transaction, String query, Map<String,Object> parameters )
             throws QueryExecutionException
     {
-        TransactionalContextFactory contextFactory = contextFactorySupplier.get();
         TransactionalContext context =
                 contextFactory.newContext( QueryEngineProvider.describe(), transaction, query, parameters );
         return spi.executeQuery( query, parameters, context );

community/kernel/src/main/java/org/neo4j/kernel/impl/factory/GraphDatabaseFacadeFactory.java

@@ -21,16 +21,21 @@
 
 import java.io.File;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 import org.neo4j.graphdb.config.Setting;
 import org.neo4j.graphdb.security.URLAccessRule;
 import org.neo4j.helpers.Exceptions;
 import org.neo4j.kernel.AvailabilityGuard;
+import org.neo4j.kernel.NeoStoreDataSource;
 import org.neo4j.kernel.configuration.Settings;
 import org.neo4j.kernel.extension.KernelExtensionFactory;
 import org.neo4j.kernel.impl.coreapi.CoreAPIAvailabilityGuard;
 import org.neo4j.kernel.impl.query.QueryEngineProvider;
+import org.neo4j.kernel.impl.query.QueryExecutionEngine;
+import org.neo4j.kernel.impl.transaction.state.DataSourceManager;
 import org.neo4j.kernel.monitoring.Monitors;
 import org.neo4j.logging.LogProvider;
 import org.neo4j.logging.Logger;
@@ -40,6 +45,7 @@
 import static org.neo4j.kernel.configuration.Settings.illegalValueMessage;
 import static org.neo4j.kernel.configuration.Settings.matches;
 import static org.neo4j.kernel.configuration.Settings.setting;
+import static org.neo4j.kernel.impl.query.QueryEngineProvider.noEngine;
 
 /**
  * This is the main factory for creating database instances. It delegates creation to three different modules
@@ -133,14 +139,40 @@ public GraphDatabaseFacade initFacade( File storeDir, Map<String,String> params,
     {
         PlatformModule platform = createPlatform( storeDir, params, dependencies, graphDatabaseFacade );
         EditionModule edition = editionFactory.apply( platform );
-        final DataSourceModule dataSource = createDataSource( dependencies, platform, edition );
+
+        AtomicReference<QueryExecutionEngine> queryEngine = new AtomicReference<>( noEngine() );
+        final DataSourceModule dataSource = createDataSource( platform, edition, queryEngine::get );
         Logger msgLog = platform.logging.getInternalLog( getClass() ).infoLogger();
         CoreAPIAvailabilityGuard coreAPIAvailabilityGuard = edition.coreAPIAvailabilityGuard;
 
         // Start it
-        graphDatabaseFacade.init( new ClassicCoreSPI( platform, dataSource, msgLog, coreAPIAvailabilityGuard ),
+        graphDatabaseFacade.initSPI( new ClassicCoreSPI( platform, dataSource, msgLog, coreAPIAvailabilityGuard ),
                 platform.config );
 
+        platform.dataSourceManager.addListener( new DataSourceManager.Listener()
+        {
+            private QueryExecutionEngine engine;
+
+            @Override
+            public void registered( NeoStoreDataSource dataSource )
+            {
+                if ( engine == null )
+                {
+                    engine = QueryEngineProvider.initialize(
+                            platform.dependencies, platform.graphDatabaseFacade, dependencies.executionEngines() );
+                    graphDatabaseFacade.initTransactionalContextFactoryFromSPI();
+                }
+
+                queryEngine.set( engine );
+            }
+
+            @Override
+            public void unregistered( NeoStoreDataSource dataSource )
+            {
+                queryEngine.set( noEngine() );
+            }
+        } );
+
         Throwable error = null;
         try
         {
@@ -191,10 +223,10 @@ protected PlatformModule createPlatform( File storeDir, Map<String,String> param
     /**
      * Create the datasource module. Override to replace with custom module.
      */
-    protected DataSourceModule createDataSource( final Dependencies dependencies, final PlatformModule platformModule,
-            EditionModule editionModule )
+    protected DataSourceModule createDataSource( final PlatformModule platformModule, EditionModule editionModule,
+            Supplier<QueryExecutionEngine> queryEngine )
     {
-        return new DataSourceModule( dependencies, platformModule, editionModule );
+        return new DataSourceModule( platformModule, editionModule, queryEngine );
     }
 
     private void enableAvailabilityLogging( AvailabilityGuard availabilityGuard, final Logger msgLog )

community/kernel/src/main/java/org/neo4j/kernel/impl/proc/ProcedureGDSFactory.java

@@ -71,8 +71,9 @@ public GraphDatabaseService apply( Context context ) throws ProcedureException
         KernelTransaction transaction = context.get( Context.KERNEL_TRANSACTION );
         Thread owningThread = context.get( Context.THREAD );
         GraphDatabaseFacade facade = new GraphDatabaseFacade();
-        facade.init( new ProcedureGDBFacadeSPI( owningThread, transaction, queryExecutor, storeDir, resolver,
+        facade.initSPI( new ProcedureGDBFacadeSPI( owningThread, transaction, queryExecutor, storeDir, resolver,
                 AutoIndexing.UNSUPPORTED, storeId, availability, urlValidator ), config );
+        facade.initTransactionalContextFactoryFromSPI();
         return facade;
     }
 

community/kernel/src/main/java/org/neo4j/kernel/impl/query/Neo4jTransactionalContext.java

@@ -19,8 +19,6 @@
  */
 package org.neo4j.kernel.impl.query;
 
-import java.util.function.Function;
-
 import org.neo4j.graphdb.Lock;
 import org.neo4j.graphdb.PropertyContainer;
 import org.neo4j.kernel.GraphDatabaseQueryService;
@@ -46,6 +44,7 @@ public class Neo4jTransactionalContext implements TransactionalContext
     private final AccessMode mode;
     private final DbmsOperations.Factory dbmsOperationsFactory;
     private final Guard guard;
+    private final TransactionalContextFactory factory;
 
     private InternalTransaction transaction;
     private Statement statement;
@@ -64,8 +63,9 @@ public Neo4jTransactionalContext(
             PropertyContainerLocker locker,
             ThreadToStatementContextBridge txBridge,
             DbmsOperations.Factory dbmsOperationsFactory,
-            Guard guard )
-    {
+            Guard guard,
+            TransactionalContextFactory factory
+    ) {
         this.graph = graph;
         this.transaction = initialTransaction;
         this.transactionType = transactionType;
@@ -76,6 +76,7 @@ public Neo4jTransactionalContext(
         this.txBridge = txBridge;
         this.dbmsOperationsFactory = dbmsOperationsFactory;
         this.guard = guard;
+        this.factory = factory;
     }
 
     @Override
@@ -203,20 +204,7 @@ public TransactionalContext getOrBeginNewIfClosed()
         else
         {
             InternalTransaction transaction = graph.beginTransaction( transactionType, mode );
-            Statement statement = txBridge.get();
-            statement.queryRegistration().registerExecutingQuery( executingQuery );
-            return new Neo4jTransactionalContext(
-                    graph,
-                    transaction,
-                    transaction.transactionType(),
-                    transaction.mode(),
-                    statement,
-                    executingQuery,
-                    locker,
-                    graph.getDependencyResolver().resolveDependency( ThreadToStatementContextBridge.class ),
-                    graph.getDependencyResolver().resolveDependency( DbmsOperations.Factory.class ),
-                    guard
-            );
+            return factory.newContext( this.executingQuery, transaction );
         }
     }