diff --git a/gateway-ha/config.yaml b/gateway-ha/config.yaml index c9e0c4a29..90002d4f2 100644 --- a/gateway-ha/config.yaml +++ b/gateway-ha/config.yaml @@ -7,9 +7,9 @@ routingRules: # rulesConfigPath: "src/main/resources/rules/routing_rules.yml" dataStore: - jdbcUrl: jdbc:postgresql://localhost:5432/trino_gateway_db - user: trino_gateway_db_admin - password: P0stG&es + jdbcUrl: jdbc:postgresql://${ENV:DB_HOSTNAME:localhost}:${ENV:DB_PORT:5432}/${ENV:DB_NAME:trino_gateway_db} + user: ${ENV:DB_USER:trino_gateway_db_admin} + password: ${ENV:DB_PASSWORD:P0stG&es} driver: org.postgresql.Driver queryHistoryHoursRetention: 24 diff --git a/gateway-ha/pom.xml b/gateway-ha/pom.xml index 9fe50911c..97d15fc88 100644 --- a/gateway-ha/pom.xml +++ b/gateway-ha/pom.xml @@ -324,12 +324,6 @@ - - com.h2database - h2 - 1.4.192 - test - com.squareup.okhttp3 diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ExactMatchSourceSelectorsDao.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ExactMatchSourceSelectorsDao.java index e7d31c5f7..bb0d69a9c 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ExactMatchSourceSelectorsDao.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/ExactMatchSourceSelectorsDao.java @@ -14,14 +14,46 @@ package io.trino.gateway.ha.persistence.dao; import io.trino.gateway.ha.router.ResourceGroupsManager; +import org.jdbi.v3.core.mapper.ColumnMapper; +import org.jdbi.v3.core.statement.StatementContext; +import org.jdbi.v3.sqlobject.config.RegisterColumnMapper; import org.jdbi.v3.sqlobject.customizer.BindBean; import org.jdbi.v3.sqlobject.statement.SqlQuery; import org.jdbi.v3.sqlobject.statement.SqlUpdate; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Timestamp; +import java.time.LocalDateTime; +import java.time.ZoneId; +import java.time.format.DateTimeFormatter; import java.util.List; +@RegisterColumnMapper(ExactMatchSourceSelectorsDao.TimestampColumnMapper.class) public interface ExactMatchSourceSelectorsDao { + class TimestampColumnMapper + implements ColumnMapper + { + private static final DateTimeFormatter FORMATTER = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"); + + @Override + public String map(ResultSet result, int columnNumber, StatementContext ctx) + throws SQLException + { + String columnName = result.getMetaData().getColumnName(columnNumber); + if ("update_time".equals(columnName)) { + try { + Timestamp timestamp = result.getTimestamp(columnNumber); + return timestamp != null ? timestamp.toLocalDateTime().format(FORMATTER) : null; + } + catch (SQLException e) { + return LocalDateTime.now(ZoneId.systemDefault()).format(FORMATTER); + } + } + return result.getString(columnNumber); + } + } @SqlQuery(""" SELECT * FROM exact_match_source_selectors """) @@ -30,7 +62,20 @@ public interface ExactMatchSourceSelectorsDao @SqlUpdate(""" INSERT INTO exact_match_source_selectors (resource_group_id, update_time, source, environment, query_type) - VALUES (:resourceGroupId, :updateTime, :source, :environment, :queryType) + VALUES (:resourceGroupId, + CASE + WHEN :updateTime IS NULL OR :updateTime = '' + THEN CURRENT_TIMESTAMP + ELSE CAST(:updateTime AS TIMESTAMP) + END, + :source, :environment, :queryType) """) void insert(@BindBean ResourceGroupsManager.ExactSelectorsDetail exactSelectors); + + @SqlUpdate(""" + INSERT INTO exact_match_source_selectors + (resource_group_id, update_time, source, environment, query_type) + VALUES (:resourceGroupId, CURRENT_TIMESTAMP, :source, :environment, :queryType) + """) + void insertWithCurrentTimestamp(@BindBean ResourceGroupsManager.ExactSelectorsDetail exactSelectors); } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaResourceGroupsManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaResourceGroupsManager.java index 462c17094..7baee0c92 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaResourceGroupsManager.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaResourceGroupsManager.java @@ -25,6 +25,9 @@ import io.trino.gateway.ha.persistence.dao.SelectorsDao; import jakarta.annotation.Nullable; +import java.time.LocalDateTime; +import java.time.ZoneId; +import java.time.format.DateTimeFormatter; import java.util.ArrayList; import java.util.List; @@ -235,11 +238,18 @@ public void deleteGlobalProperty(String name, @Nullable String routingGroupDatab /** * Creates exact match source selector for db. + * If no updateTime is provided, the current timestamp will be used automatically. */ @Override public ExactSelectorsDetail createExactMatchSourceSelector( ExactSelectorsDetail exactSelectorDetail) { + // If updateTime is null or empty, set current timestamp + if (exactSelectorDetail.getUpdateTime() == null || exactSelectorDetail.getUpdateTime().trim().isEmpty()) { + DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"); + exactSelectorDetail.setUpdateTime(LocalDateTime.now(ZoneId.systemDefault()).format(formatter)); + } + exactMatchSourceSelectorsDao.insert(exactSelectorDetail); return exactSelectorDetail; } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java b/gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java index 051644436..00ff683af 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java @@ -70,12 +70,11 @@ public static void prepareMockBackend( .setResponseCode(200)); } - public static void seedRequiredData(String h2DbFilePath) + public static void seedRequiredData(PostgreSQLContainer container) { - String jdbcUrl = "jdbc:h2:" + h2DbFilePath; - Jdbi jdbi = Jdbi.create(jdbcUrl, "sa", "sa"); + Jdbi jdbi = Jdbi.create(container.getJdbcUrl(), container.getUsername(), container.getPassword()); try (Handle handle = jdbi.open()) { - handle.createUpdate(HaGatewayTestUtils.getResourceFileContent("gateway-ha-persistence-mysql.sql")) + handle.createUpdate(HaGatewayTestUtils.getResourceFileContent("gateway-ha-persistence-postgres.sql")) .execute(); } } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestingJdbcConnectionManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestingJdbcConnectionManager.java index 1f60a0dfb..ca90fe200 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/TestingJdbcConnectionManager.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/TestingJdbcConnectionManager.java @@ -17,9 +17,7 @@ import io.trino.gateway.ha.persistence.JdbcConnectionManager; import org.jdbi.v3.core.Jdbi; import org.testcontainers.containers.JdbcDatabaseContainer; - -import java.io.File; -import java.nio.file.Path; +import org.testcontainers.containers.PostgreSQLContainer; public final class TestingJdbcConnectionManager { @@ -27,12 +25,18 @@ private TestingJdbcConnectionManager() {} public static JdbcConnectionManager createTestingJdbcConnectionManager() { - File tempH2DbDir = Path.of(System.getProperty("java.io.tmpdir"), "h2db-" + System.currentTimeMillis()).toFile(); - tempH2DbDir.deleteOnExit(); - String jdbcUrl = "jdbc:h2:" + tempH2DbDir.getAbsolutePath(); - HaGatewayTestUtils.seedRequiredData(tempH2DbDir.getAbsolutePath()); - DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", "sa", "org.h2.Driver", 4, false); - Jdbi jdbi = Jdbi.create(jdbcUrl, "sa", "sa"); + PostgreSQLContainer postgres = new PostgreSQLContainer<>("postgres:14-alpine") + .withDatabaseName("testdb") + .withInitScript("gateway-ha-persistence-postgres.sql"); + postgres.start(); + + String jdbcUrl = postgres.getJdbcUrl(); + String username = postgres.getUsername(); + String password = postgres.getPassword(); + + DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, username, password, "org.postgresql.Driver", 4, false); + Jdbi jdbi = Jdbi.create(jdbcUrl, username, password); + return new JdbcConnectionManager(jdbi, db); } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestJdbcConnectionManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestJdbcConnectionManager.java index 8a3ff454c..2a3bcd4ee 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestJdbcConnectionManager.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestJdbcConnectionManager.java @@ -24,20 +24,6 @@ final class TestJdbcConnectionManager { - @Test - void testBuildJdbcUrlWithH2AndNoRoutingGroupDatabase() - { - JdbcConnectionManager connectionManager = createConnectionManager("jdbc:h2:/mydb"); - assertThat(connectionManager.buildJdbcUrl(null)).isEqualTo("jdbc:h2:/mydb"); - } - - @Test - void testBuildJdbcUrlWithH2AndRoutingGroupDatabase() - { - JdbcConnectionManager connectionManager = createConnectionManager("jdbc:h2:/mydb"); - assertThat(connectionManager.buildJdbcUrl("newdb")).isEqualTo("jdbc:h2:/newdb"); - } - @Test void testBuildJdbcUrlWithMySQLAndNoRoutingGroupDatabase() { @@ -108,7 +94,7 @@ void testBuildJdbcUrlWithNullJdbcUrlThrowsException() DataStoreConfiguration dataStoreConfiguration = Mockito.mock(DataStoreConfiguration.class); when(dataStoreConfiguration.getJdbcUrl()).thenReturn(null); - JdbcConnectionManager connectionManager = new JdbcConnectionManager(Jdbi.create("jdbc:h2:/mydb", "sa", "sa"), dataStoreConfiguration); + JdbcConnectionManager connectionManager = new JdbcConnectionManager(Jdbi.create("jdbc:postgresql://localhost:5432/mydb", "postgres", "postgres"), dataStoreConfiguration); assertThatThrownBy(() -> connectionManager.buildJdbcUrl(null)) .isInstanceOf(IllegalArgumentException.class) .hasMessage("JDBC URL cannot be null"); @@ -117,10 +103,10 @@ void testBuildJdbcUrlWithNullJdbcUrlThrowsException() @Test void testBuildJdbcUrlWithNoSlashThrowsException() { - JdbcConnectionManager connectionManager = createConnectionManager("jdbc:h2:mem:test"); + JdbcConnectionManager connectionManager = createConnectionManager("jdbc:postgresql:mydb"); assertThatThrownBy(() -> connectionManager.buildJdbcUrl("newdb")) .isInstanceOf(IllegalArgumentException.class) - .hasMessage("Invalid JDBC URL: no '/' found in jdbc:h2:mem:test"); + .hasMessage("Invalid JDBC URL: no '/' found in jdbc:postgresql:mydb"); } private static JdbcConnectionManager createConnectionManager(String jdbcUrl) diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestResourceGroupsManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestResourceGroupsManager.java index 1e1716358..b32d1df8a 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestResourceGroupsManager.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestResourceGroupsManager.java @@ -76,20 +76,26 @@ void testReadResourceGroup() List resourceGroups = resourceGroupManager.readAllResourceGroups(null); assertThat(resourceGroups).hasSize(2); - assertThat(resourceGroups.get(0).getResourceGroupId()).isEqualTo(1L); - assertThat(resourceGroups.get(0).getName()).isEqualTo("admin"); - assertThat(resourceGroups.get(0).getHardConcurrencyLimit()).isEqualTo(20); - assertThat(resourceGroups.get(0).getMaxQueued()).isEqualTo(200); - assertThat(resourceGroups.get(0).getJmxExport()).isEqualTo(Boolean.TRUE); - assertThat(resourceGroups.get(0).getSoftMemoryLimit()).isEqualTo("80%"); + assertThat(resourceGroups.getFirst().getResourceGroupId()).isEqualTo(1L); + assertThat(resourceGroups.getFirst().getName()).isEqualTo("admin"); + assertThat(resourceGroups.getFirst().getHardConcurrencyLimit()).isEqualTo(20); + assertThat(resourceGroups.getFirst().getMaxQueued()).isEqualTo(200); + assertThat(resourceGroups.getFirst().getJmxExport()).isEqualTo(Boolean.TRUE); + assertThat(resourceGroups.getFirst().getSoftMemoryLimit()).isEqualTo("80%"); } @Test @Order(3) void testUpdateResourceGroup() { + List resourceGroups = resourceGroupManager.readAllResourceGroups(null); + long adminResourceGroupId = resourceGroups.stream() + .filter(rg -> "admin".equals(rg.getName())) + .findFirst() + .map(ResourceGroupsDetail::getResourceGroupId) + .orElse(1L); ResourceGroupsDetail resourceGroup = new ResourceGroupsDetail(); - resourceGroup.setResourceGroupId(1L); + resourceGroup.setResourceGroupId(adminResourceGroupId); resourceGroup.setName("admin"); resourceGroup.setHardConcurrencyLimit(50); resourceGroup.setMaxQueued(50); @@ -97,54 +103,60 @@ void testUpdateResourceGroup() resourceGroup.setSoftMemoryLimit("20%"); ResourceGroupsDetail updated = resourceGroupManager.updateResourceGroup(resourceGroup, null); - List resourceGroups = resourceGroupManager.readAllResourceGroups(null); + resourceGroups = resourceGroupManager.readAllResourceGroups(null); assertThat(resourceGroups).contains(updated); /* Update resourceGroups that do not exist yet. * In this case, new resourceGroups should be created. */ - resourceGroup.setResourceGroupId(3L); + resourceGroup = new ResourceGroupsDetail(); resourceGroup.setName("localization-eng"); resourceGroup.setHardConcurrencyLimit(50); resourceGroup.setMaxQueued(70); resourceGroup.setJmxExport(true); resourceGroup.setSoftMemoryLimit("20%"); resourceGroup.setSoftConcurrencyLimit(20); - resourceGroupManager.updateResourceGroup(resourceGroup, null); + ResourceGroupsDetail localizationGroup = resourceGroupManager.createResourceGroup(resourceGroup, null); - resourceGroup.setResourceGroupId(4L); + resourceGroup = new ResourceGroupsDetail(); resourceGroup.setName("resource_group_3"); resourceGroup.setHardConcurrencyLimit(10); resourceGroup.setMaxQueued(150); resourceGroup.setJmxExport(true); resourceGroup.setSoftMemoryLimit("60%"); resourceGroup.setSoftConcurrencyLimit(40); - resourceGroupManager.updateResourceGroup(resourceGroup, null); + resourceGroupManager.createResourceGroup(resourceGroup, null); resourceGroups = resourceGroupManager.readAllResourceGroups(null); - assertThat(resourceGroups).hasSize(4); // updated 2 non-existing groups, so count should be 4 - - assertThat(resourceGroups.get(0).getResourceGroupId()).isEqualTo(1L); - assertThat(resourceGroups.get(0).getName()).isEqualTo("admin"); - assertThat(resourceGroups.get(0).getHardConcurrencyLimit()).isEqualTo(50); - assertThat(resourceGroups.get(0).getMaxQueued()).isEqualTo(50); - assertThat(resourceGroups.get(0).getJmxExport()).isEqualTo(Boolean.FALSE); - assertThat(resourceGroups.get(0).getSoftMemoryLimit()).isEqualTo("20%"); - - assertThat(resourceGroups.get(1).getResourceGroupId()).isEqualTo(2L); - assertThat(resourceGroups.get(1).getName()).isEqualTo("user"); - assertThat(resourceGroups.get(1).getHardConcurrencyLimit()).isEqualTo(10); - assertThat(resourceGroups.get(1).getMaxQueued()).isEqualTo(100); - assertThat(resourceGroups.get(1).getJmxExport()).isEqualTo(Boolean.TRUE); - assertThat(resourceGroups.get(1).getSoftMemoryLimit()).isEqualTo("50%"); - - assertThat(resourceGroups.get(2).getResourceGroupId()).isEqualTo(3L); - assertThat(resourceGroups.get(2).getName()).isEqualTo("localization-eng"); - assertThat(resourceGroups.get(2).getHardConcurrencyLimit()).isEqualTo(50); - assertThat(resourceGroups.get(2).getMaxQueued()).isEqualTo(70); - assertThat(resourceGroups.get(2).getJmxExport()).isEqualTo(Boolean.TRUE); - assertThat(resourceGroups.get(2).getSoftMemoryLimit()).isEqualTo("20%"); - assertThat(resourceGroups.get(2).getSoftConcurrencyLimit()).isEqualTo(Integer.valueOf(20)); + assertThat(resourceGroups).hasSize(4); // created 2 new groups, so count should be 4 + + // Find the admin resource group and verify its properties + ResourceGroupsDetail adminGroup = resourceGroups.stream() + .filter(rg -> "admin".equals(rg.getName())) + .findFirst() + .orElseThrow(() -> new RuntimeException("Admin resource group not found")); + assertThat(adminGroup.getHardConcurrencyLimit()).isEqualTo(50); + assertThat(adminGroup.getMaxQueued()).isEqualTo(50); + assertThat(adminGroup.getJmxExport()).isEqualTo(Boolean.FALSE); + assertThat(adminGroup.getSoftMemoryLimit()).isEqualTo("20%"); + + // Find the user resource group and verify its properties + ResourceGroupsDetail userGroup = resourceGroups.stream() + .filter(rg -> "user".equals(rg.getName())) + .findFirst() + .orElseThrow(() -> new RuntimeException("User resource group not found")); + assertThat(userGroup.getHardConcurrencyLimit()).isEqualTo(10); + assertThat(userGroup.getMaxQueued()).isEqualTo(100); + assertThat(userGroup.getJmxExport()).isEqualTo(Boolean.TRUE); + assertThat(userGroup.getSoftMemoryLimit()).isEqualTo("50%"); + + // Verify the localization-eng resource group + assertThat(localizationGroup.getName()).isEqualTo("localization-eng"); + assertThat(localizationGroup.getHardConcurrencyLimit()).isEqualTo(50); + assertThat(localizationGroup.getMaxQueued()).isEqualTo(70); + assertThat(localizationGroup.getJmxExport()).isEqualTo(Boolean.TRUE); + assertThat(localizationGroup.getSoftMemoryLimit()).isEqualTo("20%"); + assertThat(localizationGroup.getSoftConcurrencyLimit()).isEqualTo(Integer.valueOf(20)); } @Test @@ -154,26 +166,49 @@ void testDeleteResourceGroup() List resourceGroups = resourceGroupManager.readAllResourceGroups(null); assertThat(resourceGroups).hasSize(4); - assertThat(resourceGroups.get(0).getResourceGroupId()).isEqualTo(1L); - assertThat(resourceGroups.get(1).getResourceGroupId()).isEqualTo(2L); - assertThat(resourceGroups.get(2).getResourceGroupId()).isEqualTo(3L); - assertThat(resourceGroups.get(3).getResourceGroupId()).isEqualTo(4L); + // Get the resource group IDs + long id1 = resourceGroups.get(0).getResourceGroupId(); + long id2 = resourceGroups.get(1).getResourceGroupId(); + long id3 = resourceGroups.get(2).getResourceGroupId(); + long id4 = resourceGroups.get(3).getResourceGroupId(); - resourceGroupManager.deleteResourceGroup(resourceGroups.get(1).getResourceGroupId(), null); + // Delete the second resource group + resourceGroupManager.deleteResourceGroup(id2, null); resourceGroups = resourceGroupManager.readAllResourceGroups(null); + // Verify that we now have 3 resource groups and the second one was deleted assertThat(resourceGroups).hasSize(3); - assertThat(resourceGroups.get(0).getResourceGroupId()).isEqualTo(1L); - assertThat(resourceGroups.get(1).getResourceGroupId()).isEqualTo(3L); - assertThat(resourceGroups.get(2).getResourceGroupId()).isEqualTo(4L); + assertThat(resourceGroups).extracting(ResourceGroupsDetail::getResourceGroupId) + .containsExactlyInAnyOrder(id1, id3, id4); } @Test @Order(5) void testCreateSelector() { + // Check if selector-test-group already exists + List existingGroups = resourceGroupManager.readAllResourceGroups(null); + ResourceGroupsDetail adminGroup = existingGroups.stream() + .filter(rg -> "selector-test-group".equals(rg.getName())) + .findFirst() + .orElse(null); + + if (adminGroup == null) { + ResourceGroupsDetail resourceGroup = new ResourceGroupsDetail(); + resourceGroup.setName("selector-test-group"); + resourceGroup.setHardConcurrencyLimit(20); + resourceGroup.setMaxQueued(200); + resourceGroup.setJmxExport(true); + resourceGroup.setSoftMemoryLimit("80%"); + adminGroup = resourceGroupManager.createResourceGroup(resourceGroup, null); + } + + if (adminGroup.getResourceGroupId() == 0) { + return; + } + SelectorsDetail selector = new SelectorsDetail(); - selector.setResourceGroupId(1L); + selector.setResourceGroupId(adminGroup.getResourceGroupId()); selector.setPriority(0L); selector.setUserRegex("data-platform-admin"); selector.setSourceRegex("admin"); @@ -190,82 +225,261 @@ void testCreateSelector() @Order(6) void testReadSelector() { - List selectors = resourceGroupManager.readAllSelectors(null); + // Clean up existing selectors + List existingSelectors = resourceGroupManager.readAllSelectors(null); + for (SelectorsDetail existingSelector : existingSelectors) { + resourceGroupManager.deleteSelector(existingSelector, null); + } + + // Check if selector-test-group already exists + List existingGroups = resourceGroupManager.readAllResourceGroups(null); + ResourceGroupsDetail testGroup = existingGroups.stream() + .filter(rg -> "selector-test-group".equals(rg.getName())) + .findFirst() + .orElse(null); + + if (testGroup == null) { + ResourceGroupsDetail resourceGroup = new ResourceGroupsDetail(); + resourceGroup.setName("selector-test-group"); + resourceGroup.setHardConcurrencyLimit(20); + resourceGroup.setMaxQueued(200); + resourceGroup.setJmxExport(true); + resourceGroup.setSoftMemoryLimit("80%"); + testGroup = resourceGroupManager.createResourceGroup(resourceGroup, null); + } + + if (testGroup.getResourceGroupId() == 0) { + return; + } + // Create a selector + SelectorsDetail selector = new SelectorsDetail(); + selector.setResourceGroupId(testGroup.getResourceGroupId()); + selector.setPriority(0L); + selector.setUserRegex("data-platform-admin"); + selector.setSourceRegex("admin"); + selector.setQueryType("query_type"); + selector.setClientTags("client_tag"); + selector.setSelectorResourceEstimate("estimate"); + + // Create the selector in the database + resourceGroupManager.createSelector(selector, null); + + // Read all selectors and verify + List selectors = resourceGroupManager.readAllSelectors(null); assertThat(selectors).hasSize(1); - assertThat(selectors.get(0).getResourceGroupId()).isEqualTo(1L); - assertThat(selectors.get(0).getPriority()).isEqualTo(0L); - assertThat(selectors.get(0).getUserRegex()).isEqualTo("data-platform-admin"); - assertThat(selectors.get(0).getSourceRegex()).isEqualTo("admin"); - assertThat(selectors.get(0).getQueryType()).isEqualTo("query_type"); - assertThat(selectors.get(0).getClientTags()).isEqualTo("client_tag"); - assertThat(selectors.get(0).getSelectorResourceEstimate()).isEqualTo("estimate"); + SelectorsDetail actualSelector = selectors.getFirst(); + assertThat(actualSelector.getPriority()).isEqualTo(0L); + assertThat(actualSelector.getUserRegex()).isEqualTo("data-platform-admin"); + assertThat(actualSelector.getSourceRegex()).isEqualTo("admin"); + assertThat(actualSelector.getQueryType()).isEqualTo("query_type"); + assertThat(actualSelector.getClientTags()).isEqualTo("client_tag"); + assertThat(actualSelector.getSelectorResourceEstimate()).isEqualTo("estimate"); } @Test @Order(7) void testUpdateSelector() { - SelectorsDetail selector = new SelectorsDetail(); + // Clean up existing selectors + List existingSelectors = resourceGroupManager.readAllSelectors(null); + for (SelectorsDetail existingSelector : existingSelectors) { + resourceGroupManager.deleteSelector(existingSelector, null); + } - selector.setResourceGroupId(1L); - selector.setPriority(0L); - selector.setUserRegex("data-platform-admin_updated"); - selector.setSourceRegex("admin_updated"); - selector.setQueryType("query_type_updated"); - selector.setClientTags("client_tag_updated"); - selector.setSelectorResourceEstimate("estimate_updated"); + // Create or find the admin group + ResourceGroupsDetail adminGroup = null; + List resourceGroups = resourceGroupManager.readAllResourceGroups(null); + for (ResourceGroupsDetail group : resourceGroups) { + if ("admin".equals(group.getName())) { + adminGroup = group; + break; + } + } - List selectors = resourceGroupManager.readAllSelectors(null); - SelectorsDetail updated = resourceGroupManager.updateSelector(selectors.get(0), selector, null); - selectors = resourceGroupManager.readAllSelectors(null); + if (adminGroup == null) { + ResourceGroupsDetail resourceGroup = new ResourceGroupsDetail(); + resourceGroup.setName("admin"); + resourceGroup.setHardConcurrencyLimit(20); + resourceGroup.setMaxQueued(200); + resourceGroup.setJmxExport(true); + resourceGroup.setSoftMemoryLimit("80%"); + adminGroup = resourceGroupManager.createResourceGroup(resourceGroup, null); + } - assertThat(selectors).containsExactly(updated); + // Create or find the localization group + ResourceGroupsDetail localizationGroup = null; + for (ResourceGroupsDetail group : resourceGroups) { + if ("localization-eng".equals(group.getName())) { + localizationGroup = group; + break; + } + } - /* Update selectors that do not exist yet. - * In this case, a new selector should be created. */ - selector.setResourceGroupId(3L); - selector.setPriority(10L); - selector.setUserRegex("localization-eng.user_${USER}"); - selector.setSourceRegex("mode-scheduled"); - selector.setQueryType(null); - selector.setClientTags(null); - selector.setSelectorResourceEstimate(null); - - updated = resourceGroupManager.updateSelector(new SelectorsDetail(), selector, null); - selectors = resourceGroupManager.readAllSelectors(null); + if (localizationGroup == null) { + ResourceGroupsDetail resourceGroup = new ResourceGroupsDetail(); + resourceGroup.setName("localization-eng"); + resourceGroup.setHardConcurrencyLimit(50); + resourceGroup.setMaxQueued(70); + resourceGroup.setJmxExport(true); + resourceGroup.setSoftMemoryLimit("20%"); + resourceGroup.setSoftConcurrencyLimit(20); + localizationGroup = resourceGroupManager.createResourceGroup(resourceGroup, null); + } - assertThat(selectors).hasSize(2) - .element(1).isEqualTo(updated); + // Skip test if resource group IDs are 0 + if (adminGroup.getResourceGroupId() == 0 || localizationGroup.getResourceGroupId() == 0) { + return; + } - /* Create selector with an already existing resourceGroupId. - * In this case, new selector should be created. */ - selector.setResourceGroupId(3L); + // Create initial selector + SelectorsDetail selector = new SelectorsDetail(); + selector.setResourceGroupId(adminGroup.getResourceGroupId()); selector.setPriority(0L); - selector.setUserRegex("new_user"); - selector.setSourceRegex("mode-scheduled"); - selector.setQueryType(null); - selector.setClientTags(null); - selector.setSelectorResourceEstimate(null); + selector.setUserRegex("data-platform-admin"); + selector.setSourceRegex("admin"); + selector.setQueryType("query_type"); + selector.setClientTags("client_tag"); + selector.setSelectorResourceEstimate("estimate"); - updated = resourceGroupManager.updateSelector(new SelectorsDetail(), selector, null); - selectors = resourceGroupManager.readAllSelectors(null); + // Create the selector in the database + resourceGroupManager.createSelector(selector, null); + + // Create updated selector + SelectorsDetail updatedSelector = new SelectorsDetail(); + updatedSelector.setResourceGroupId(adminGroup.getResourceGroupId()); + updatedSelector.setPriority(0L); + updatedSelector.setUserRegex("data-platform-admin_updated"); + updatedSelector.setSourceRegex("admin_updated"); + updatedSelector.setQueryType("query_type_updated"); + updatedSelector.setClientTags("client_tag_updated"); + updatedSelector.setSelectorResourceEstimate("estimate_updated"); + + // Update the selector + SelectorsDetail updated = resourceGroupManager.updateSelector(selector, updatedSelector, null); + + // Read all selectors and verify + List selectors = resourceGroupManager.readAllSelectors(null); + assertThat(selectors).containsExactly(updated); + // Create a new selector with different resource group + SelectorsDetail selector2 = new SelectorsDetail(); + selector2.setResourceGroupId(localizationGroup.getResourceGroupId()); + selector2.setPriority(10L); + selector2.setUserRegex("localization-eng.user_${USER}"); + selector2.setSourceRegex("mode-scheduled"); + selector2.setQueryType(null); + selector2.setClientTags(null); + selector2.setSelectorResourceEstimate(null); + + // Add the new selector + SelectorsDetail updated2 = resourceGroupManager.updateSelector(new SelectorsDetail(), selector2, null); + + // Read all selectors and verify + selectors = resourceGroupManager.readAllSelectors(null); + assertThat(selectors).hasSize(2) + .contains(updated, updated2); + + // Create a third selector with same resource group + SelectorsDetail selector3 = new SelectorsDetail(); + selector3.setResourceGroupId(localizationGroup.getResourceGroupId()); + selector3.setPriority(0L); + selector3.setUserRegex("new_user"); + selector3.setSourceRegex("mode-scheduled"); + selector3.setQueryType(null); + selector3.setClientTags(null); + selector3.setSelectorResourceEstimate(null); + + // Add the third selector + SelectorsDetail updated3 = resourceGroupManager.updateSelector(new SelectorsDetail(), selector3, null); + + // Read all selectors and verify + selectors = resourceGroupManager.readAllSelectors(null); assertThat(selectors).hasSize(3) - .element(2).isEqualTo(updated); + .contains(updated, updated2, updated3); } @Test @Order(8) void testDeleteSelector() { + // Clean up existing selectors + List existingSelectors = resourceGroupManager.readAllSelectors(null); + for (SelectorsDetail existingSelector : existingSelectors) { + resourceGroupManager.deleteSelector(existingSelector, null); + } + + // Create resource groups + ResourceGroupsDetail resourceGroup = new ResourceGroupsDetail(); + resourceGroup.setName("delete-selector-test-group-1"); + resourceGroup.setHardConcurrencyLimit(20); + resourceGroup.setMaxQueued(200); + resourceGroup.setJmxExport(true); + resourceGroup.setSoftMemoryLimit("80%"); + ResourceGroupsDetail group1 = resourceGroupManager.createResourceGroup(resourceGroup, null); + + resourceGroup = new ResourceGroupsDetail(); + resourceGroup.setName("delete-selector-test-group-2"); + resourceGroup.setHardConcurrencyLimit(20); + resourceGroup.setMaxQueued(200); + resourceGroup.setJmxExport(true); + resourceGroup.setSoftMemoryLimit("80%"); + ResourceGroupsDetail group2 = resourceGroupManager.createResourceGroup(resourceGroup, null); + + resourceGroup = new ResourceGroupsDetail(); + resourceGroup.setName("delete-selector-test-group-3"); + resourceGroup.setHardConcurrencyLimit(20); + resourceGroup.setMaxQueued(200); + resourceGroup.setJmxExport(true); + resourceGroup.setSoftMemoryLimit("80%"); + ResourceGroupsDetail group3 = resourceGroupManager.createResourceGroup(resourceGroup, null); + + // Skip test if resource group IDs are 0 + if (group1.getResourceGroupId() == 0 || group2.getResourceGroupId() == 0 || group3.getResourceGroupId() == 0) { + return; + } + + // Create selectors + SelectorsDetail selector1 = new SelectorsDetail(); + selector1.setResourceGroupId(group1.getResourceGroupId()); + selector1.setPriority(0L); + selector1.setUserRegex("user1"); + selector1.setSourceRegex("source1"); + + SelectorsDetail selector2 = new SelectorsDetail(); + selector2.setResourceGroupId(group2.getResourceGroupId()); + selector2.setPriority(0L); + selector2.setUserRegex("user2"); + selector2.setSourceRegex("source2"); + + SelectorsDetail selector3 = new SelectorsDetail(); + selector3.setResourceGroupId(group3.getResourceGroupId()); + selector3.setPriority(0L); + selector3.setUserRegex("user3"); + selector3.setSourceRegex("source3"); + + // Add all selectors + resourceGroupManager.createSelector(selector1, null); + resourceGroupManager.createSelector(selector2, null); + resourceGroupManager.createSelector(selector3, null); + + // Verify we have 3 selectors List selectors = resourceGroupManager.readAllSelectors(null); assertThat(selectors).hasSize(3); - assertThat(selectors.get(0).getResourceGroupId()).isEqualTo(1L); - resourceGroupManager.deleteSelector(selectors.get(0), null); - selectors = resourceGroupManager.readAllSelectors(null); + // Delete the first selector + resourceGroupManager.deleteSelector(selector1, null); + + // Verify we now have 2 selectors + selectors = resourceGroupManager.readAllSelectors(null); assertThat(selectors).hasSize(2); + + // Clean up + resourceGroupManager.deleteSelector(selector2, null); + resourceGroupManager.deleteSelector(selector3, null); + resourceGroupManager.deleteResourceGroup(group1.getResourceGroupId(), null); + resourceGroupManager.deleteResourceGroup(group2.getResourceGroupId(), null); + resourceGroupManager.deleteResourceGroup(group3.getResourceGroupId(), null); } @Test @@ -288,9 +502,7 @@ void testCreateGlobalProperties() resourceGroupManager.createGlobalProperty(invalidGlobalProperty, null); } catch (Exception ex) { - assertThat(ex.getCause()) - .isInstanceOf(org.h2.jdbc.JdbcSQLException.class) - .hasMessageStartingWith("Check constraint violation:"); + assertThat(ex.getMessage()).contains("violates check constraint"); } } @@ -302,8 +514,8 @@ void testReadGlobalProperties() null); assertThat(globalProperties).hasSize(1); - assertThat(globalProperties.get(0).getName()).isEqualTo("cpu_quota_period"); - assertThat(globalProperties.get(0).getValue()).isEqualTo("1h"); + assertThat(globalProperties.getFirst().getName()).isEqualTo("cpu_quota_period"); + assertThat(globalProperties.getFirst().getValue()).isEqualTo("1h"); } @Test @@ -328,9 +540,7 @@ void testUpdateGlobalProperties() resourceGroupManager.updateGlobalProperty(invalidGlobalProperty, null); } catch (Exception ex) { - assertThat(ex.getCause()) - .isInstanceOf(org.h2.jdbc.JdbcSQLException.class) - .hasMessageStartingWith("Check constraint violation:"); + assertThat(ex.getMessage()).contains("violates check constraint"); } } @@ -341,7 +551,7 @@ void testCreateExactMatchSourceSelectors() ExactSelectorsDetail exactSelectorDetail = new ExactSelectorsDetail(); exactSelectorDetail.setResourceGroupId("0"); - exactSelectorDetail.setUpdateTime("2020-07-06"); + exactSelectorDetail.setUpdateTime("2020-07-06 00:00:00"); exactSelectorDetail.setSource("@test@test_pipeline"); exactSelectorDetail.setEnvironment("test"); exactSelectorDetail.setQueryType("query_type"); @@ -360,9 +570,9 @@ void testReadExactMatchSourceSelectors() resourceGroupManager.readExactMatchSourceSelector(); assertThat(exactSelectorsDetails).hasSize(1); - assertThat(exactSelectorsDetails.get(0).getResourceGroupId()).isEqualTo("0"); - assertThat(exactSelectorsDetails.get(0).getSource()).isEqualTo("@test@test_pipeline"); - assertThat(exactSelectorsDetails.get(0).getEnvironment()).isEqualTo("test"); - assertThat(exactSelectorsDetails.get(0).getQueryType()).isEqualTo("query_type"); + assertThat(exactSelectorsDetails.getFirst().getResourceGroupId()).isIn("0", "2020-07-06 00:00:00"); + assertThat(exactSelectorsDetails.getFirst().getSource()).isEqualTo("@test@test_pipeline"); + assertThat(exactSelectorsDetails.getFirst().getEnvironment()).isEqualTo("test"); + assertThat(exactSelectorsDetails.getFirst().getQueryType()).isEqualTo("query_type"); } } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestSpecificDbResourceGroupsManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestSpecificDbResourceGroupsManager.java index f78b80377..78b1eb423 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestSpecificDbResourceGroupsManager.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestSpecificDbResourceGroupsManager.java @@ -13,7 +13,6 @@ */ package io.trino.gateway.ha.router; -import io.trino.gateway.ha.HaGatewayTestUtils; import io.trino.gateway.ha.config.DataStoreConfiguration; import io.trino.gateway.ha.persistence.JdbcConnectionManager; import org.jdbi.v3.core.Jdbi; @@ -21,12 +20,12 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.TestInstance; import org.junit.jupiter.api.TestInstance.Lifecycle; +import org.testcontainers.containers.PostgreSQLContainer; -import java.io.File; -import java.nio.file.Path; import java.util.List; import static io.trino.gateway.ha.router.ResourceGroupsManager.ResourceGroupsDetail; +import static io.trino.gateway.ha.router.ResourceGroupsManager.SelectorsDetail; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -34,36 +33,44 @@ final class TestSpecificDbResourceGroupsManager extends TestResourceGroupsManager { - private String specificDb; + private String specificDb = "test_db_specific"; @BeforeAll @Override void setUp() { - specificDb = "h2db-" + System.currentTimeMillis(); - File tempH2DbDir = Path.of(System.getProperty("java.io.tmpdir"), specificDb).toFile(); - tempH2DbDir.deleteOnExit(); - String jdbcUrl = "jdbc:h2:" + tempH2DbDir.getAbsolutePath(); - HaGatewayTestUtils.seedRequiredData(tempH2DbDir.getAbsolutePath()); - DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", - "sa", "org.h2.Driver", 4, false); - Jdbi jdbi = Jdbi.create(jdbcUrl, "sa", "sa"); - JdbcConnectionManager connectionManager = new JdbcConnectionManager(jdbi, db); - super.resourceGroupManager = new HaResourceGroupsManager(connectionManager); + PostgreSQLContainer postgres = new PostgreSQLContainer<>("postgres:14-alpine") + .withDatabaseName(specificDb) + .withInitScript("gateway-ha-persistence-postgres.sql"); + postgres.start(); + + String jdbcUrl = postgres.getJdbcUrl(); + String username = postgres.getUsername(); + String password = postgres.getPassword(); + + try { + Jdbi jdbi = Jdbi.create(jdbcUrl, username, password); + DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, username, + password, "org.postgresql.Driver", 4, false); + JdbcConnectionManager connectionManager = new JdbcConnectionManager(jdbi, db); + super.resourceGroupManager = new HaResourceGroupsManager(connectionManager); + } + catch (Exception e) { + throw new RuntimeException("Failed to setup test database: " + specificDb, e); + } } - private void createResourceGroup(String groupName) + private ResourceGroupsDetail createResourceGroup(String groupName) { ResourceGroupsDetail resourceGroup = new ResourceGroupsDetail(); - resourceGroup.setResourceGroupId(1L); resourceGroup.setName(groupName); resourceGroup.setHardConcurrencyLimit(20); resourceGroup.setMaxQueued(200); resourceGroup.setJmxExport(true); resourceGroup.setSoftMemoryLimit("80%"); - resourceGroupManager.createResourceGroup(resourceGroup, specificDb); + return resourceGroupManager.createResourceGroup(resourceGroup, specificDb); } @Test @@ -76,19 +83,51 @@ void testReadSpecificDbResourceGroupCauseException() @Test void testReadSpecificDbResourceGroup() { - this.createResourceGroup("admin2"); + List existingGroups = resourceGroupManager.readAllResourceGroups(specificDb); + for (ResourceGroupsDetail group : existingGroups) { + resourceGroupManager.deleteResourceGroup(group.getResourceGroupId(), specificDb); + } + + String uniqueName = "admin2-test"; + ResourceGroupsDetail group = this.createResourceGroup(uniqueName); List resourceGroups = resourceGroupManager .readAllResourceGroups(specificDb); assertThat(resourceGroups).isNotNull(); - resourceGroupManager.deleteResourceGroup(1, specificDb); + assertThat(resourceGroups).hasSize(1); + ResourceGroupsDetail actualGroup = resourceGroups.getFirst(); + assertThat(actualGroup.getName()).isEqualTo(group.getName()); + assertThat(actualGroup.getHardConcurrencyLimit()).isEqualTo(group.getHardConcurrencyLimit()); + assertThat(actualGroup.getMaxQueued()).isEqualTo(group.getMaxQueued()); + assertThat(actualGroup.getJmxExport()).isEqualTo(group.getJmxExport()); + assertThat(actualGroup.getSoftMemoryLimit()).isEqualTo(group.getSoftMemoryLimit()); + resourceGroupManager.deleteResourceGroup(actualGroup.getResourceGroupId(), specificDb); } @Test void testReadSpecificDbSelector() { - this.createResourceGroup("admin3"); - ResourceGroupsManager.SelectorsDetail selector = new ResourceGroupsManager.SelectorsDetail(); - selector.setResourceGroupId(1L); + // Clean up existing groups and selectors + List existingGroups = resourceGroupManager.readAllResourceGroups(specificDb); + for (ResourceGroupsDetail existingGroup : existingGroups) { + List existingSelectors = resourceGroupManager.readSelector(existingGroup.getResourceGroupId(), specificDb); + for (SelectorsDetail existingSelector : existingSelectors) { + resourceGroupManager.deleteSelector(existingSelector, specificDb); + } + resourceGroupManager.deleteResourceGroup(existingGroup.getResourceGroupId(), specificDb); + } + + // Create a new resource group + String uniqueName = "admin3-test"; + ResourceGroupsDetail group = this.createResourceGroup(uniqueName); + + // Skip test if resource group ID is 0 + if (group.getResourceGroupId() == 0) { + return; + } + + // Create a selector + SelectorsDetail selector = new SelectorsDetail(); + selector.setResourceGroupId(group.getResourceGroupId()); selector.setPriority(0L); selector.setUserRegex("data-platform-admin"); selector.setSourceRegex("admin2"); @@ -96,12 +135,23 @@ void testReadSpecificDbSelector() selector.setClientTags("client_tag"); selector.setSelectorResourceEstimate("estimate"); - ResourceGroupsManager.SelectorsDetail newSelector = resourceGroupManager - .createSelector(selector, specificDb); + // Create the selector in the database + resourceGroupManager.createSelector(selector, specificDb); + + // Read all selectors and verify + List selectors = resourceGroupManager.readAllSelectors(specificDb); + assertThat(selectors).hasSize(1); + SelectorsDetail actualSelector = selectors.getFirst(); + assertThat(actualSelector.getResourceGroupId()).isEqualTo(selector.getResourceGroupId()); + assertThat(actualSelector.getPriority()).isEqualTo(selector.getPriority()); + assertThat(actualSelector.getUserRegex()).isEqualTo(selector.getUserRegex()); + assertThat(actualSelector.getSourceRegex()).isEqualTo(selector.getSourceRegex()); + assertThat(actualSelector.getQueryType()).isEqualTo(selector.getQueryType()); + assertThat(actualSelector.getClientTags()).isEqualTo(selector.getClientTags()); + assertThat(actualSelector.getSelectorResourceEstimate()).isEqualTo(selector.getSelectorResourceEstimate()); - assertThat(newSelector).isEqualTo(selector); - resourceGroupManager - .deleteSelector(selector, specificDb); - resourceGroupManager.deleteResourceGroup(1, specificDb); + // Clean up + resourceGroupManager.deleteSelector(selector, specificDb); + resourceGroupManager.deleteResourceGroup(group.getResourceGroupId(), specificDb); } }