diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index c0053a6ff..a621f5aa1 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -34,7 +34,3 @@ services: timeout: 1s retries: 60 start_period: 10s - volumes: - - target: /docker-entrypoint-initdb.d/1-gateway-ha-persistence-postgres.sql - source: ../gateway-ha/src/main/resources/gateway-ha-persistence-postgres.sql - type: bind diff --git a/docs/installation.md b/docs/installation.md index aaaf98bb8..ba601c45f 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -38,16 +38,37 @@ distribution is installed. ### Backend database -Trino Gateway requires a MySQL or PostgreSQL database. +Trino Gateway requires a MySQL or PostgreSQL database. Database initialization +is performed automatically when the Trino Gateway process starts. Migrations +are performed using `Flyway`. -Use the following scripts in the `gateway-ha/src/main/resources/` folder to -initialize the database: - -* `gateway-ha-persistence-mysql.sql` for MySQL -* `gateway-ha-persistence-postgres.sql` for PostgreSQL +The migration files can viewed in the `gateway-ha/src/main/resources/` folder. +Each database type supported has its own sub-folder. The files are also included in the JAR file. +If you do not want migrations to be performed automatically on startup, then +you can set `runMigrationsEnabled` to `false` in the data store configuration. +For example: + +```yaml +dataStore: + jdbcUrl: jdbc:postgresql://postgres:5432/trino_gateway_db + user: USER + password: PASSWORD + driver: org.postgresql.Driver + queryHistoryHoursRetention: 24 + runMigrationsEnabled: false +``` + +`Flyway` uses a transactional lock in databases that support it such as +[PostgreSQL](https://documentation.red-gate.com/fd/postgresql-database-235241807.html#). +In the scenario where multiple Trino Gateway instances are running and sharing +the same backend database, the first Trino Gateway instance to start will get +the lock and run the database migrations with `Flyway`. Other Trino Gateway +instances might fail during startup while migrations are running but once migrations +are completed they will start as expected. + ### Trino clusters The proxied Trino clusters behind the Trino Gateway must support the Trino JDBC diff --git a/gateway-ha/pom.xml b/gateway-ha/pom.xml index 7149d36e0..8ef0ea2cb 100644 --- a/gateway-ha/pom.xml +++ b/gateway-ha/pom.xml @@ -21,6 +21,7 @@ https://registry.npmmirror.com + 11.0.1 4.1.0 5.14.2 4.12.0 @@ -237,6 +238,12 @@ 1.78.1 + + org.flywaydb + flyway-core + ${dep.flyway.version} + + org.glassfish.jersey.core jersey-server @@ -290,6 +297,20 @@ runtime + + org.flywaydb + flyway-database-postgresql + ${dep.flyway.version} + runtime + + + + org.flywaydb + flyway-mysql + ${dep.flyway.version} + runtime + + org.mvel mvel2 @@ -371,6 +392,12 @@ test + + org.testcontainers + jdbc + test + + org.testcontainers mysql diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/HaGatewayLauncher.java b/gateway-ha/src/main/java/io/trino/gateway/ha/HaGatewayLauncher.java index 51b05fd93..59b3b7872 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/HaGatewayLauncher.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/HaGatewayLauncher.java @@ -32,6 +32,7 @@ import io.airlift.units.Duration; import io.trino.gateway.baseapp.BaseApp; import io.trino.gateway.ha.config.HaGatewayConfiguration; +import io.trino.gateway.ha.persistence.FlywayMigration; import org.weakref.jmx.guice.MBeanModule; import java.nio.file.Files; @@ -109,6 +110,7 @@ public static void main(String[] args) } String config = Files.readString(Path.of(args[0])); HaGatewayConfiguration haGatewayConfiguration = objectMapper.readValue(replaceEnvironmentVariables(config), HaGatewayConfiguration.class); + FlywayMigration.migrate(haGatewayConfiguration.getDataStore()); List modules = addModules(haGatewayConfiguration); new HaGatewayLauncher().start(modules, haGatewayConfiguration); } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/config/DataStoreConfiguration.java b/gateway-ha/src/main/java/io/trino/gateway/ha/config/DataStoreConfiguration.java index cd8ca14bf..cd04c5335 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/config/DataStoreConfiguration.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/config/DataStoreConfiguration.java @@ -20,14 +20,16 @@ public class DataStoreConfiguration private String password; private String driver; private Integer queryHistoryHoursRetention = 4; + private boolean runMigrationsEnabled = true; - public DataStoreConfiguration(String jdbcUrl, String user, String password, String driver, Integer queryHistoryHoursRetention) + public DataStoreConfiguration(String jdbcUrl, String user, String password, String driver, Integer queryHistoryHoursRetention, boolean runMigrationsEnabled) { this.jdbcUrl = jdbcUrl; this.user = user; this.password = password; this.driver = driver; this.queryHistoryHoursRetention = queryHistoryHoursRetention; + this.runMigrationsEnabled = runMigrationsEnabled; } public DataStoreConfiguration() {} @@ -81,4 +83,14 @@ public void setQueryHistoryHoursRetention(Integer queryHistoryHoursRetention) { this.queryHistoryHoursRetention = queryHistoryHoursRetention; } + + public boolean isRunMigrationsEnabled() + { + return this.runMigrationsEnabled; + } + + public void setRunMigrationsEnabled(boolean runMigrationsEnabled) + { + this.runMigrationsEnabled = runMigrationsEnabled; + } } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/FlywayMigration.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/FlywayMigration.java new file mode 100644 index 000000000..3294af2db --- /dev/null +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/FlywayMigration.java @@ -0,0 +1,57 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.gateway.ha.persistence; + +import io.airlift.log.Logger; +import io.trino.gateway.ha.config.DataStoreConfiguration; +import org.flywaydb.core.Flyway; +import org.flywaydb.core.api.output.MigrateResult; + +import static java.lang.String.format; + +public class FlywayMigration +{ + private static final Logger log = Logger.get(FlywayMigration.class); + + private FlywayMigration() {} + + private static String getLocation(String configDbUrl) + { + if (configDbUrl.startsWith("jdbc:postgresql")) { + return "postgresql"; + } + if (configDbUrl.startsWith("jdbc:mysql")) { + return "mysql"; + } + throw new IllegalArgumentException(format("Invalid JDBC URL: %s. Only PostgreSQL and MySQL are supported.", configDbUrl)); + } + + public static void migrate(DataStoreConfiguration config) + { + if (!config.isRunMigrationsEnabled()) { + log.info("Skip migrations as automatic migrations are disabled"); + return; + } + log.info("Performing migrations..."); + Flyway flyway = Flyway.configure() + .dataSource(config.getJdbcUrl(), config.getUser(), config.getPassword()) + .locations(getLocation(config.getJdbcUrl())) + .baselineOnMigrate(true) + .baselineVersion("0") + .load(); + + MigrateResult migrations = flyway.migrate(); + log.info("Performed %s migrations", migrations.migrationsExecuted); + } +} diff --git a/gateway-ha/src/main/resources/mysql/V1__create_schema.sql b/gateway-ha/src/main/resources/mysql/V1__create_schema.sql new file mode 100644 index 000000000..51d7badf6 --- /dev/null +++ b/gateway-ha/src/main/resources/mysql/V1__create_schema.sql @@ -0,0 +1,78 @@ +CREATE TABLE IF NOT EXISTS gateway_backend ( +name VARCHAR(256) PRIMARY KEY, +routing_group VARCHAR (256), +backend_url VARCHAR (256), +external_url VARCHAR (256), +active BOOLEAN +); + +CREATE TABLE IF NOT EXISTS query_history ( +query_id VARCHAR(256) PRIMARY KEY, +query_text VARCHAR (256), +created bigint, +backend_url VARCHAR (256), +user_name VARCHAR(256), +source VARCHAR(256) +); +CREATE INDEX query_history_created_idx ON query_history(created); + +CREATE TABLE IF NOT EXISTS resource_groups ( + resource_group_id BIGINT NOT NULL AUTO_INCREMENT, + name VARCHAR(250) NOT NULL UNIQUE, + + -- OPTIONAL POLICY CONTROLS + parent BIGINT NULL, + jmx_export BOOLEAN NULL, + scheduling_policy VARCHAR(128) NULL, + scheduling_weight INT NULL, + + -- REQUIRED QUOTAS + soft_memory_limit VARCHAR(128) NOT NULL, + max_queued INT NOT NULL, + hard_concurrency_limit INT NOT NULL, + + -- OPTIONAL QUOTAS + soft_concurrency_limit INT NULL, + soft_cpu_limit VARCHAR(128) NULL, + hard_cpu_limit VARCHAR(128) NULL, + environment VARCHAR(128) NULL, + + PRIMARY KEY(resource_group_id), + FOREIGN KEY (parent) REFERENCES resource_groups (resource_group_id) +); + +CREATE TABLE IF NOT EXISTS selectors ( + resource_group_id BIGINT NOT NULL, + priority BIGINT NOT NULL, + + -- Regex fields -- these will be used as a regular expression pattern to + -- match against the field of the same name on queries + user_regex VARCHAR(512), + source_regex VARCHAR(512), + + -- Selector fields -- these must match exactly. + query_type VARCHAR(512), + client_tags VARCHAR(512), + selector_resource_estimate VARCHAR(1024), + + FOREIGN KEY (resource_group_id) REFERENCES resource_groups(resource_group_id) +); + +CREATE TABLE IF NOT EXISTS resource_groups_global_properties ( + name VARCHAR(128) NOT NULL PRIMARY KEY, + value VARCHAR(512) NULL, + CHECK (name in ('cpu_quota_period')) +); + +CREATE TABLE IF NOT EXISTS exact_match_source_selectors ( + resource_group_id VARCHAR(256) NOT NULL, + update_time DATETIME NOT NULL, + + -- Selector fields which must exactly match a query + source VARCHAR(512) NOT NULL, + environment VARCHAR(128), + query_type VARCHAR(512), + + PRIMARY KEY (environment, source(128), query_type), + UNIQUE (source(128), environment, query_type(128), resource_group_id) +); diff --git a/gateway-ha/src/main/resources/postgresql/V1__create_schema.sql b/gateway-ha/src/main/resources/postgresql/V1__create_schema.sql new file mode 100644 index 000000000..a875a0fa5 --- /dev/null +++ b/gateway-ha/src/main/resources/postgresql/V1__create_schema.sql @@ -0,0 +1,78 @@ +CREATE TABLE IF NOT EXISTS gateway_backend ( +name VARCHAR(256) PRIMARY KEY, +routing_group VARCHAR (256), +backend_url VARCHAR (256), +external_url VARCHAR (256), +active BOOLEAN +); + +CREATE TABLE IF NOT EXISTS query_history ( +query_id VARCHAR(256) PRIMARY KEY, +query_text VARCHAR (256), +created bigint, +backend_url VARCHAR (256), +user_name VARCHAR(256), +source VARCHAR(256) +); +CREATE INDEX IF NOT EXISTS query_history_created_idx ON query_history(created); + +CREATE TABLE IF NOT EXISTS resource_groups ( + resource_group_id SERIAL, + name VARCHAR(250) NOT NULL UNIQUE, + + -- OPTIONAL POLICY CONTROLS + parent BIGINT NULL, + jmx_export BOOLEAN NULL, + scheduling_policy VARCHAR(128) NULL, + scheduling_weight INT NULL, + + -- REQUIRED QUOTAS + soft_memory_limit VARCHAR(128) NOT NULL, + max_queued INT NOT NULL, + hard_concurrency_limit INT NOT NULL, + + -- OPTIONAL QUOTAS + soft_concurrency_limit INT NULL, + soft_cpu_limit VARCHAR(128) NULL, + hard_cpu_limit VARCHAR(128) NULL, + environment VARCHAR(128) NULL, + + PRIMARY KEY(resource_group_id), + FOREIGN KEY (parent) REFERENCES resource_groups (resource_group_id) +); + +CREATE TABLE IF NOT EXISTS selectors ( + resource_group_id BIGINT NOT NULL, + priority BIGINT NOT NULL, + + -- Regex fields -- these will be used as a regular expression pattern to + -- match against the field of the same name on queries + user_regex VARCHAR(512), + source_regex VARCHAR(512), + + -- Selector fields -- these must match exactly. + query_type VARCHAR(512), + client_tags VARCHAR(512), + selector_resource_estimate VARCHAR(1024), + + FOREIGN KEY (resource_group_id) REFERENCES resource_groups(resource_group_id) +); + +CREATE TABLE IF NOT EXISTS resource_groups_global_properties ( + name VARCHAR(128) NOT NULL PRIMARY KEY, + value VARCHAR(512) NULL, + CHECK (name in ('cpu_quota_period')) +); + +CREATE TABLE IF NOT EXISTS exact_match_source_selectors ( + resource_group_id VARCHAR(256) NOT NULL, + update_time TIMESTAMP NOT NULL, + + -- Selector fields which must exactly match a query + source VARCHAR(512) NOT NULL, + environment VARCHAR(128), + query_type VARCHAR(128), -- (reduced from 512) + + PRIMARY KEY (environment, source, query_type), + UNIQUE (source, environment, query_type, resource_group_id) +); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestTrinoResource.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestTrinoResource.java index 856ebf2a9..7e599f30d 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/TestTrinoResource.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/TestTrinoResource.java @@ -61,7 +61,7 @@ void setup() // Setup resource group manager String jdbcUrl = "jdbc:h2:" + testConfig.h2DbFilePath(); - DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", "sa", "org.h2.Driver", 4); + DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", "sa", "org.h2.Driver", 4, false); Jdbi jdbi = Jdbi.create(jdbcUrl, "sa", "sa"); connectionManager = new JdbcConnectionManager(jdbi, db); resourceGroupManager = new HaResourceGroupsManager(connectionManager); 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 d4b36cc0a..5ff7d63df 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 @@ -30,7 +30,7 @@ public static JdbcConnectionManager createTestingJdbcConnectionManager() tempH2DbDir.deleteOnExit(); String jdbcUrl = "jdbc:h2:" + tempH2DbDir.getAbsolutePath(); HaGatewayTestUtils.seedRequiredData(new HaGatewayTestUtils.TestConfig("", tempH2DbDir.getAbsolutePath())); - DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", "sa", "org.h2.Driver", 4); + DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", "sa", "org.h2.Driver", 4, false); Jdbi jdbi = Jdbi.create(jdbcUrl, "sa", "sa"); return new JdbcConnectionManager(jdbi, db); } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/BaseTestDatabaseMigrations.java b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/BaseTestDatabaseMigrations.java new file mode 100644 index 000000000..168baed6a --- /dev/null +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/BaseTestDatabaseMigrations.java @@ -0,0 +1,157 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.gateway.ha.persistence; + +import io.trino.gateway.ha.config.DataStoreConfiguration; +import org.jdbi.v3.core.Handle; +import org.jdbi.v3.core.Jdbi; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.parallel.Execution; +import org.junit.jupiter.api.parallel.Isolated; +import org.testcontainers.containers.JdbcDatabaseContainer; + +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS; +import static org.junit.jupiter.api.parallel.ExecutionMode.SAME_THREAD; + +@TestInstance(PER_CLASS) +@Execution(SAME_THREAD) +@Isolated +public abstract class BaseTestDatabaseMigrations +{ + protected final JdbcDatabaseContainer container = startContainer(); + protected final Jdbi jdbi = Jdbi.create(container.getJdbcUrl(), container.getUsername(), container.getPassword()); + + protected abstract JdbcDatabaseContainer startContainer(); + + protected abstract String getDriver(); + + protected abstract void createGatewaySchema(); + + @AfterAll + public final void close() + { + container.close(); + } + + @Test + public void testMigrationWithEmptyDatabase() + { + DataStoreConfiguration config = new DataStoreConfiguration( + container.getJdbcUrl(), + container.getUsername(), + container.getPassword(), + getDriver(), + 4, + true); + FlywayMigration.migrate(config); + verifyGatewaySchema(0); + + dropAllTables(); + } + + @Test + public void testMigrationWithNonemptyDatabase() + { + DataStoreConfiguration config = new DataStoreConfiguration( + container.getJdbcUrl(), + container.getUsername(), + container.getPassword(), + getDriver(), + 4, + true); + String t1Create = "CREATE TABLE t1 (id INT)"; + String t2Create = "CREATE TABLE t2 (id INT)"; + Handle jdbiHandle = jdbi.open(); + jdbiHandle.execute(t1Create); + jdbiHandle.execute(t2Create); + FlywayMigration.migrate(config); + verifyGatewaySchema(0); + String t1Drop = "DROP TABLE t1"; + String t2Drop = "DROP TABLE t2"; + jdbiHandle.execute(t1Drop); + jdbiHandle.execute(t2Drop); + jdbiHandle.close(); + + dropAllTables(); + } + + @Test + public void testMigrationWithExistingGatewaySchema() + { + createGatewaySchema(); + // add a row to one of the existing tables before migration + jdbi.withHandle(handle -> + handle.execute("INSERT INTO resource_groups_global_properties VALUES ('a_name', 'a_value')")); + DataStoreConfiguration config = new DataStoreConfiguration( + container.getJdbcUrl(), + container.getUsername(), + container.getPassword(), + getDriver(), + 4, + true); + FlywayMigration.migrate(config); + verifyGatewaySchema(1); + dropAllTables(); + } + + protected void verifyGatewaySchema(int expectedPropertiesCount) + { + verifyResultSetCount("SELECT name FROM gateway_backend", 0); + verifyResultSetCount("SELECT query_id FROM query_history", 0); + verifyResultSetCount("SELECT name FROM resource_groups_global_properties", expectedPropertiesCount); + verifyResultSetCount("SELECT name FROM resource_groups", 0); + verifyResultSetCount("SELECT user_regex FROM selectors", 0); + verifyResultSetCount("SELECT environment FROM exact_match_source_selectors", 0); + } + + private void verifyResultSetCount(String sql, int expectedCount) + { + List results = jdbi.withHandle(handle -> + handle.createQuery(sql).mapTo(String.class).list()); + assertThat(results).hasSize(expectedCount); + } + + protected void dropAllTables() + { + String gatewayBackendTable = "DROP TABLE IF EXISTS gateway_backend"; + String queryHistoryTable = "DROP TABLE IF EXISTS query_history"; + String propertiesTable = "DROP TABLE IF EXISTS resource_groups_global_properties"; + String resourceGroupsTable = "DROP TABLE IF EXISTS resource_groups"; + String selectorsTable = "DROP TABLE IF EXISTS selectors"; + String exactMatchTable = "DROP TABLE IF EXISTS exact_match_source_selectors"; + String flywayHistoryTable = "DROP TABLE IF EXISTS flyway_schema_history"; + Handle jdbiHandle = jdbi.open(); + String sql = String.format("SELECT 1 FROM information_schema.tables WHERE table_schema = '%s'", getTestSchema()); + verifyResultSetCount(sql, 7); + jdbiHandle.execute(gatewayBackendTable); + jdbiHandle.execute(queryHistoryTable); + jdbiHandle.execute(propertiesTable); + jdbiHandle.execute(selectorsTable); + jdbiHandle.execute(resourceGroupsTable); + jdbiHandle.execute(exactMatchTable); + jdbiHandle.execute(flywayHistoryTable); + verifyResultSetCount(sql, 0); + jdbiHandle.close(); + } + + protected String getTestSchema() + { + return "public"; + } +} diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsMySql.java b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsMySql.java new file mode 100644 index 000000000..67b16be79 --- /dev/null +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsMySql.java @@ -0,0 +1,111 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.gateway.ha.persistence; + +import org.jdbi.v3.core.Handle; +import org.testcontainers.containers.JdbcDatabaseContainer; +import org.testcontainers.containers.MySQLContainer; + +public class TestDatabaseMigrationsMySql + extends BaseTestDatabaseMigrations +{ + @Override + protected final JdbcDatabaseContainer startContainer() + { + JdbcDatabaseContainer container = new MySQLContainer<>("mysql:8.0.36"); + container.start(); + return container; + } + + @Override + protected final String getDriver() + { + return "com.mysql.cj.jdbc.Driver"; + } + + @Override + protected final String getTestSchema() + { + return "test"; + } + + @Override + protected void createGatewaySchema() + { + String gatewayBackendTable = "CREATE TABLE gateway_backend (\n" + + " name VARCHAR(256) PRIMARY KEY,\n" + + " routing_group VARCHAR (256),\n" + + " backend_url VARCHAR (256),\n" + + " external_url VARCHAR (256),\n" + + " active BOOLEAN\n" + + ");"; + String queryHistoryTable = "CREATE TABLE query_history (\n" + + " query_id VARCHAR(256) PRIMARY KEY,\n" + + " query_text VARCHAR (256),\n" + + " created bigint,\n" + + " backend_url VARCHAR (256),\n" + + " user_name VARCHAR(256),\n" + + " source VARCHAR(256)\n" + + ");"; + String propertiesTable = "CREATE TABLE resource_groups_global_properties (\n" + + " name VARCHAR(128) NOT NULL PRIMARY KEY,\n" + + " value VARCHAR(512) NULL,\n" + + " CHECK (name in ('cpu_quota_period', 'a_name', 'a_value'))\n" + + ");"; + String resourceGroupsTable = "CREATE TABLE resource_groups (\n" + + " resource_group_id BIGINT NOT NULL AUTO_INCREMENT,\n" + + " name VARCHAR(250) NOT NULL,\n" + + " soft_memory_limit VARCHAR(128) NOT NULL,\n" + + " max_queued INT NOT NULL,\n" + + " soft_concurrency_limit INT NULL,\n" + + " hard_concurrency_limit INT NOT NULL,\n" + + " scheduling_policy VARCHAR(128) NULL,\n" + + " scheduling_weight INT NULL,\n" + + " jmx_export BOOLEAN NULL,\n" + + " soft_cpu_limit VARCHAR(128) NULL,\n" + + " hard_cpu_limit VARCHAR(128) NULL,\n" + + " parent BIGINT NULL,\n" + + " environment VARCHAR(128) NULL,\n" + + " PRIMARY KEY (resource_group_id),\n" + + " FOREIGN KEY (parent) REFERENCES resource_groups (resource_group_id) ON DELETE CASCADE\n" + + ");"; + String selectorsTable = "CREATE TABLE selectors (\n" + + " resource_group_id BIGINT NOT NULL,\n" + + " priority BIGINT NOT NULL,\n" + + " user_regex VARCHAR(512),\n" + + " source_regex VARCHAR(512),\n" + + " query_type VARCHAR(512),\n" + + " client_tags VARCHAR(512),\n" + + " selector_resource_estimate VARCHAR(1024),\n" + + " FOREIGN KEY (resource_group_id) REFERENCES resource_groups (resource_group_id) ON DELETE CASCADE\n" + + ");"; + String exactMatchSourceSelectorsTable = "CREATE TABLE exact_match_source_selectors (\n" + + " resource_group_id VARCHAR(256) NOT NULL,\n" + + " update_time DATETIME NOT NULL,\n" + + " source VARCHAR(512) NOT NULL,\n" + + " environment VARCHAR(128),\n" + + " query_type VARCHAR(512),\n" + + " PRIMARY KEY (environment, source(128), query_type),\n" + + " UNIQUE (source(128), environment, query_type(128), resource_group_id)\n" + + ");"; + Handle jdbiHandle = jdbi.open(); + jdbiHandle.execute(gatewayBackendTable); + jdbiHandle.execute(queryHistoryTable); + jdbiHandle.execute(propertiesTable); + jdbiHandle.execute(resourceGroupsTable); + jdbiHandle.execute(selectorsTable); + jdbiHandle.execute(exactMatchSourceSelectorsTable); + jdbiHandle.close(); + } +} diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsPostgresql.java b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsPostgresql.java new file mode 100644 index 000000000..85f1c1fc9 --- /dev/null +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsPostgresql.java @@ -0,0 +1,105 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.gateway.ha.persistence; + +import org.jdbi.v3.core.Handle; +import org.testcontainers.containers.JdbcDatabaseContainer; +import org.testcontainers.containers.PostgreSQLContainer; + +public class TestDatabaseMigrationsPostgresql + extends BaseTestDatabaseMigrations +{ + @Override + protected final JdbcDatabaseContainer startContainer() + { + JdbcDatabaseContainer container = new PostgreSQLContainer<>("postgres:11"); + container.start(); + return container; + } + + @Override + protected final String getDriver() + { + return "org.postgresql.Driver"; + } + + @Override + protected void createGatewaySchema() + { + String gatewayBackendTable = "CREATE TABLE gateway_backend (\n" + + " name VARCHAR(256) PRIMARY KEY,\n" + + " routing_group VARCHAR (256),\n" + + " backend_url VARCHAR (256),\n" + + " external_url VARCHAR (256),\n" + + " active BOOLEAN\n" + + ");"; + String queryHistoryTable = "CREATE TABLE query_history (\n" + + " query_id VARCHAR(256) PRIMARY KEY,\n" + + " query_text VARCHAR (256),\n" + + " created bigint,\n" + + " backend_url VARCHAR (256),\n" + + " user_name VARCHAR(256),\n" + + " source VARCHAR(256)\n" + + ");"; + String propertiesTable = "CREATE TABLE resource_groups_global_properties (\n" + + " name VARCHAR(128) NOT NULL PRIMARY KEY,\n" + + " value VARCHAR(512) NULL,\n" + + " CHECK (name in ('cpu_quota_period', 'a_name', 'a_value'))\n" + + ");"; + String resourceGroupsTable = "CREATE TABLE resource_groups (\n" + + " resource_group_id SERIAL,\n" + + " name VARCHAR(250) NOT NULL,\n" + + " soft_memory_limit VARCHAR(128) NOT NULL,\n" + + " max_queued INT NOT NULL,\n" + + " soft_concurrency_limit INT NULL,\n" + + " hard_concurrency_limit INT NOT NULL,\n" + + " scheduling_policy VARCHAR(128) NULL,\n" + + " scheduling_weight INT NULL,\n" + + " jmx_export BOOLEAN NULL,\n" + + " soft_cpu_limit VARCHAR(128) NULL,\n" + + " hard_cpu_limit VARCHAR(128) NULL,\n" + + " parent BIGINT NULL,\n" + + " environment VARCHAR(128) NULL,\n" + + " PRIMARY KEY (resource_group_id),\n" + + " FOREIGN KEY (parent) REFERENCES resource_groups (resource_group_id) ON DELETE CASCADE\n" + + ");"; + String selectorsTable = "CREATE TABLE selectors (\n" + + " resource_group_id BIGINT NOT NULL,\n" + + " priority BIGINT NOT NULL,\n" + + " user_regex VARCHAR(512),\n" + + " source_regex VARCHAR(512),\n" + + " query_type VARCHAR(512),\n" + + " client_tags VARCHAR(512),\n" + + " selector_resource_estimate VARCHAR(1024),\n" + + " FOREIGN KEY (resource_group_id) REFERENCES resource_groups (resource_group_id) ON DELETE CASCADE\n" + + ");"; + String exactMatchSourceSelectorsTable = "CREATE TABLE exact_match_source_selectors (\n" + + " resource_group_id VARCHAR(256) NOT NULL,\n" + + " update_time TIMESTAMP NOT NULL,\n" + + " source VARCHAR(512) NOT NULL,\n" + + " environment VARCHAR(128),\n" + + " query_type VARCHAR(512),\n" + + " PRIMARY KEY (environment, source, query_type),\n" + + " UNIQUE (source, environment, query_type, resource_group_id)\n" + + ");"; + Handle jdbiHandle = jdbi.open(); + jdbiHandle.execute(gatewayBackendTable); + jdbiHandle.execute(queryHistoryTable); + jdbiHandle.execute(propertiesTable); + jdbiHandle.execute(resourceGroupsTable); + jdbiHandle.execute(selectorsTable); + jdbiHandle.execute(exactMatchSourceSelectorsTable); + jdbiHandle.close(); + } +} 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 68322ace2..e2ce6096a 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 @@ -47,7 +47,7 @@ void setUp() HaGatewayTestUtils.seedRequiredData( new HaGatewayTestUtils.TestConfig("", tempH2DbDir.getAbsolutePath())); DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", - "sa", "org.h2.Driver", 4); + "sa", "org.h2.Driver", 4, false); Jdbi jdbi = Jdbi.create(jdbcUrl, "sa", "sa"); JdbcConnectionManager connectionManager = new JdbcConnectionManager(jdbi, db); super.resourceGroupManager = new HaResourceGroupsManager(connectionManager); diff --git a/gateway-ha/src/test/resources/auth/auth-test-config.yml b/gateway-ha/src/test/resources/auth/auth-test-config.yml index f5bc2c31a..4d0c16da3 100644 --- a/gateway-ha/src/test/resources/auth/auth-test-config.yml +++ b/gateway-ha/src/test/resources/auth/auth-test-config.yml @@ -7,6 +7,7 @@ dataStore: user: sa password: sa driver: org.h2.Driver + runMigrationsEnabled: false modules: - io.trino.gateway.ha.module.HaGatewayProviderModule diff --git a/gateway-ha/src/test/resources/auth/oauth-test-config.yml b/gateway-ha/src/test/resources/auth/oauth-test-config.yml index 638684fa0..c9b3fb31a 100644 --- a/gateway-ha/src/test/resources/auth/oauth-test-config.yml +++ b/gateway-ha/src/test/resources/auth/oauth-test-config.yml @@ -11,6 +11,7 @@ dataStore: user: sa password: sa driver: org.h2.Driver + runMigrationsEnabled: false modules: - io.trino.gateway.ha.module.HaGatewayProviderModule diff --git a/gateway-ha/src/test/resources/test-config-template.yml b/gateway-ha/src/test/resources/test-config-template.yml index 81ea1b461..cd842b3a8 100644 --- a/gateway-ha/src/test/resources/test-config-template.yml +++ b/gateway-ha/src/test/resources/test-config-template.yml @@ -8,6 +8,7 @@ dataStore: user: sa password: sa driver: org.h2.Driver + runMigrationsEnabled: false modules: - io.trino.gateway.ha.module.HaGatewayProviderModule diff --git a/gateway-ha/src/test/resources/test-config-with-routing-template.yml b/gateway-ha/src/test/resources/test-config-with-routing-template.yml index 817b93e20..bcd999038 100644 --- a/gateway-ha/src/test/resources/test-config-with-routing-template.yml +++ b/gateway-ha/src/test/resources/test-config-with-routing-template.yml @@ -7,6 +7,7 @@ dataStore: user: sa password: sa driver: org.h2.Driver + runMigrationsEnabled: false modules: - io.trino.gateway.ha.module.HaGatewayProviderModule diff --git a/gateway-ha/src/test/resources/test-config-without-x-forwarded-template.yml b/gateway-ha/src/test/resources/test-config-without-x-forwarded-template.yml index 4be161476..a44f8e64d 100644 --- a/gateway-ha/src/test/resources/test-config-without-x-forwarded-template.yml +++ b/gateway-ha/src/test/resources/test-config-without-x-forwarded-template.yml @@ -7,6 +7,7 @@ dataStore: user: sa password: sa driver: org.h2.Driver + runMigrationsEnabled: false modules: - io.trino.gateway.ha.module.HaGatewayProviderModule