Skip to main content

drizzle_core/sql/
comment.rs

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