Fix tests with prepopulated catalog store

homar · kokosing · commit b7a6a36b840b · 2025-10-20T14:00:07.000+02:00
Unfortunately io.trino.connector.ConnectorServicesProvider#loadInitialCatalogs
can't be called from TestingTrinoServer as it is too early and connector plugins
added in tests are not available yet. That results in initialization of catalog
failing for a wrong reason - lack of proper connector factory instead of
catalog being incorrectly defined.
diff --git a/core/trino-main/src/main/java/io/trino/server/testing/TestingTrinoServer.java b/core/trino-main/src/main/java/io/trino/server/testing/TestingTrinoServer.java
@@ -438,10 +438,6 @@ private TestingTrinoServer(
         spoolingConfiguration.ifPresent(config ->
                 spoolingManagerRegistry.loadSpoolingManager(config.factoryName(), config.configuration()));
 
-        catalogStoreManager.ifPresent(CatalogStoreManager::loadConfiguredCatalogStore);
-        ConnectorServicesProvider connectorServicesProvider = injector.getInstance(ConnectorServicesProvider.class);
-        connectorServicesProvider.loadInitialCatalogs();
-
         EventListenerManager eventListenerManager = injector.getInstance(EventListenerManager.class);
         eventListeners.forEach(eventListenerManager::addEventListener);
 
diff --git a/core/trino-main/src/main/java/io/trino/testing/StandaloneQueryRunner.java b/core/trino-main/src/main/java/io/trino/testing/StandaloneQueryRunner.java
@@ -20,6 +20,7 @@
 import io.opentelemetry.sdk.trace.data.SpanData;
 import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
 import io.trino.Session;
+import io.trino.connector.ConnectorServicesProvider;
 import io.trino.cost.StatsCalculator;
 import io.trino.execution.FailureInjector.InjectedFailureType;
 import io.trino.execution.QueryManagerConfig;
@@ -88,6 +89,7 @@ public StandaloneQueryRunner(Session defaultSession, Consumer<TestingTrinoServer
                         .buildOrThrow());
         serverProcessor.accept(builder);
         this.server = builder.build();
+        server.getInstance(Key.get(ConnectorServicesProvider.class)).loadInitialCatalogs();
 
         this.trinoClient = new TestingDirectTrinoClient(
                 server.getDispatchManager(),
diff --git a/testing/trino-testing/src/main/java/io/trino/testing/DistributedQueryRunner.java b/testing/trino-testing/src/main/java/io/trino/testing/DistributedQueryRunner.java
@@ -30,6 +30,7 @@
 import io.trino.Session.SessionBuilder;
 import io.trino.client.ClientSession;
 import io.trino.client.StatementClient;
+import io.trino.connector.ConnectorServicesProvider;
 import io.trino.connector.CoordinatorDynamicCatalogManager;
 import io.trino.cost.StatsCalculator;
 import io.trino.execution.FailureInjector.InjectedFailureType;
@@ -997,6 +998,7 @@ public void queryCompleted(QueryCompletedEvent queryCompletedEvent)
                 closeAllSuppress(e, queryRunner);
                 throw e;
             }
+            queryRunner.getCoordinator().getInstance(Key.get(ConnectorServicesProvider.class)).loadInitialCatalogs();
 
             return queryRunner;
         }
diff --git a/testing/trino-testing/src/main/java/io/trino/testing/QueryAssertions.java b/testing/trino-testing/src/main/java/io/trino/testing/QueryAssertions.java
@@ -470,7 +470,7 @@ public static void assertQueryFails(QueryRunner queryRunner, Session session, @L
         }
     }
 
