Skip to main content

drizzle_core/sql/
comment.rs

1//! sqlcommenter helpers: attach trace or context metadata to queries as SQL
2//! comments.
3//!
4//! Mirrors upstream `drizzle-orm`'s `sql.comment()` / `sqlCommenter()` helpers
5//! so that observability sidecars that parse SQL comments (Google Cloud SQL
6//! Insights, Sqlcommenter, etc.) see the same on-the-wire format.
7//!
8//! - [`comment`] wraps a free-form string in `/* ... */`, sanitising any
9//!   embedded `/*` or `*/` sequences so they can't terminate the comment.
10//! - [`comment_tags`] URL-encodes each key/value (matching JS
11//!   `encodeURIComponent`, plus a `'` → `\'` escape), sorts alphabetically,
12//!   joins with `,`, and wraps in `/* ... */`.
13//!
14//! Both helpers return an empty [`SQL`] fragment when the input reduces to
15//! nothing (empty string, or a map whose values are all empty/skipped).
16
17use crate::prelude::*;
18use crate::sql::SQL;
19use crate::traits::SQLParam;
20use percent_encoding::{AsciiSet, CONTROLS, utf8_percent_encode};
21
22/// Byte set to percent-encode for `encodeURIComponent` semantics.
23///
24/// `encodeURIComponent` preserves `A-Z a-z 0-9 - _ . ! ~ * ' ( )` and
25/// percent-encodes every other byte of the UTF-8 encoding. `percent-encoding`
26/// works inversely — you enumerate bytes *to* encode — so we start from
27/// `CONTROLS` (0x00-0x1F + 0x7F) and add every printable ASCII byte that is
28/// *not* in the unreserved set.
29const COMPONENT: &AsciiSet = &CONTROLS
30    // space + ASCII printables that are not in the unreserved set.
31    .add(b' ')
32    .add(b'"')
33    .add(b'#')
34    .add(b'$')
35    .add(b'%')
36    .add(b'&')
37    .add(b'+')
38    .add(b',')
39    .add(b'/')
40    .add(b':')
41    .add(b';')
42    .add(b'<')
43    .add(b'=')
44    .add(b'>')
45    .add(b'?')
46    .add(b'@')
47    .add(b'[')
48    .add(b'\\')
49    .add(b']')
50    .add(b'^')
51    .add(b'`')
52    .add(b'{')
53    .add(b'|')
54    .add(b'}');
55
56/// Creates a free-form SQL comment: `/*text*/`.
57///
58/// The text cannot end the comment early: `/*` becomes `/ *` and `*/`
59/// becomes `* /`. An empty input gives an empty fragment.
60///
61/// Most code calls the `.comment(...)` method on a dialect's query builder
62/// instead.
63///
64/// # Examples
65///
66/// ```
67/// use drizzle_core::{SQL, sql::comment};
68/// # use drizzle_core::{Dialect, SQLParam, SQLiteDialect};
69/// # use std::borrow::Cow;
70/// # #[derive(Debug, Clone, PartialEq)]
71/// # struct Value(i64);
72/// # impl SQLParam for Value {
73/// #     const DIALECT: Dialect = Dialect::SQLite;
74/// #     type DialectMarker = SQLiteDialect;
75/// # }
76/// # impl From<Value> for Cow<'_, Value> {
77/// #     fn from(value: Value) -> Self { Cow::Owned(value) }
78/// # }
79///
80/// let sql: SQL<'_, Value> = comment("route=/users */ DROP");
81/// assert_eq!(sql.sql(), "/*route=/users * / DROP*/");
82/// ```
83pub fn comment<'a, V: SQLParam>(text: impl AsRef<str>) -> SQL<'a, V> {
84    let text = text.as_ref();
85    if text.is_empty() {
86        return SQL::empty();
87    }
88    let sanitized = sanitize_string_input(text);
89    let mut out = String::with_capacity(sanitized.len() + 4);
90    out.push_str("/*");
91    out.push_str(&sanitized);
92    out.push_str("*/");
93    SQL::raw(out)
94}
95
96/// Creates a sqlcommenter tag comment: `/*key='value',...*/`.
97///
98/// Keys and values are URL-encoded like JavaScript's `encodeURIComponent`,
99/// with `'` also escaped as `\'`. Pairs are sorted by their encoded form
100/// and joined with `,`. Pairs with an empty value are skipped. If no pair
101/// is left, the result is an empty fragment.
102///
103/// # Examples
104///
105/// ```
106/// use drizzle_core::{SQL, sql::comment_tags};
107/// # use drizzle_core::{Dialect, SQLParam, SQLiteDialect};
108/// # use std::borrow::Cow;
109/// # #[derive(Debug, Clone, PartialEq)]
110/// # struct Value(i64);
111/// # impl SQLParam for Value {
112/// #     const DIALECT: Dialect = Dialect::SQLite;
113/// #     type DialectMarker = SQLiteDialect;
114/// # }
115/// # impl From<Value> for Cow<'_, Value> {
116/// #     fn from(value: Value) -> Self { Cow::Owned(value) }
117/// # }
118///
119/// let sql: SQL<'_, Value> = comment_tags([("route", "/users"), ("action", "list")]);
120/// assert_eq!(sql.sql(), "/*action='list',route='%2Fusers'*/");
121/// ```
122pub fn comment_tags<'a, V, I, K, Val>(pairs: I) -> SQL<'a, V>
123where
124    V: SQLParam,
125    I: IntoIterator<Item = (K, Val)>,
126    K: AsRef<str>,
127    Val: AsRef<str>,
128{
129    let mut parts: Vec<String> = Vec::new();
130    for (k, v) in pairs {
131        let v = v.as_ref();
132        if v.is_empty() {
133            continue;
134        }
135        let ek = sanitize_object_element(k.as_ref());
136        let ev = sanitize_object_element(v);
137        let mut entry = String::with_capacity(ek.len() + ev.len() + 3);
138        entry.push_str(&ek);
139        entry.push_str("='");
140        entry.push_str(&ev);
141        entry.push('\'');
142        parts.push(entry);
143    }
144    if parts.is_empty() {
145        return SQL::empty();
146    }
147    parts.sort();
148
149    let total: usize = parts.iter().map(String::len).sum::<usize>() + parts.len() + 3;
150    let mut out = String::with_capacity(total);
151    out.push_str("/*");
152    for (i, p) in parts.iter().enumerate() {
153        if i > 0 {
154            out.push(',');
155        }
156        out.push_str(p);
157    }
158    out.push_str("*/");
159    SQL::raw(out)
160}
161
162/// Sanitise a free-form comment string so it can't terminate the enclosing
163/// `/* ... */` block. Replaces `/*` with `/ *` and `*/` with `* /`, in that
164/// order — matching upstream JS behaviour byte-for-byte.
165#[inline]
166fn sanitize_string_input(input: &str) -> String {
167    input.replace("/*", "/ *").replace("*/", "* /")
168}
169
170/// Sanitise a key or value for a tag-style comment. URL-encodes using
171/// `encodeURIComponent` semantics via [`percent_encoding::utf8_percent_encode`]
172/// against [`COMPONENT`], then escapes any remaining `'` as `\'` (since `'` is
173/// in the unreserved set that `encodeURIComponent` preserves but it would
174/// clash with the surrounding `'...'` wrapping).
175#[inline]
176fn sanitize_object_element(s: &str) -> String {
177    let encoded = utf8_percent_encode(s, COMPONENT).to_string();
178    if encoded.as_bytes().contains(&b'\'') {
179        encoded.replace('\'', "\\'")
180    } else {
181        encoded
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188    use crate::dialect::{Dialect, SQLiteDialect};
189
190    // A minimal SQLParam we can use in tests without pulling in a driver dep.
191    #[derive(Clone, Debug)]
192    struct TestValue;
193
194    impl SQLParam for TestValue {
195        const DIALECT: Dialect = Dialect::SQLite;
196        type DialectMarker = SQLiteDialect;
197    }
198
199    fn render(s: SQL<'_, TestValue>) -> String {
200        s.sql()
201    }
202
203    #[test]
204    fn empty_string_yields_empty_sql() {
205        let s: SQL<'_, TestValue> = comment("");
206        assert_eq!(render(s), "");
207    }
208
209    #[test]
210    fn plain_string_is_wrapped() {
211        let s: SQL<'_, TestValue> = comment("hello world");
212        assert_eq!(render(s), "/*hello world*/");
213    }
214
215    #[test]
216    fn string_input_sanitises_comment_terminators() {
217        let s: SQL<'_, TestValue> = comment("/* nested */ end");
218        assert_eq!(render(s), "/*/ * nested * / end*/");
219    }
220
221    #[test]
222    fn tags_are_sorted_and_url_encoded() {
223        let s: SQL<'_, TestValue> = comment_tags([("route", "/users/:id"), ("action", "update")]);
224        assert_eq!(render(s), "/*action='update',route='%2Fusers%2F%3Aid'*/");
225    }
226
227    #[test]
228    fn empty_values_are_skipped() {
229        let s: SQL<'_, TestValue> = comment_tags([("a", ""), ("b", "ok")]);
230        assert_eq!(render(s), "/*b='ok'*/");
231    }
232
233    #[test]
234    fn all_empty_yields_empty_sql() {
235        let s: SQL<'_, TestValue> = comment_tags([("a", ""), ("b", "")]);
236        assert_eq!(render(s), "");
237    }
238
239    #[test]
240    fn quote_in_value_is_escaped_after_url_encoding() {
241        let s: SQL<'_, TestValue> = comment_tags([("k", "it's")]);
242        assert_eq!(render(s), r"/*k='it\'s'*/");
243    }
244
245    #[test]
246    fn multibyte_utf8_is_percent_encoded_per_byte() {
247        // "é" is U+00E9 = C3 A9 in UTF-8.
248        let s: SQL<'_, TestValue> = comment_tags([("name", "café")]);
249        assert_eq!(render(s), "/*name='caf%C3%A9'*/");
250    }
251
252    #[test]
253    fn unreserved_set_is_preserved() {
254        // ALPHA/DIGIT plus -_.!~*'()
255        let s: SQL<'_, TestValue> = comment_tags([("k", "abcXYZ012-_.!~*()")]);
256        // The single-quote is not in the input; this just confirms unreserved
257        // bytes are left alone.
258        assert_eq!(render(s), "/*k='abcXYZ012-_.!~*()'*/");
259    }
260}