-
Notifications
You must be signed in to change notification settings - Fork 3.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add PostgreSQL migration files for DB resource groups
- Loading branch information
Showing
11 changed files
with
313 additions
and
212 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
5 changes: 5 additions & 0 deletions
5
.../src/main/resources/db/migration/postgresql/V1__add_resource_groups_global_properties.sql
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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')) | ||
); |
16 changes: 16 additions & 0 deletions
16
...rce-group-managers/src/main/resources/db/migration/postgresql/V2__add_resource_groups.sql
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | ||
); |
10 changes: 10 additions & 0 deletions
10
...-resource-group-managers/src/main/resources/db/migration/postgresql/V3__add_selectors.sql
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | ||
); |
9 changes: 9 additions & 0 deletions
9
...agers/src/main/resources/db/migration/postgresql/V4__add_exact_match_source_selectors.sql
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) | ||
); |
117 changes: 117 additions & 0 deletions
117
.../test/java/io/trino/plugin/resourcegroups/db/BaseTestDbResourceGroupsFlywayMigration.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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(); | ||
} | ||
} |
Oops, something went wrong.