-    protected static void assertQueryReturnsEmptyResult(QueryRunner queryRunner, Session session, @Language("SQL") String sql)
+    public static void assertQueryReturnsEmptyResult(QueryRunner queryRunner, Session session, @Language("SQL") String sql)
     {
         QueryId queryId = null;
         try {
diff --git a/testing/trino-tests/src/test/java/io/trino/connector/TestDynamicCatalogs.java b/testing/trino-tests/src/test/java/io/trino/connector/TestDynamicCatalogs.java
@@ -25,6 +25,7 @@
 import io.trino.spi.catalog.CatalogProperties;
 import io.trino.spi.catalog.CatalogStore;
 import io.trino.spi.catalog.CatalogStoreFactory;
+import io.trino.spi.connector.CatalogVersion;
 import io.trino.spi.connector.ConnectorName;
 import io.trino.testing.DistributedQueryRunner;
 import io.trino.testing.H2QueryRunner;
@@ -42,16 +43,20 @@
 import static io.trino.connector.FileCatalogStore.computeCatalogVersion;
 import static io.trino.testing.QueryAssertions.assertQuery;
 import static io.trino.testing.QueryAssertions.assertQueryFails;
+import static io.trino.testing.QueryAssertions.assertQueryReturnsEmptyResult;
 import static io.trino.testing.QueryAssertions.assertUpdate;
 import static io.trino.testing.TestingNames.randomNameSuffix;
 import static io.trino.testing.TestingSession.testSession;
+import static java.util.Objects.requireNonNull;
 import static org.junit.jupiter.api.parallel.ExecutionMode.SAME_THREAD;
 
 @Execution(SAME_THREAD)
 public class TestDynamicCatalogs
 {
     private static final String BROKEN_CATALOG = "broken_catalog";
+    private static final String PREPOPULATED_CATALOG = "prepopulated_catalog";
     private static final CatalogName BROKEN_CATALOG_NAME = new CatalogName(BROKEN_CATALOG);
+    private static final CatalogName PREPOPULATED_CATALOG_NAME = new CatalogName(PREPOPULATED_CATALOG);
     private static final ConnectorName MEMORY_CONNECTOR_NAME = new ConnectorName("memory");
 
     @Test
@@ -81,16 +86,21 @@ public void testNewHealthyCatalog()
     }
 
     @Test
-    public void testNewUnhealthyCatalog()
+    public void testPrepopulatedUnhealthyCatalog()
             throws Exception
     {
         Session session = testSession();
+        ImmutableMap<String, String> properties = ImmutableMap.of("non_existing", "false");
         QueryRunner queryRunner = DistributedQueryRunner.builder(session)
-                .setAdditionalModule(new TestCatalogStoreModule())
+                .setAdditionalModule(new TestCatalogStoreModule(ImmutableMap.of(BROKEN_CATALOG_NAME, new CatalogProperties(
+                        BROKEN_CATALOG_NAME,
+                        computeCatalogVersion(BROKEN_CATALOG_NAME, MEMORY_CONNECTOR_NAME, properties),
+                        MEMORY_CONNECTOR_NAME,
+                        properties))))
+                .setAdditionalSetup(runner -> runner.installPlugin(new MemoryPlugin()))
                 .setCoordinatorProperties(ImmutableMap.of("catalog.store", "prepopulated_memory"))
                 .setWorkerCount(0)
                 .build();
-        queryRunner.installPlugin(new MemoryPlugin());
         queryRunner.createCatalog("healthy_catalog", "memory", ImmutableMap.of("memory.max-data-per-node", "128MB"));
         H2QueryRunner h2QueryRunner = new H2QueryRunner();
 
@@ -103,29 +113,71 @@ public void testNewUnhealthyCatalog()
         assertQuery(queryRunner, session, "SHOW CATALOGS", h2QueryRunner, "VALUES 'healthy_catalog', 'system'", false, false);
     }
 
+    @Test
+    public void testPrepopulatedHealthyCatalog()
+            throws Exception
+    {
+        Session session = testSession();
+        ImmutableMap<String, String> properties = ImmutableMap.of("memory.max-data-per-node", "128MB");
+        QueryRunner queryRunner = DistributedQueryRunner.builder(session)
+                .setAdditionalModule(new TestCatalogStoreModule(ImmutableMap.of(PREPOPULATED_CATALOG_NAME, new CatalogProperties(
+                        PREPOPULATED_CATALOG_NAME,
+                        new CatalogVersion("abc123"),
+                        MEMORY_CONNECTOR_NAME,
+                        properties))))
+                .setAdditionalSetup(runner -> runner.installPlugin(new MemoryPlugin()))
+                .setCoordinatorProperties(ImmutableMap.of("catalog.store", "prepopulated_memory"))
+                .setWorkerCount(0)
+                .build();
+        queryRunner.createCatalog("healthy_catalog", "memory", ImmutableMap.of("memory.max-data-per-node", "128MB"));
+        H2QueryRunner h2QueryRunner = new H2QueryRunner();
+
+        assertQuery(queryRunner, session, "SHOW CATALOGS", h2QueryRunner, "VALUES 'healthy_catalog', '" + PREPOPULATED_CATALOG + "', 'system'", false, false);
+        assertUpdate(queryRunner, session, "CREATE TABLE %s.default.test_table (age INT)".formatted(PREPOPULATED_CATALOG), OptionalLong.empty(), Optional.empty());
+        assertQueryReturnsEmptyResult(queryRunner, session, "SELECT * FROM %s.default.test_table".formatted(PREPOPULATED_CATALOG));
+        assertQueryFails(queryRunner, session, "CREATE CATALOG %s USING memory WITH (\"memory.max-data-per-node\" = '128MB')".formatted(PREPOPULATED_CATALOG), ".*Catalog '%s' already exists.*".formatted(PREPOPULATED_CATALOG));
+
+        assertUpdate(queryRunner, session, "DROP CATALOG " + PREPOPULATED_CATALOG, OptionalLong.empty(), Optional.empty());
+        assertQuery(queryRunner, session, "SHOW CATALOGS", h2QueryRunner, "VALUES 'healthy_catalog', 'system'", false, false);
+    }
+
     public static class TestCatalogStoreModule
             extends AbstractConfigurationAwareModule
     {
+        private final Map<CatalogName, CatalogProperties> prepopulatedCatalogs;
+
+        public TestCatalogStoreModule(Map<CatalogName, CatalogProperties> prepopulatedCatalogs)
+        {
+            this.prepopulatedCatalogs = requireNonNull(prepopulatedCatalogs, "prepopulatedCatalogs is null");
+        }
+
         @Override
         protected void setup(Binder binder)
         {
             if (buildConfigObject(ServerConfig.class).isCoordinator()) {
-                install(new PrepopulatedInMemoryCatalogStoreModule());
+                install(new PrepopulatedInMemoryCatalogStoreModule(prepopulatedCatalogs));
             }
         }
     }
 
     private static class PrepopulatedInMemoryCatalogStoreModule
             extends AbstractConfigurationAwareModule
     {
+        private final Map<CatalogName, CatalogProperties> prepopulatedCatalogs;
+
+        public PrepopulatedInMemoryCatalogStoreModule(Map<CatalogName, CatalogProperties> prepopulatedCatalogs)
+        {
+            this.prepopulatedCatalogs = requireNonNull(prepopulatedCatalogs, "prepopulatedCatalogs is null");
+        }
+
         @Override
         protected void setup(Binder binder) {}
 
         @Provides
         @Singleton
         public PrepopulatedInMemoryCatalogStoreFactory createDbCatalogStoreFactory(CatalogStoreManager catalogStoreManager)
         {
-            PrepopulatedInMemoryCatalogStoreFactory factory = new PrepopulatedInMemoryCatalogStoreFactory();
+            PrepopulatedInMemoryCatalogStoreFactory factory = new PrepopulatedInMemoryCatalogStoreFactory(prepopulatedCatalogs);
             catalogStoreManager.addCatalogStoreFactory(factory);
             return factory;
         }
@@ -134,6 +186,13 @@ public PrepopulatedInMemoryCatalogStoreFactory createDbCatalogStoreFactory(Catal
     private static class PrepopulatedInMemoryCatalogStoreFactory
             implements CatalogStoreFactory
     {
+        private final Map<CatalogName, CatalogProperties> prepopulatedCatalogs;
+
+        public PrepopulatedInMemoryCatalogStoreFactory(Map<CatalogName, CatalogProperties> prepopulatedCatalogs)
+        {
+            this.prepopulatedCatalogs = requireNonNull(prepopulatedCatalogs, "prepopulatedCatalogs is null");
+        }
+
         @Override
         public String getName()
         {
@@ -143,36 +202,40 @@ public String getName()
         @Override
         public CatalogStore create(Map<String, String> config)
         {
-            return new PrepopulatedInMemoryCatalogStore();
+            return new PrepopulatedInMemoryCatalogStore(prepopulatedCatalogs);
         }
     }
 
     private static class PrepopulatedInMemoryCatalogStore
             extends InMemoryCatalogStore
     {
+        private final Map<CatalogName, CatalogProperties> prepopulatedCatalogs;
+
+        public PrepopulatedInMemoryCatalogStore(Map<CatalogName, CatalogProperties> prepopulatedCatalogs)
+        {
+            this.prepopulatedCatalogs = requireNonNull(prepopulatedCatalogs, "prepopulatedCatalogs is null");
+        }
+
         @Override
         public Collection<StoredCatalog> getCatalogs()
         {
             Collection<StoredCatalog> catalogs = super.getCatalogs();
             List<StoredCatalog> catalogsCopy = new ArrayList<>(catalogs);
-            catalogsCopy.add(new StoredCatalog()
-            {
-                @Override
-                public CatalogName name()
-                {
-                    return new CatalogName("broken_catalog");
-                }
-
-                @Override
-                public CatalogProperties loadProperties()
+            prepopulatedCatalogs.forEach((catalogName, catalogProperties) -> {
+                catalogsCopy.add(new StoredCatalog()
                 {
-                    ImmutableMap<String, String> properties = ImmutableMap.of("non_existing", "false");
-                    return new CatalogProperties(
-                            BROKEN_CATALOG_NAME,
-                            computeCatalogVersion(BROKEN_CATALOG_NAME, MEMORY_CONNECTOR_NAME, properties),
-                            MEMORY_CONNECTOR_NAME,
-                            properties);
-                }
+                    @Override
+                    public CatalogName name()
+                    {
+                        return catalogName;
+                    }
+
+                    @Override
+                    public CatalogProperties loadProperties()
+                    {
+                        return catalogProperties;
+                    }
+                });
             });
             return catalogsCopy;
         }
diff --git a/testing/trino-tests/src/test/java/io/trino/connector/system/TestSystemMetadataCatalogTable.java b/testing/trino-tests/src/test/java/io/trino/connector/system/TestSystemMetadataCatalogTable.java
@@ -17,6 +17,10 @@
 import io.trino.Session;
 import io.trino.connector.TestDynamicCatalogs;
 import io.trino.plugin.memory.MemoryPlugin;
+import io.trino.spi.catalog.CatalogName;
+import io.trino.spi.catalog.CatalogProperties;
+import io.trino.spi.connector.CatalogVersion;
+import io.trino.spi.connector.ConnectorName;
 import io.trino.testing.AbstractTestQueryFramework;
 import io.trino.testing.DistributedQueryRunner;
 import io.trino.testing.QueryRunner;
@@ -30,13 +34,20 @@
 public class TestSystemMetadataCatalogTable
         extends AbstractTestQueryFramework
 {
+    private static final String BROKEN_CATALOG = "broken_catalog";
+
     @Override
     protected QueryRunner createQueryRunner()
             throws Exception
     {
         Session session = testSessionBuilder().build();
+        ImmutableMap<String, String> properties = ImmutableMap.of("non_existing", "false");
         QueryRunner queryRunner = DistributedQueryRunner.builder(session)
-                .setAdditionalModule(new TestDynamicCatalogs.TestCatalogStoreModule())
+                .setAdditionalModule(new TestDynamicCatalogs.TestCatalogStoreModule(ImmutableMap.of(new CatalogName(BROKEN_CATALOG), new CatalogProperties(
+                        new CatalogName(BROKEN_CATALOG),
+                        new CatalogVersion("abc123"),
+                        new ConnectorName("memory"),
+                        properties))))
                 .setCoordinatorProperties(ImmutableMap.of("catalog.store", "prepopulated_memory"))
                 .setWorkerCount(0)
                 .build();

Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,7 @@`
`30`	`30`	`import io.trino.Session.SessionBuilder;`
`31`	`31`	`import io.trino.client.ClientSession;`
`32`	`32`	`import io.trino.client.StatementClient;`
	`33`	`+import io.trino.connector.ConnectorServicesProvider;`
`33`	`34`	`import io.trino.connector.CoordinatorDynamicCatalogManager;`
`34`	`35`	`import io.trino.cost.StatsCalculator;`
`35`	`36`	`import io.trino.execution.FailureInjector.InjectedFailureType;`
`@@ -997,6 +998,7 @@ public void queryCompleted(QueryCompletedEvent queryCompletedEvent)`
`997`	`998`	`closeAllSuppress(e, queryRunner);`
`998`	`999`	`throw e;`
`999`	`1000`	`}`
	`1001`	`+ queryRunner.getCoordinator().getInstance(Key.get(ConnectorServicesProvider.class)).loadInitialCatalogs();`
`1000`	`1002`
`1001`	`1003`	`return queryRunner;`
`1002`	`1004`	`}`
Original file line number	Diff line number	Diff line change
`@@ -470,7 +470,7 @@ public static void assertQueryFails(QueryRunner queryRunner, Session session, @L`
`470`	`470`	`}`
`471`	`471`	`}`
`472`	`472`
`473`		`- protected static void assertQueryReturnsEmptyResult(QueryRunner queryRunner, Session session, @Language("SQL") String sql)`
	`473`	`+ public static void assertQueryReturnsEmptyResult(QueryRunner queryRunner, Session session, @Language("SQL") String sql)`
`474`	`474`	`{`
`475`	`475`	`QueryId queryId = null;`
`476`	`476`	`try {`