Skip to content

Commit 531929a

Browse files
posullivposulliv
authored
Use flyway to perform database migrations on startup
1 parent 1fab8ad commit 531929a

19 files changed

+663
-14
lines changed

‎docker/docker-compose.yml‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,3 @@ services:
3434
timeout: 1s
3535
retries: 60
3636
start_period: 10s
37-
volumes:
38-
- target: /docker-entrypoint-initdb.d/1-gateway-ha-persistence-postgres.sql
39-
source: ../gateway-ha/src/main/resources/gateway-ha-persistence-postgres.sql
40-
type: bind

‎docs/installation.md‎

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -38,16 +38,37 @@ distribution is installed.
3838

3939
### Backend database
4040

41-
Trino Gateway requires a MySQL or PostgreSQL database.
41+
Trino Gateway requires a MySQL or PostgreSQL database. Database initialization
42+
is performed automatically when the Trino Gateway process starts. Migrations
43+
are performed using `Flyway`.
4244

43-
Use the following scripts in the `gateway-ha/src/main/resources/` folder to
44-
initialize the database:
45-
46-
* `gateway-ha-persistence-mysql.sql` for MySQL
47-
* `gateway-ha-persistence-postgres.sql` for PostgreSQL
45+
The migration files can viewed in the `gateway-ha/src/main/resources/` folder.
46+
Each database type supported has its own sub-folder.
4847

4948
The files are also included in the JAR file.
5049

50+
If you do not want migrations to be performed automatically on startup, then
51+
you can set `runMigrationsEnabled` to `false` in the data store configuration.
52+
For example:
53+
54+
```yaml
55+
dataStore:
56+
jdbcUrl: jdbc:postgresql://postgres:5432/trino_gateway_db
57+
user: USER
58+
password: PASSWORD
59+
driver: org.postgresql.Driver
60+
queryHistoryHoursRetention: 24
61+
runMigrationsEnabled: false
62+
```
63+
64+
`Flyway` uses a transactional lock in databases that support it such as
65+
[PostgreSQL](https://documentation.red-gate.com/fd/postgresql-database-235241807.html#).
66+
In the scenario where multiple Trino Gateway instances are running and sharing
67+
the same backend database, the first Trino Gateway instance to start will get
68+
the lock and run the database migrations with `Flyway`. Other Trino Gateway
69+
instances might fail during startup while migrations are running but once migrations
70+
are completed they will start as expected.
71+
5172
### Trino clusters
5273

5374
The proxied Trino clusters behind the Trino Gateway must support the Trino JDBC

‎gateway-ha/pom.xml‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
<frontend.pnpmRegistryURL>https://registry.npmmirror.com</frontend.pnpmRegistryURL>
2222

2323
<!-- dependency versions -->
24+
<dep.flyway.version>11.0.1</dep.flyway.version>
2425
<dep.jeasy.version>4.1.0</dep.jeasy.version>
2526
<dep.mockito.version>5.14.2</dep.mockito.version>
2627
<dep.okhttp3.version>4.12.0</dep.okhttp3.version>
@@ -237,6 +238,12 @@
237238
<version>1.78.1</version>
238239
</dependency>
239240

241+
<dependency>
242+
<groupId>org.flywaydb</groupId>
243+
<artifactId>flyway-core</artifactId>
244+
<version>${dep.flyway.version}</version>
245+
</dependency>
246+
240247
<dependency>
241248
<groupId>org.glassfish.jersey.core</groupId>
242249
<artifactId>jersey-server</artifactId>
@@ -290,6 +297,20 @@
290297
<scope>runtime</scope>
291298
</dependency>
292299

300+
<dependency>
301+
<groupId>org.flywaydb</groupId>
302+
<artifactId>flyway-database-postgresql</artifactId>
303+
<version>${dep.flyway.version}</version>
304+
<scope>runtime</scope>
305+
</dependency>
306+
307+
<dependency>
308+
<groupId>org.flywaydb</groupId>
309+
<artifactId>flyway-mysql</artifactId>
310+
<version>${dep.flyway.version}</version>
311+
<scope>runtime</scope>
312+
</dependency>
313+
293314
<dependency>
294315
<groupId>org.mvel</groupId>
295316
<artifactId>mvel2</artifactId>
@@ -371,6 +392,12 @@
371392
<scope>test</scope>
372393
</dependency>
373394

395+
<dependency>
396+
<groupId>org.testcontainers</groupId>
397+
<artifactId>jdbc</artifactId>
398+
<scope>test</scope>
399+
</dependency>
400+
374401
<dependency>
375402
<groupId>org.testcontainers</groupId>
376403
<artifactId>mysql</artifactId>

‎gateway-ha/src/main/java/io/trino/gateway/ha/HaGatewayLauncher.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import io.airlift.units.Duration;
3333
import io.trino.gateway.baseapp.BaseApp;
3434
import io.trino.gateway.ha.config.HaGatewayConfiguration;
35+
import io.trino.gateway.ha.persistence.FlywayMigration;
3536
import org.weakref.jmx.guice.MBeanModule;
3637

3738
import java.nio.file.Files;
@@ -109,6 +110,7 @@ public static void main(String[] args)
109110
}
110111
String config = Files.readString(Path.of(args[0]));
111112
HaGatewayConfiguration haGatewayConfiguration = objectMapper.readValue(replaceEnvironmentVariables(config), HaGatewayConfiguration.class);
113+
FlywayMigration.migrate(haGatewayConfiguration.getDataStore());
112114
List<Module> modules = addModules(haGatewayConfiguration);
113115
new HaGatewayLauncher().start(modules, haGatewayConfiguration);
114116
}

