Skip to main content

cufflink_pgcodec/
session.rs

1//! The settings statements run under on a session: the service's schema
2//! first on the search path, the statement's time bound, and the role it
3//! runs as. Whatever holds the session sends these texts, so the
4//! platform's sessions, its migration executor and the native test host
5//! set up a session alike.
6
7use cufflink_types::manifest::validate::folded_ident;
8use std::time::Duration;
9
10/// Where a session's settings apply.
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum Scope {
13    /// `SET LOCAL`: until the transaction ends.
14    Transaction,
15    /// `SET`: until the session's reset runs.
16    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
28/// `name` as a quoted SQL identifier, exactly as written.
29pub fn quote_ident(name: &str) -> String {
30    format!("\"{}\"", name.replace('"', "\"\""))
31}
32
33/// `SET LOCAL ROLE` or `SET ROLE`, as `scope` says.
34pub fn set_role(scope: Scope, role: &str) -> String {
35    format!("{} ROLE {}", scope.set(), quote_ident(role))
36}
37
38/// What a guest's statements run under: the schema is written as every
39/// statement writes it ([`folded_ident`]), so the path a guest's SQL
40/// resolves in, the schema a deploy created and the one a migration ran in
41/// are the same schema. `role` is the one a session imposes on its
42/// connection, where it does.
43pub 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
61/// Undoes a [`Scope::Session`] [`prelude`].
62pub 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
70/// What a service's SQL migrations run under: its schema, the statement
71/// and lock time bounds, and the role the migrations run as where the run
72/// imposes one.
73pub 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
93/// Undoes a [`Scope::Session`] [`migration_prelude`].
94pub const MIGRATION_RESET: &str =
95    "RESET ROLE; RESET search_path; RESET statement_timeout; RESET lock_timeout";
96
97/// Ends the role a migration ran as within its transaction, so the
98/// history row is written as the session's own role.
99pub 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("note_service", None, Scope::Transaction, TIMEOUT),
111            "SET LOCAL search_path TO \"note_service\", public; SET LOCAL statement_timeout = 29500"
112        );
113        assert_eq!(
114            prelude("note_service", None, Scope::Session, TIMEOUT),
115            "SET search_path TO \"note_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_note_service");
131        assert_eq!(
132            prelude("note_service", role, Scope::Transaction, TIMEOUT),
133            "SET LOCAL search_path TO \"note_service\", public; \
134             SET LOCAL statement_timeout = 29500; \
135             SET LOCAL ROLE \"svc_acme_note_service\""
136        );
137        assert_eq!(
138            prelude("note_service", role, Scope::Session, TIMEOUT),
139            "SET search_path TO \"note_service\", public; SET statement_timeout = 29500; \
140             SET ROLE \"svc_acme_note_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("note_service", None, Scope::Transaction, statement, lock),
157            "SET LOCAL search_path TO \"note_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}