-
Notifications
You must be signed in to change notification settings - Fork 1.2k
/
lib.rs
125 lines (111 loc) 路 3.93 KB
/
lib.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
use std::io;
use std::time::Duration;
use anyhow::Result;
use futures::{Future, TryFutureExt};
use sqlx::{AnyConnection, Connection};
use crate::opt::{Command, ConnectOpts, DatabaseCommand, MigrateCommand};
mod database;
mod metadata;
// mod migration;
// mod migrator;
#[cfg(feature = "completions")]
mod completions;
mod migrate;
mod opt;
mod prepare;
pub use crate::opt::Opt;
pub async fn run(opt: Opt) -> Result<()> {
match opt.command {
Command::Migrate(migrate) => match migrate.command {
MigrateCommand::Add {
source,
description,
reversible,
sequential,
timestamp,
} => migrate::add(&source, &description, reversible, sequential, timestamp).await?,
MigrateCommand::Run {
source,
dry_run,
ignore_missing,
connect_opts,
} => migrate::run(&source, &connect_opts, dry_run, *ignore_missing).await?,
MigrateCommand::Revert {
source,
dry_run,
ignore_missing,
connect_opts,
} => migrate::revert(&source, &connect_opts, dry_run, *ignore_missing).await?,
MigrateCommand::Info {
source,
connect_opts,
} => migrate::info(&source, &connect_opts).await?,
MigrateCommand::BuildScript { source, force } => migrate::build_script(&source, force)?,
},
Command::Database(database) => match database.command {
DatabaseCommand::Create { connect_opts } => database::create(&connect_opts).await?,
DatabaseCommand::Drop {
confirmation,
connect_opts,
} => database::drop(&connect_opts, !confirmation.yes).await?,
DatabaseCommand::Reset {
confirmation,
source,
connect_opts,
} => database::reset(&source, &connect_opts, !confirmation.yes).await?,
DatabaseCommand::Setup {
source,
connect_opts,
} => database::setup(&source, &connect_opts).await?,
},
Command::Prepare {
check,
workspace,
connect_opts,
args,
} => prepare::run(check, workspace, connect_opts, args).await?,
#[cfg(feature = "completions")]
Command::Completions { shell } => completions::run(shell),
};
Ok(())
}
/// Attempt to connect to the database server, retrying up to `ops.connect_timeout`.
async fn connect(opts: &ConnectOpts) -> sqlx::Result<AnyConnection> {
retry_connect_errors(opts, AnyConnection::connect).await
}
/// Attempt an operation that may return errors like `ConnectionRefused`,
/// retrying up until `ops.connect_timeout`.
///
/// The closure is passed `&ops.database_url` for easy composition.
async fn retry_connect_errors<'a, F, Fut, T>(
opts: &'a ConnectOpts,
mut connect: F,
) -> sqlx::Result<T>
where
F: FnMut(&'a str) -> Fut,
Fut: Future<Output = sqlx::Result<T>> + 'a,
{
sqlx::any::install_default_drivers();
backoff::future::retry(
backoff::ExponentialBackoffBuilder::new()
.with_max_elapsed_time(Some(Duration::from_secs(opts.connect_timeout)))
.build(),
|| {
connect(&opts.database_url).map_err(|e| -> backoff::Error<sqlx::Error> {
match e {
sqlx::Error::Io(ref ioe) => match ioe.kind() {
io::ErrorKind::ConnectionRefused
| io::ErrorKind::ConnectionReset
| io::ErrorKind::ConnectionAborted => {
return backoff::Error::transient(e);
}
_ => (),
},
_ => (),
}
backoff::Error::permanent(e)
})
},
)
.await
}