‎gateway-ha/src/main/java/io/trino/gateway/ha/config/DataStoreConfiguration.java‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,16 @@ public class DataStoreConfiguration
2020
private String password;
2121
private String driver;
2222
private Integer queryHistoryHoursRetention = 4;
23+
private boolean runMigrationsEnabled = true;
2324

24-
public DataStoreConfiguration(String jdbcUrl, String user, String password, String driver, Integer queryHistoryHoursRetention)
25+
public DataStoreConfiguration(String jdbcUrl, String user, String password, String driver, Integer queryHistoryHoursRetention, boolean runMigrationsEnabled)
2526
{
2627
this.jdbcUrl = jdbcUrl;
2728
this.user = user;
2829
this.password = password;
2930
this.driver = driver;
3031
this.queryHistoryHoursRetention = queryHistoryHoursRetention;
32+
this.runMigrationsEnabled = runMigrationsEnabled;
3133
}
3234

3335
public DataStoreConfiguration() {}
@@ -81,4 +83,14 @@ public void setQueryHistoryHoursRetention(Integer queryHistoryHoursRetention)
8183
{
8284
this.queryHistoryHoursRetention = queryHistoryHoursRetention;
8385
}
86+
87+
public boolean isRunMigrationsEnabled()
88+
{
89+
return this.runMigrationsEnabled;
90+
}
91+
92+
public void setRunMigrationsEnabled(boolean runMigrationsEnabled)
93+
{
94+
this.runMigrationsEnabled = runMigrationsEnabled;
95+
}
8496
}

‎gateway-ha/src/main/java/io/trino/gateway/ha/persistence/FlywayMigration.java‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package io.trino.gateway.ha.persistence;
15+
16+
import io.airlift.log.Logger;
17+
import io.trino.gateway.ha.config.DataStoreConfiguration;
18+
import org.flywaydb.core.Flyway;
19+
import org.flywaydb.core.api.output.MigrateResult;
20+
21+
import static java.lang.String.format;
22+
23+
public class FlywayMigration
24+
{
25+
private static final Logger log = Logger.get(FlywayMigration.class);
26+
27+
private FlywayMigration() {}
28+
29+
private static String getLocation(String configDbUrl)
30+
{
31+
if (configDbUrl.startsWith("jdbc:postgresql")) {
32+
return "postgresql";
33+
}
34+
if (configDbUrl.startsWith("jdbc:mysql")) {
35+
return "mysql";
36+
}
37+
throw new IllegalArgumentException(format("Invalid JDBC URL: %s. Only PostgreSQL and MySQL are supported.", configDbUrl));
38+
}
39+
40+
public static void migrate(DataStoreConfiguration config)
41+
{
42+
if (!config.isRunMigrationsEnabled()) {
43+
log.info("Skip migrations as automatic migrations are disabled");
44+
return;
45+
}
46+
log.info("Performing migrations...");
47+
Flyway flyway = Flyway.configure()
48+
.dataSource(config.getJdbcUrl(), config.getUser(), config.getPassword())
49+
.locations(getLocation(config.getJdbcUrl()))
50+
.baselineOnMigrate(true)
51+
.baselineVersion("0")
52+
.load();
53+
54+
MigrateResult migrations = flyway.migrate();
55+
log.info("Performed %s migrations", migrations.migrationsExecuted);
56+
}
57+
}

