Skip to content

Commit

Permalink
Add PostgreSQL migration files for DB resource groups
Browse files Browse the repository at this point in the history
  • Loading branch information
posulliv authored and hashhar committed Jan 20, 2022
1 parent f009a46 commit 417c1d2
Show file tree
Hide file tree
Showing 11 changed files with 313 additions and 212 deletions.
18 changes: 18 additions & 0 deletions plugin/trino-resource-group-managers/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,12 @@
<scope>runtime</scope>
</dependency>

<dependency>
<groupId>org.postgresql</groupId>
<artifactId>postgresql</artifactId>
<scope>runtime</scope>
</dependency>

<!-- Trino SPI -->
<dependency>
<groupId>io.trino</groupId>
Expand Down Expand Up @@ -175,12 +181,24 @@
<scope>test</scope>
</dependency>

<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>jdbc</artifactId>
<scope>test</scope>
</dependency>

<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>mysql</artifactId>
<scope>test</scope>
</dependency>

<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>postgresql</artifactId>
<scope>test</scope>
</dependency>

<dependency>
<groupId>org.testng</groupId>
<artifactId>testng</artifactId>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
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);
Expand All @@ -25,12 +27,28 @@ private FlywayMigration()
{
}

private static String getLocation(String configDbUrl)
{
if (configDbUrl.startsWith("jdbc:postgresql")) {
return "/db/migration/postgresql";
}
else if (configDbUrl.startsWith("jdbc:oracle")) {
return "/db/migration/oracle";
}
else if (configDbUrl.startsWith("jdbc:mysql")) {
return "/db/migration/mysql";
}
// validation is not performed in DbResourceGroupConfig because DB backed
// resource group tests use the h2 database.
throw new IllegalArgumentException(format("Invalid JDBC URL: %s. Only PostgreSQL, MySQL, and Oracle are supported.", configDbUrl));
}

public static void migrate(DbResourceGroupConfig config)
{
log.info("Performing migrations...");
Flyway flyway = Flyway.configure()
.dataSource(config.getConfigDbUrl(), config.getConfigDbUser(), config.getConfigDbPassword())
.locations("/db/migration/mysql")
.locations(getLocation(config.getConfigDbUrl()))
.baselineOnMigrate(true)
.baselineVersion("0")
.load();
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
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'))
);
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
CREATE TABLE IF NOT EXISTS resource_groups (
resource_group_id BIGSERIAL PRIMARY KEY,
name VARCHAR(250) NOT NULL,
soft_memory_limit VARCHAR(128) NOT NULL,
max_queued INT NOT NULL,
soft_concurrency_limit INT,
hard_concurrency_limit INT NOT NULL,
scheduling_policy VARCHAR(128),
scheduling_weight INT,
jmx_export BOOLEAN,
soft_cpu_limit VARCHAR(128),
hard_cpu_limit VARCHAR(128),
parent BIGINT,
environment VARCHAR(128),
FOREIGN KEY (parent) REFERENCES resource_groups (resource_group_id) ON DELETE CASCADE
);
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
CREATE TABLE IF NOT EXISTS selectors (
resource_group_id BIGINT NOT NULL,
priority BIGINT NOT NULL,
user_regex VARCHAR(512),
source_regex VARCHAR(512),
query_type VARCHAR(512),
client_tags VARCHAR(512),
selector_resource_estimate VARCHAR(1024),
FOREIGN KEY (resource_group_id) REFERENCES resource_groups (resource_group_id) ON DELETE CASCADE
);
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
CREATE TABLE IF NOT EXISTS exact_match_source_selectors(
environment VARCHAR(128),
source VARCHAR(512) NOT NULL,
query_type VARCHAR(512),
update_time TIMESTAMP NOT NULL,
resource_group_id VARCHAR(256) NOT NULL,
PRIMARY KEY (environment, source, resource_group_id),
UNIQUE (source, environment, query_type, resource_group_id)
);
Original file line number Diff line number Diff line change
@@ -0,0 +1,117 @@
/*
* 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.plugin.resourcegroups.db;

import org.jdbi.v3.core.Handle;
import org.jdbi.v3.core.Jdbi;
import org.testcontainers.containers.JdbcDatabaseContainer;
import org.testng.annotations.AfterClass;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;

import java.util.List;

import static org.testng.Assert.assertEquals;

@Test(singleThreaded = true)
public abstract class BaseTestDbResourceGroupsFlywayMigration
{
protected JdbcDatabaseContainer<?> container;
protected Jdbi jdbi;

@BeforeClass
public final void setup()
{
container = startContainer();
jdbi = Jdbi.create(container.getJdbcUrl(), container.getUsername(), container.getPassword());
}

protected abstract JdbcDatabaseContainer<?> startContainer();

@AfterClass(alwaysRun = true)
public final void close()
{
container.close();
}

@AfterMethod(alwaysRun = true)
public void cleanup()
{
dropAllTables();
}

@Test
public void testMigrationWithEmptyDatabase()
{
DbResourceGroupConfig config = new DbResourceGroupConfig()
.setConfigDbUrl(container.getJdbcUrl())
.setConfigDbUser(container.getUsername())
.setConfigDbPassword(container.getPassword());
FlywayMigration.migrate(config);
verifyResourceGroupsSchema(0);
}

@Test
public void testMigrationWithNonEmptyDatabase()
{
String t1Create = "CREATE TABLE t1 (id INT)";
String t2Create = "CREATE TABLE t2 (id INT)";
Handle jdbiHandle = jdbi.open();
jdbiHandle.execute(t1Create);
jdbiHandle.execute(t2Create);
DbResourceGroupConfig config = new DbResourceGroupConfig()
.setConfigDbUrl(container.getJdbcUrl())
.setConfigDbUser(container.getUsername())
.setConfigDbPassword(container.getPassword());
FlywayMigration.migrate(config);
verifyResourceGroupsSchema(0);
String t1Drop = "DROP TABLE t1";
String t2Drop = "DROP TABLE t2";
jdbiHandle.execute(t1Drop);
jdbiHandle.execute(t2Drop);
jdbiHandle.close();
}

protected void verifyResourceGroupsSchema(long expectedPropertiesCount)
{
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, long expectedCount)
{
List<String> results = jdbi.withHandle(handle ->
handle.createQuery(sql).mapTo(String.class).list());
assertEquals(results.size(), expectedCount);
}

protected void dropAllTables()
{
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();
jdbiHandle.execute(propertiesTable);
jdbiHandle.execute(selectorsTable);
jdbiHandle.execute(resourceGroupsTable);
jdbiHandle.execute(exactMatchTable);
jdbiHandle.execute(flywayHistoryTable);
jdbiHandle.close();
}
}
Loading

0 comments on commit 417c1d2

Please sign in to comment.