Skip to content

Commit 57087cb

Browse files
tarkahtarkah
committed
add migrate for aurora
1 parent fa7f2af commit 57087cb

File tree

2 files changed

+364
-1
lines changed

2 files changed

+364
-1
lines changed

‎sqlx-core/src/aurora/migrate.rs

Lines changed: 362 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,362 @@
1+
use crate::aurora::{Aurora, AuroraConnectOptions, AuroraConnection, AuroraDbType};
2+
use crate::connection::{ConnectOptions, Connection};
3+
use crate::error::Error;
4+
use crate::executor::Executor;
5+
use crate::migrate::MigrateError;
6+
use crate::migrate::Migration;
7+
use crate::migrate::{Migrate, MigrateDatabase};
8+
use crate::query::query;
9+
use crate::query_as::query_as;
10+
use crate::query_scalar::query_scalar;
11+
use crc::crc32;
12+
use futures_core::future::BoxFuture;
13+
use std::str::FromStr;
14+
use std::time::Duration;
15+
use std::time::Instant;
16+
17+
fn parse_for_maintenance(uri: &str) -> Result<(AuroraConnectOptions, String, AuroraDbType), Error> {
18+
let mut options = AuroraConnectOptions::from_str(uri)?;
19+
20+
let db_type = if let Some(db_type) = options.db_type {
21+
db_type
22+
} else {
23+
return Err(Error::Configuration(
24+
"DATABASE_URL does not specify a db type".into(),
25+
));
26+
};
27+
28+
let database = if let Some(database) = &options.database {
29+
database.to_owned()
30+
} else {
31+
return Err(Error::Configuration(
32+
"DATABASE_URL does not specify a database".into(),
33+
));
34+
};
35+
36+
match db_type {
37+
AuroraDbType::MySQL => options.database = None,
38+
AuroraDbType::Postgres => options.database = Some("postgres".into()),
39+
}
40+
41+
Ok((options, database, db_type))
42+
}
43+
44+
impl MigrateDatabase for Aurora {
45+
fn create_database(uri: &str) -> BoxFuture<'_, Result<(), Error>> {
46+
Box::pin(async move {
47+
let (options, database, db_type) = parse_for_maintenance(uri)?;
48+
let mut conn = options.connect().await?;
49+
50+
let sql = match db_type {
51+
AuroraDbType::MySQL => format!("CREATE DATABASE `{}`", database),
52+
AuroraDbType::Postgres => {
53+
format!("CREATE DATABASE \"{}\"", database.replace('"', "\"\""))
54+
}
55+
};
56+
57+
let _ = conn.execute(&*sql).await?;
58+
59+
Ok(())
60+
})
61+
}
62+
63+
fn database_exists(uri: &str) -> BoxFuture<'_, Result<bool, Error>> {
64+
Box::pin(async move {
65+
let (options, database, db_type) = parse_for_maintenance(uri)?;
66+
let mut conn = options.connect().await?;
67+
68+
let sql = match db_type {
69+
AuroraDbType::MySQL => {
70+
"select exists(SELECT 1 from INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = ?)"
71+
}
72+
AuroraDbType::Postgres => {
73+
"select exists(SELECT 1 from pg_database WHERE datname = $1)"
74+
}
75+
};
76+
77+
let exists: bool = query_scalar(sql)
78+
.bind(database)
79+
.fetch_one(&mut conn)
80+
.await?;
81+
82+
Ok(exists)
83+
})
84+
}
85+
86+
fn drop_database(uri: &str) -> BoxFuture<'_, Result<(), Error>> {
87+
Box::pin(async move {
88+
let (options, database, db_type) = parse_for_maintenance(uri)?;
89+
let mut conn = options.connect().await?;
90+
91+
let sql = match db_type {
92+
AuroraDbType::MySQL => {
93+
format!("DROP DATABASE IF EXISTS `{}`", database,)
94+
}
95+
AuroraDbType::Postgres => {
96+
format!(
97+
"DROP DATABASE IF EXISTS \"{}\"",
98+
database.replace('"', "\"\"")
99+
)
100+
}
101+
};
102+
103+
let _ = conn.execute(&*sql).await?;
104+
105+
Ok(())
106+
})
107+
}
108+
}
109+
110+
impl Migrate for AuroraConnection {
111+
fn ensure_migrations_table(&mut self) -> BoxFuture<'_, Result<(), MigrateError>> {
112+
Box::pin(async move {
113+
let sql = match self.db_type {
114+
AuroraDbType::MySQL => {
115+
r#"
116+
CREATE TABLE IF NOT EXISTS _sqlx_migrations (
117+
version BIGINT PRIMARY KEY,
118+
description TEXT NOT NULL,
119+
installed_on TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
120+
success BOOLEAN NOT NULL,
121+
checksum BLOB NOT NULL,
122+
execution_time BIGINT NOT NULL
123+
);
124+
"#
125+
}
126+
AuroraDbType::Postgres => {
127+
r#"
128+
CREATE TABLE IF NOT EXISTS _sqlx_migrations (
129+
version BIGINT PRIMARY KEY,
130+
description TEXT NOT NULL,
131+
installed_on TIMESTAMPTZ NOT NULL DEFAULT now(),
132+
success BOOLEAN NOT NULL,
133+
checksum BYTEA NOT NULL,
134+
execution_time BIGINT NOT NULL
135+
);
136+
"#
137+
}
138+
};
139+
140+
// language=SQL
141+
self.execute(sql).await?;
142+
143+
Ok(())
144+
})
145+
}
146+
147+
fn version(&mut self) -> BoxFuture<'_, Result<Option<(i64, bool)>, MigrateError>> {
148+
Box::pin(async move {
149+
// language=SQL
150+
let row = query_as(
151+
"SELECT version, NOT success FROM _sqlx_migrations ORDER BY version DESC LIMIT 1",
152+
)
153+
.fetch_optional(self)
154+
.await?;
155+
156+
Ok(row)
157+
})
158+
}
159+
160+
fn lock(&mut self) -> BoxFuture<'_, Result<(), MigrateError>> {
161+
Box::pin(async move {
162+
let database_name = current_database(self).await?;
163+
let lock_id = generate_lock_id(&database_name);
164+
165+
let sql = match self.db_type {
166+
AuroraDbType::MySQL => "SELECT GET_LOCK(?, -1)",
167+
AuroraDbType::Postgres => "SELECT pg_advisory_lock($1)",
168+
};
169+
170+
// language=SQL
171+
let query = query(sql);
172+
173+
match self.db_type {
174+
AuroraDbType::MySQL => {
175+
query.bind(format!("{:x}", lock_id)).execute(self).await?;
176+
}
177+
AuroraDbType::Postgres => {
178+
query.bind(lock_id).execute(self).await?;
179+
}
180+
};
181+
182+
Ok(())
183+
})
184+
}
185+
186+
fn unlock(&mut self) -> BoxFuture<'_, Result<(), MigrateError>> {
187+
Box::pin(async move {
188+
let database_name = current_database(self).await?;
189+
let lock_id = generate_lock_id(&database_name);
190+
191+
let sql = match self.db_type {
192+
AuroraDbType::MySQL => "SELECT RELEASE_LOCK(?)",
193+
AuroraDbType::Postgres => "SELECT pg_advisory_unlock($1)",
194+
};
195+
196+
// language=SQL
197+
let query = query(sql);
198+
199+
match self.db_type {
200+
AuroraDbType::MySQL => {
201+
query.bind(format!("{:x}", lock_id)).execute(self).await?;
202+
}
203+
AuroraDbType::Postgres => {
204+
query.bind(lock_id).execute(self).await?;
205+
}
206+
};
207+
208+
Ok(())
209+
})
210+
}
211+
212+
fn validate<'e: 'm, 'm>(
213+
&'e mut self,
214+
migration: &'m Migration,
215+
) -> BoxFuture<'m, Result<(), MigrateError>> {
216+
Box::pin(async move {
217+
let sql = match self.db_type {
218+
AuroraDbType::MySQL => "SELECT checksum FROM _sqlx_migrations WHERE version = ?",
219+
AuroraDbType::Postgres => {
220+
"SELECT checksum FROM _sqlx_migrations WHERE version = $1"
221+
}
222+
};
223+
224+
// language=SQL
225+
let checksum: Option<Vec<u8>> = query_scalar(sql)
226+
.bind(migration.version)
227+
.fetch_optional(self)
228+
.await?;
229+
230+
if let Some(checksum) = checksum {
231+
if checksum == &*migration.checksum {
232+
Ok(())
233+
} else {
234+
Err(MigrateError::VersionMismatch(migration.version))
235+
}
236+
} else {
237+
Err(MigrateError::VersionMissing(migration.version))
238+
}
239+
})
240+
}
241+
242+
fn apply<'e: 'm, 'm>(
243+
&'e mut self,
244+
migration: &'m Migration,
245+
) -> BoxFuture<'m, Result<Duration, MigrateError>> {
246+
Box::pin(async move {
247+
match self.db_type {
248+
AuroraDbType::MySQL => {
249+
let start = Instant::now();
250+
251+
let res = self.execute(&*migration.sql).await;
252+
253+
let elapsed = start.elapsed();
254+
255+
// language=MySQL
256+
let _ = query(
257+
r#"
258+
INSERT INTO _sqlx_migrations ( version, description, success, checksum, execution_time )
259+
VALUES ( ?, ?, ?, ?, ? )
260+
"#,
261+
)
262+
.bind(migration.version)
263+
.bind(&*migration.description)
264+
.bind(res.is_ok())
265+
.bind(&*migration.checksum)
266+
.bind(elapsed.as_nanos() as i64)
267+
.execute(self)
268+
.await?;
269+
270+
res?;
271+
272+
Ok(elapsed)
273+
}
274+
AuroraDbType::Postgres => {
275+
let mut tx = self.begin().await?;
276+
let start = Instant::now();
277+
278+
let _ = tx.execute(&*migration.sql).await?;
279+
280+
tx.commit().await?;
281+
282+
let elapsed = start.elapsed();
283+
284+
// language=SQL
285+
let _ = query(
286+
r#"
287+
INSERT INTO _sqlx_migrations ( version, description, success, checksum, execution_time )
288+
VALUES ( $1, $2, TRUE, $3, $4 )
289+
"#,
290+
)
291+
.bind(migration.version)
292+
.bind(&*migration.description)
293+
.bind(&*migration.checksum)
294+
.bind(elapsed.as_nanos() as i64)
295+
.execute(self)
296+
.await?;
297+
298+
Ok(elapsed)
299+
}
300+
}
301+
})
302+
}
303+
304+
fn revert<'e: 'm, 'm>(
305+
&'e mut self,
306+
migration: &'m Migration,
307+
) -> BoxFuture<'m, Result<Duration, MigrateError>> {
308+
Box::pin(async move {
309+
match self.db_type {
310+
AuroraDbType::MySQL => {
311+
let start = Instant::now();
312+
313+
self.execute(&*migration.sql).await?;
314+
315+
let elapsed = start.elapsed();
316+
317+
// language=SQL
318+
let _ = query(r#"DELETE FROM _sqlx_migrations WHERE version = ?"#)
319+
.bind(migration.version)
320+
.execute(self)
321+
.await?;
322+
323+
Ok(elapsed)
324+
}
325+
AuroraDbType::Postgres => {
326+
let mut tx = self.begin().await?;
327+
let start = Instant::now();
328+
329+
let _ = tx.execute(&*migration.sql).await?;
330+
331+
tx.commit().await?;
332+
333+
let elapsed = start.elapsed();
334+
335+
// language=SQL
336+
let _ = query(r#"DELETE FROM _sqlx_migrations WHERE version = $1"#)
337+
.bind(migration.version)
338+
.execute(self)
339+
.await?;
340+
341+
Ok(elapsed)
342+
}
343+
}
344+
})
345+
}
346+
}
347+
348+
async fn current_database(conn: &mut AuroraConnection) -> Result<String, MigrateError> {
349+
let sql = match conn.db_type {
350+
AuroraDbType::MySQL => "SELECT DATABASE()",
351+
AuroraDbType::Postgres => "SELECT current_database()",
352+
};
353+
354+
// language=SQL
355+
Ok(query_scalar(sql).fetch_one(conn).await?)
356+
}
357+
358+
// inspired from rails: https://github.com/rails/rails/blob/6e49cc77ab3d16c06e12f93158eaf3e507d4120e/activerecord/lib/active_record/migration.rb#L1308
359+
fn generate_lock_id(database_name: &str) -> i64 {
360+
// 0x3d32ad9e chosen by fair dice roll
361+
0x3d32ad9e * (crc32::checksum_ieee(database_name.as_bytes()) as i64)
362+
}

‎sqlx-core/src/aurora/mod.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ mod connection;
44
mod database;
55
mod done;
66
mod error;
7+
mod migrate;
78
mod options;
89
mod row;
910
mod statement;
@@ -18,7 +19,7 @@ pub use connection::AuroraConnection;
1819
pub use database::Aurora;
1920
pub use done::AuroraDone;
2021
pub use error::AuroraDatabaseError;
21-
pub use options::AuroraConnectOptions;
22+
pub use options::{AuroraConnectOptions, AuroraDbType};
2223
pub use row::AuroraRow;
2324
pub use statement::AuroraStatement;
2425
pub use transaction::AuroraTransactionManager;

0 commit comments

Comments
 (0)