‎gateway-ha/src/main/resources/mysql/V1__create_schema.sql‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
CREATE TABLE IF NOT EXISTS gateway_backend (
2+
name VARCHAR(256) PRIMARY KEY,
3+
routing_group VARCHAR (256),
4+
backend_url VARCHAR (256),
5+
external_url VARCHAR (256),
6+
active BOOLEAN
7+
);
8+
9+
CREATE TABLE IF NOT EXISTS query_history (
10+
query_id VARCHAR(256) PRIMARY KEY,
11+
query_text VARCHAR (256),
12+
created bigint,
13+
backend_url VARCHAR (256),
14+
user_name VARCHAR(256),
15+
source VARCHAR(256)
16+
);
17+
CREATE INDEX query_history_created_idx ON query_history(created);
18+
19+
CREATE TABLE IF NOT EXISTS resource_groups (
20+
resource_group_id BIGINT NOT NULL AUTO_INCREMENT,
21+
name VARCHAR(250) NOT NULL UNIQUE,
22+
23+
-- OPTIONAL POLICY CONTROLS
24+
parent BIGINT NULL,
25+
jmx_export BOOLEAN NULL,
26+
scheduling_policy VARCHAR(128) NULL,
27+
scheduling_weight INT NULL,
28+
29+
-- REQUIRED QUOTAS
30+
soft_memory_limit VARCHAR(128) NOT NULL,
31+
max_queued INT NOT NULL,
32+
hard_concurrency_limit INT NOT NULL,
33+
34+
-- OPTIONAL QUOTAS
35+
soft_concurrency_limit INT NULL,
36+
soft_cpu_limit VARCHAR(128) NULL,
37+
hard_cpu_limit VARCHAR(128) NULL,
38+
environment VARCHAR(128) NULL,
39+
40+
PRIMARY KEY(resource_group_id),
41+
FOREIGN KEY (parent) REFERENCES resource_groups (resource_group_id)
42+
);
43+
44+
CREATE TABLE IF NOT EXISTS selectors (
45+
resource_group_id BIGINT NOT NULL,
46+
priority BIGINT NOT NULL,
47+
48+
-- Regex fields -- these will be used as a regular expression pattern to
49+
-- match against the field of the same name on queries
50+
user_regex VARCHAR(512),
51+
source_regex VARCHAR(512),
52+
53+
-- Selector fields -- these must match exactly.
54+
query_type VARCHAR(512),
55+
client_tags VARCHAR(512),
56+
selector_resource_estimate VARCHAR(1024),
57+
58+
FOREIGN KEY (resource_group_id) REFERENCES resource_groups(resource_group_id)
59+
);
60+
61+
CREATE TABLE IF NOT EXISTS resource_groups_global_properties (
62+
name VARCHAR(128) NOT NULL PRIMARY KEY,
63+
value VARCHAR(512) NULL,
64+
CHECK (name in ('cpu_quota_period'))
65+
);
66+
67+
CREATE TABLE IF NOT EXISTS exact_match_source_selectors (
68+
resource_group_id VARCHAR(256) NOT NULL,
69+
update_time DATETIME NOT NULL,
70+
71+
-- Selector fields which must exactly match a query
72+
source VARCHAR(512) NOT NULL,
73+
environment VARCHAR(128),
74+
query_type VARCHAR(512),
75+
76+
PRIMARY KEY (environment, source(128), query_type),
77+
UNIQUE (source(128), environment, query_type(128), resource_group_id)
78+
);

