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
//! Transaction control methods for PostgreSQL connection.
use super::{PgConnection, PgError, PgResult};
/// Quote a SQL identifier (for savepoint names).
/// Wraps in double-quotes and escapes embedded double-quotes.
fn quote_savepoint_name(name: &str) -> PgResult<String> {
if name.is_empty() {
return Err(PgError::Query("savepoint name is empty".to_string()));
}
if name.contains('\0') {
return Err(PgError::Query(
"savepoint name contains NUL byte".to_string(),
));
}
Ok(format!("\"{}\"", name.replace('"', "\"\"")))
}
impl PgConnection {
/// Begin a new transaction.
/// After calling this, all queries run within the transaction
/// until `commit()` or `rollback()` is called.
pub async fn begin_transaction(&mut self) -> PgResult<()> {
self.execute_simple("BEGIN").await
}
/// Commit the current transaction.
/// Makes all changes since `begin_transaction()` permanent.
pub async fn commit(&mut self) -> PgResult<()> {
self.execute_simple("COMMIT").await
}
/// Rollback the current transaction.
/// Discards all changes since `begin_transaction()`.
pub async fn rollback(&mut self) -> PgResult<()> {
self.execute_simple("ROLLBACK").await
}
/// Create a named savepoint within the current transaction.
/// Savepoints allow partial rollback within a transaction.
/// Use `rollback_to()` to return to this savepoint.
pub async fn savepoint(&mut self, name: &str) -> PgResult<()> {
self.execute_simple(&format!("SAVEPOINT {}", quote_savepoint_name(name)?))
.await
}
/// Rollback to a previously created savepoint.
/// Discards all changes since the named savepoint was created,
/// but keeps the transaction open.
pub async fn rollback_to(&mut self, name: &str) -> PgResult<()> {
self.execute_simple(&format!(
"ROLLBACK TO SAVEPOINT {}",
quote_savepoint_name(name)?
))
.await
}
/// Release a savepoint (free resources, if no longer needed).
pub async fn release_savepoint(&mut self, name: &str) -> PgResult<()> {
self.execute_simple(&format!(
"RELEASE SAVEPOINT {}",
quote_savepoint_name(name)?
))
.await
}
}
#[cfg(test)]
mod tests {
use super::quote_savepoint_name;
#[test]
fn quote_savepoint_name_escapes_quotes() {
assert_eq!(quote_savepoint_name("sp\"1").unwrap(), "\"sp\"\"1\"");
}
#[test]
fn quote_savepoint_name_rejects_empty_or_nul() {
assert!(quote_savepoint_name("").is_err());
assert!(quote_savepoint_name("sp\0shadow").is_err());
}
}