1use cufflink_types::manifest::validate::folded_ident;
8use std::time::Duration;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum Scope {
13 Transaction,
15 Session,
17}
18
19impl Scope {
20 pub fn set(self) -> &'static str {
21 match self {
22 Scope::Transaction => "SET LOCAL",
23 Scope::Session => "SET",
24 }
25 }
26}
27
28pub fn quote_ident(name: &str) -> String {
30 format!("\"{}\"", name.replace('"', "\"\""))
31}
32
33pub fn set_role(scope: Scope, role: &str) -> String {
35 format!("{} ROLE {}", scope.set(), quote_ident(role))
36}
37
38pub fn prelude(
44 schema: &str,
45 role: Option<&str>,
46 scope: Scope,
47 statement_timeout: Duration,
48) -> String {
49 let set = scope.set();
50 let settings = format!(
51 "{set} search_path TO {}, public; {set} statement_timeout = {}",
52 folded_ident(schema),
53 statement_timeout.as_millis()
54 );
55 match role {
56 Some(role) => format!("{settings}; {}", set_role(scope, role)),
57 None => settings,
58 }
59}
60
61pub fn reset(role_imposed: bool) -> &'static str {
63 if role_imposed {
64 "RESET search_path; RESET statement_timeout; RESET ROLE"
65 } else {
66 "RESET search_path; RESET statement_timeout"
67 }
68}
69
70pub fn migration_prelude(
74 schema: &str,
75 role: Option<&str>,
76 scope: Scope,
77 statement_timeout: Duration,
78 lock_timeout: Duration,
79) -> String {
80 let set = scope.set();
81 let role = role
82 .map(|role| format!("{};", set_role(scope, role)))
83 .unwrap_or_default();
84 format!(
85 "{set} search_path TO {}, public; {set} statement_timeout = {}; \
86 {set} lock_timeout = {}; {role}",
87 folded_ident(schema),
88 statement_timeout.as_millis(),
89 lock_timeout.as_millis(),
90 )
91}
92
93pub const MIGRATION_RESET: &str =
95 "RESET ROLE; RESET search_path; RESET statement_timeout; RESET lock_timeout";
96
97pub const ROLE_NONE: &str = "SET LOCAL ROLE NONE";
100
101#[cfg(test)]
102mod tests {
103 use super::*;
104
105 const TIMEOUT: Duration = Duration::from_millis(29_500);
106
107 #[test]
108 fn a_guest_prelude_puts_the_schema_first_and_bounds_the_statement() {
109 assert_eq!(
110 prelude("loan_service", None, Scope::Transaction, TIMEOUT),
111 "SET LOCAL search_path TO \"loan_service\", public; SET LOCAL statement_timeout = 29500"
112 );
113 assert_eq!(
114 prelude("loan_service", None, Scope::Session, TIMEOUT),
115 "SET search_path TO \"loan_service\", public; SET statement_timeout = 29500"
116 );
117 for (schema, path) in [("Order_Desk_2", "\"order_desk_2\""), ("order", "\"order\"")] {
118 assert_eq!(
119 prelude(schema, None, Scope::Transaction, TIMEOUT),
120 format!(
121 "SET LOCAL search_path TO {path}, public; SET LOCAL statement_timeout = 29500"
122 ),
123 );
124 }
125 assert_eq!(reset(false), "RESET search_path; RESET statement_timeout");
126 }
127
128 #[test]
129 fn a_role_is_set_last_and_reset_with_the_rest() {
130 let role = Some("svc_acme_loan_service");
131 assert_eq!(
132 prelude("loan_service", role, Scope::Transaction, TIMEOUT),
133 "SET LOCAL search_path TO \"loan_service\", public; \
134 SET LOCAL statement_timeout = 29500; \
135 SET LOCAL ROLE \"svc_acme_loan_service\""
136 );
137 assert_eq!(
138 prelude("loan_service", role, Scope::Session, TIMEOUT),
139 "SET search_path TO \"loan_service\", public; SET statement_timeout = 29500; \
140 SET ROLE \"svc_acme_loan_service\""
141 );
142 assert_eq!(
143 set_role(Scope::Session, "Odd\"Role"),
144 "SET ROLE \"Odd\"\"Role\""
145 );
146 assert_eq!(
147 reset(true),
148 "RESET search_path; RESET statement_timeout; RESET ROLE"
149 );
150 }
151
152 #[test]
153 fn a_migration_prelude_bounds_its_locks_and_names_its_role() {
154 let (statement, lock) = (Duration::from_secs(600), Duration::from_secs(10));
155 assert_eq!(
156 migration_prelude("loan_service", None, Scope::Transaction, statement, lock),
157 "SET LOCAL search_path TO \"loan_service\", public; SET LOCAL statement_timeout = \
158 600000; SET LOCAL lock_timeout = 10000; "
159 );
160 assert_eq!(
161 migration_prelude(
162 "Notes",
163 Some("svc_t_notes"),
164 Scope::Session,
165 statement,
166 lock
167 ),
168 "SET search_path TO \"notes\", public; SET statement_timeout = 600000; \
169 SET lock_timeout = 10000; SET ROLE \"svc_t_notes\";"
170 );
171 }
172
173 #[test]
174 fn identifiers_are_quoted_as_written() {
175 assert_eq!(quote_ident("svc_a"), "\"svc_a\"");
176 assert_eq!(quote_ident("Odd\"Name"), "\"Odd\"\"Name\"");
177 assert_eq!(quote_ident(""), "\"\"");
178 }
179}