‎gateway-ha/src/main/resources/postgresql/V1__create_schema.sql‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
CREATE TABLE IF NOT EXISTS gateway_backend (
2+
name VARCHAR(256) PRIMARY KEY,
3+
routing_group VARCHAR (256),
4+
backend_url VARCHAR (256),
5+
external_url VARCHAR (256),
6+
active BOOLEAN
7+
);
8+
9+
CREATE TABLE IF NOT EXISTS query_history (
10+
query_id VARCHAR(256) PRIMARY KEY,
11+
query_text VARCHAR (256),
12+
created bigint,
13+
backend_url VARCHAR (256),
14+
user_name VARCHAR(256),
15+
source VARCHAR(256)
16+
);
17+
CREATE INDEX IF NOT EXISTS query_history_created_idx ON query_history(created);
18+
19+
CREATE TABLE IF NOT EXISTS resource_groups (
20+
resource_group_id SERIAL,
21+
name VARCHAR(250) NOT NULL UNIQUE,
22+
23+
-- OPTIONAL POLICY CONTROLS
24+
parent BIGINT NULL,
25+
jmx_export BOOLEAN NULL,
26+
scheduling_policy VARCHAR(128) NULL,
27+
scheduling_weight INT NULL,
28+
29+
-- REQUIRED QUOTAS
30+
soft_memory_limit VARCHAR(128) NOT NULL,
31+
max_queued INT NOT NULL,
32+
hard_concurrency_limit INT NOT NULL,
33+
34+
-- OPTIONAL QUOTAS
35+
soft_concurrency_limit INT NULL,
36+
soft_cpu_limit VARCHAR(128) NULL,
37+
hard_cpu_limit VARCHAR(128) NULL,
38+
environment VARCHAR(128) NULL,
39+
40+
PRIMARY KEY(resource_group_id),
41+
FOREIGN KEY (parent) REFERENCES resource_groups (resource_group_id)
42+
);
43+
44+
CREATE TABLE IF NOT EXISTS selectors (
45+
resource_group_id BIGINT NOT NULL,
46+
priority BIGINT NOT NULL,
47+
48+
-- Regex fields -- these will be used as a regular expression pattern to
49+
-- match against the field of the same name on queries
50+
user_regex VARCHAR(512),
51+
source_regex VARCHAR(512),
52+
53+
-- Selector fields -- these must match exactly.
54+
query_type VARCHAR(512),
55+
client_tags VARCHAR(512),
56+
selector_resource_estimate VARCHAR(1024),
57+
58+
FOREIGN KEY (resource_group_id) REFERENCES resource_groups(resource_group_id)
59+
);
60+
61+
CREATE TABLE IF NOT EXISTS resource_groups_global_properties (
62+
name VARCHAR(128) NOT NULL PRIMARY KEY,
63+
value VARCHAR(512) NULL,
64+
CHECK (name in ('cpu_quota_period'))
65+
);
66+
67+
CREATE TABLE IF NOT EXISTS exact_match_source_selectors (
68+
resource_group_id VARCHAR(256) NOT NULL,
69+
update_time TIMESTAMP NOT NULL,
70+
71+
-- Selector fields which must exactly match a query
72+
source VARCHAR(512) NOT NULL,
73+
environment VARCHAR(128),
74+
query_type VARCHAR(128), -- (reduced from 512)
75+
76+
PRIMARY KEY (environment, source, query_type),
77+
UNIQUE (source, environment, query_type, resource_group_id)
78+
);

‎gateway-ha/src/test/java/io/trino/gateway/ha/TestTrinoResource.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ void setup()
6161

6262
// Setup resource group manager
6363
String jdbcUrl = "jdbc:h2:" + testConfig.h2DbFilePath();
64-
DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", "sa", "org.h2.Driver", 4);
64+
DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", "sa", "org.h2.Driver", 4, false);
6565
Jdbi jdbi = Jdbi.create(jdbcUrl, "sa", "sa");
6666
connectionManager = new JdbcConnectionManager(jdbi, db);
6767
resourceGroupManager = new HaResourceGroupsManager(connectionManager);

‎gateway-ha/src/test/java/io/trino/gateway/ha/TestingJdbcConnectionManager.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ public static JdbcConnectionManager createTestingJdbcConnectionManager()
3030
tempH2DbDir.deleteOnExit();
3131
String jdbcUrl = "jdbc:h2:" + tempH2DbDir.getAbsolutePath();
3232
HaGatewayTestUtils.seedRequiredData(new HaGatewayTestUtils.TestConfig("", tempH2DbDir.getAbsolutePath()));
33-
DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", "sa", "org.h2.Driver", 4);
33+
DataStoreConfiguration db = new DataStoreConfiguration(jdbcUrl, "sa", "sa", "org.h2.Driver", 4, false);
3434
Jdbi jdbi = Jdbi.create(jdbcUrl, "sa", "sa");
3535
return new JdbcConnectionManager(jdbi, db);
3636
}

0 commit comments

Comments
 (0)