Skip to main content

connectorx/
utils.rs

1use std::ops::{Deref, DerefMut};
2use url::{form_urlencoded, Url};
3
4/// Remove the given parameters from a URL query, leaving every other parameter
5/// byte-for-byte as the caller wrote it.
6///
7/// `Url::query_pairs_mut` cannot be used for this: it serializes the query as
8/// `application/x-www-form-urlencoded`, which encodes a space as `+`. Drivers that
9/// percent-decode the query (rust-postgres, for one) then read that `+` literally,
10/// so `?options=-c statement_timeout=1s` reaches the server as `-c+statement_timeout=1s`.
11pub(crate) fn remove_query_params(url: &Url, params: &[&str]) -> Url {
12    let mut stripped = url.clone();
13    match url.query() {
14        None => stripped,
15        Some(query) => {
16            let kept: Vec<&str> = query
17                .split('&')
18                .filter(|segment| !segment.is_empty())
19                .filter(|segment| {
20                    let raw_key = segment.split('=').next().unwrap_or_default();
21                    // compare decoded keys, as `query_pairs` would
22                    let key = form_urlencoded::parse(raw_key.as_bytes())
23                        .next()
24                        .map(|(key, _)| key)
25                        .unwrap_or_default();
26                    !params.contains(&&*key)
27                })
28                .collect();
29
30            let query = kept.join("&");
31            stripped.set_query(match query.is_empty() {
32                true => None,
33                false => Some(&query),
34            });
35            stripped
36        }
37    }
38}
39
40pub struct DummyBox<T>(pub T);
41
42impl<T> Deref for DummyBox<T> {
43    type Target = T;
44
45    fn deref(&self) -> &Self::Target {
46        &self.0
47    }
48}
49
50impl<T> DerefMut for DummyBox<T> {
51    fn deref_mut(&mut self) -> &mut Self::Target {
52        &mut self.0
53    }
54}
55
56#[cfg(feature = "dst_arrow")]
57pub fn decimal_to_i128(mut v: rust_decimal::Decimal, scale: u32) -> anyhow::Result<i128> {
58    v.rescale(scale);
59
60    let v_scale = v.scale();
61    if v_scale != scale as u32 {
62        return Err(anyhow::anyhow!(
63            "decimal scale is not equal to expected scale, got: {} expected: {}",
64            v_scale,
65            scale
66        ));
67    }
68
69    Ok(v.mantissa())
70}
71
72#[cfg(test)]
73mod tests {
74    use super::remove_query_params;
75    use url::Url;
76
77    fn strip(uri: &str, params: &[&str]) -> String {
78        remove_query_params(&Url::parse(uri).unwrap(), params).to_string()
79    }
80
81    #[test]
82    fn preserves_spaces_in_values() {
83        assert_eq!(
84            strip(
85                "postgresql://u:p@host/db?options=-c%20statement_timeout%3D1s&cxprotocol=binary",
86                &["cxprotocol"]
87            ),
88            "postgresql://u:p@host/db?options=-c%20statement_timeout%3D1s"
89        );
90    }
91
92    #[test]
93    fn does_not_reencode_remaining_params() {
94        assert_eq!(
95            strip("mysql://host/db?a=x+y&b=%2Fz&flag", &["nothing"]),
96            "mysql://host/db?a=x+y&b=%2Fz&flag"
97        );
98    }
99
100    #[test]
101    fn removes_every_requested_param() {
102        assert_eq!(
103            strip(
104                "postgresql://host/db?sslcert=a&keep=1&sslkey=b&sslrootcert=c",
105                &["sslcert", "sslkey", "sslrootcert"]
106            ),
107            "postgresql://host/db?keep=1"
108        );
109    }
110
111    #[test]
112    fn drops_the_query_when_nothing_is_left() {
113        assert_eq!(
114            strip("postgresql://host/db?cxprotocol=binary", &["cxprotocol"]),
115            "postgresql://host/db"
116        );
117    }
118
119    #[test]
120    fn handles_a_missing_query() {
121        assert_eq!(
122            strip("postgresql://host/db", &["cxprotocol"]),
123            "postgresql://host/db"
124        );
125    }
126
127    #[test]
128    fn matches_percent_encoded_keys() {
129        assert_eq!(
130            strip(
131                "postgresql://host/db?cx%70rotocol=binary&a=1",
132                &["cxprotocol"]
133            ),
134            "postgresql://host/db?a=1"
135        );
136    }
137
138    #[test]
139    fn keeps_duplicate_keys_that_are_not_removed() {
140        assert_eq!(
141            strip("mysql://host/db?a=1&a=2&cxprotocol=text", &["cxprotocol"]),
142            "mysql://host/db?a=1&a=2"
143        );
144    }
145}