krabka_trace_context/
sqlcommenter.rs

1//! Reads a W3C trace context out of a [sqlcommenter] tag on a SQL statement.
2//!
3//! OpenTelemetry-instrumented database drivers append
4//! `/*traceparent='00-<32 hex>-<16 hex>-<2 hex>'*/` to the statements they send.
5//! `PostgreSQL` and Krabka's own lexer both skip `--` comments and `/* */`
6//! comments, and emit no token for them. The tag changes no AST, so nothing
7//! rewrites the SQL text before the parser reads it.
8//!
9//! [sqlcommenter]: https://google.github.io/sqlcommenter/
10
11use crate::{
12    carrier::parse_traceparent,
13    propagation::{TRACEPARENT, TRACESTATE},
14};
15
16/// A trace context read out of a sqlcommenter tag, borrowed from the statement.
17///
18/// The `traceparent` has already passed validation against the W3C format.
19/// The fields hold the values exactly as they appeared in the comment. The
20/// sqlcommenter format percent-encodes them. That does nothing to a
21/// `traceparent`, but it can leave a `tracestate` encoded, and
22/// [`crate::TraceCarrier::from_w3c`] drops a `tracestate` it cannot parse.
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct SqlCommenterTrace<'a> {
25    /// The `traceparent` list value, guaranteed to parse.
26    pub traceparent: &'a str,
27    /// The `tracestate` list value, when the comment carried one.
28    pub tracestate: Option<&'a str>,
29}
30
31/// Extract a sqlcommenter trace context from `sql`, if it carries one.
32///
33/// This function examines only genuine comment regions.
34/// `SELECT '/*traceparent=…*/'` is a string literal, not a comment, and gives
35/// `None`. That is the reason this function walks the statement instead of a
36/// match on the text.
37///
38/// The cost when no tag is present is one substring search. The scan starts
39/// only when the word `traceparent` appears somewhere in `sql`.
40#[must_use]
41pub fn extract_sqlcommenter(sql: &str) -> Option<SqlCommenterTrace<'_>> {
42    sql.find(TRACEPARENT)?;
43    scan_comments(sql)
44}
45
46/// Walk `sql` and try to read a trace context out of each comment region.
47///
48/// The walk skips string literals, quoted identifiers, and dollar-quoted
49/// bodies. Any unterminated construct stops the scan, because the statement
50/// will not parse either. A guess at the structure of a truncated string or
51/// comment is exactly what makes a scanner mistake a literal for a comment.
52fn scan_comments(sql: &str) -> Option<SqlCommenterTrace<'_>> {
53    let bytes = sql.as_bytes();
54    let mut index = 0;
55
56    while let Some(&byte) = bytes.get(index) {
57        match byte {
58            b'\'' => index = skip_quoted(bytes, index, b'\'')?,
59            b'"' => index = skip_quoted(bytes, index, b'"')?,
60            b'$' => match dollar_quote_delimiter(sql, index) {
61                // Not a dollar quote — a positional parameter such as `$1`.
62                None => index += 1,
63                Some(delimiter) => {
64                    let body = index + delimiter.len();
65                    index = body + sql.get(body..)?.find(delimiter)? + delimiter.len();
66                }
67            },
68            b'-' if bytes.get(index + 1) == Some(&b'-') => {
69                let start = index + 2;
70                let end = line_comment_end(bytes, start);
71                if let Some(found) = read_comment(sql.get(start..end)?) {
72                    return Some(found);
73                }
74                index = end;
75            }
76            b'/' if bytes.get(index + 1) == Some(&b'*') => {
77                let (end, closed) = block_comment_end(bytes, index);
78                if !closed {
79                    return None;
80                }
81                if let Some(found) = read_comment(sql.get(index + 2..end - 2)?) {
82                    return Some(found);
83                }
84                index = end;
85            }
86            _ => index += 1,
87        }
88    }
89
90    None
91}
92
93/// Read a `traceparent` out of one comment body.
94///
95/// This function also reads the `tracestate` when the body carries one. It
96/// returns `None` when the comment carries no tag, or a tag that fails
97/// validation.
98fn read_comment(content: &str) -> Option<SqlCommenterTrace<'_>> {
99    let traceparent = find_field(content, TRACEPARENT)?;
100    parse_traceparent(traceparent).ok()?;
101    Some(SqlCommenterTrace {
102        traceparent,
103        tracestate: find_field(content, TRACESTATE),
104    })
105}
106
107/// Find `key=value` inside a comment body and return the value.
108///
109/// An identifier byte must not come before the key, so `db_traceparent=…`
110/// cannot pass as `traceparent=…`. The search skips an occurrence that no `=`
111/// follows, and does not stop. That behaviour lets a nested comment such as
112/// `/* /*traceparent='…'*/ */` sit in the same body as the tag. It also lets
113/// the other sqlcommenter keys, such as `route` and `db_driver`, sit there.
114fn find_field<'a>(content: &'a str, key: &str) -> Option<&'a str> {
115    let bytes = content.as_bytes();
116    let mut from = 0;
117
118    while let Some(offset) = content.get(from..).and_then(|rest| rest.find(key)) {
119        let start = from + offset;
120        let end = start + key.len();
121        let preceded_by_identifier = start
122            .checked_sub(1)
123            .and_then(|before| bytes.get(before))
124            .is_some_and(|&byte| is_identifier_byte(byte));
125        if !preceded_by_identifier && let Some(value) = read_value(content, end) {
126            return Some(value);
127        }
128        from = end;
129    }
130
131    None
132}
133
134/// Read the value of a `key=value` pair whose key ends at `index`.
135///
136/// The sqlcommenter specification puts values in single quotes. This function
137/// also accepts an unquoted value, which runs to the next delimiter.
138fn read_value(content: &str, index: usize) -> Option<&str> {
139    let bytes = content.as_bytes();
140    let mut cursor = skip_spaces(bytes, index);
141    if bytes.get(cursor) != Some(&b'=') {
142        return None;
143    }
144    cursor = skip_spaces(bytes, cursor + 1);
145
146    if bytes.get(cursor) == Some(&b'\'') {
147        let start = cursor + 1;
148        let end = start + content.get(start..)?.find('\'')?;
149        return content.get(start..end);
150    }
151
152    let start = cursor;
153    let mut end = cursor;
154    while bytes
155        .get(end)
156        .is_some_and(|&byte| !is_value_terminator(byte))
157    {
158        end += 1;
159    }
160    content.get(start..end).filter(|value| !value.is_empty())
161}
162
163fn skip_spaces(bytes: &[u8], mut index: usize) -> usize {
164    while bytes.get(index).is_some_and(u8::is_ascii_whitespace) {
165        index += 1;
166    }
167    index
168}
169
170fn is_identifier_byte(byte: u8) -> bool {
171    byte.is_ascii_alphanumeric() || byte == b'_'
172}
173
174fn is_value_terminator(byte: u8) -> bool {
175    byte.is_ascii_whitespace() || matches!(byte, b',' | b'*' | b'/' | b'\'' | b'"')
176}
177
178/// Skip a `'…'` string literal or a `"…"` quoted identifier at `index`.
179///
180/// The skip honours the doubled-quote escape. Returns `None` when the literal
181/// or the identifier is unterminated, which stops the scan: the text that
182/// remains is not statement structure.
183fn skip_quoted(bytes: &[u8], mut index: usize, quote: u8) -> Option<usize> {
184    index += 1;
185    while let Some(&byte) = bytes.get(index) {
186        if byte != quote {
187            index += 1;
188        } else if bytes.get(index + 1) == Some(&quote) {
189            index += 2;
190        } else {
191            return Some(index + 1);
192        }
193    }
194    None
195}
196
197/// The `$tag$` delimiter that opens at `index`.
198///
199/// Returns `None` when `index` does not begin a dollar-quoted body. A
200/// positional parameter such as `$1` is one such case, because its tag would
201/// have to start with a digit.
202fn dollar_quote_delimiter(sql: &str, index: usize) -> Option<&str> {
203    let bytes = sql.as_bytes();
204    let mut end = index + 1;
205    if bytes.get(end).is_some_and(u8::is_ascii_digit) {
206        return None;
207    }
208    while bytes.get(end).is_some_and(|&byte| is_identifier_byte(byte)) {
209        end += 1;
210    }
211    if bytes.get(end) == Some(&b'$') {
212        sql.get(index..=end)
213    } else {
214        None
215    }
216}
217
218fn line_comment_end(bytes: &[u8], mut index: usize) -> usize {
219    while bytes.get(index).is_some_and(|&byte| byte != b'\n') {
220        index += 1;
221    }
222    index
223}
224
225/// Find the end of the block comment that opens at `index`.
226///
227/// The search honours `PostgreSQL`'s comment nesting, which matches
228/// `krabka_pgparser`'s lexer. Returns the offset just past the comment, and
229/// whether the comment was closed.
230fn block_comment_end(bytes: &[u8], index: usize) -> (usize, bool) {
231    let mut depth = 1usize;
232    let mut cursor = index + 2;
233
234    while let Some(&byte) = bytes.get(cursor) {
235        if byte == b'/' && bytes.get(cursor + 1) == Some(&b'*') {
236            depth += 1;
237            cursor += 2;
238        } else if byte == b'*' && bytes.get(cursor + 1) == Some(&b'/') {
239            depth -= 1;
240            cursor += 2;
241            if depth == 0 {
242                return (cursor, true);
243            }
244        } else {
245            cursor += 1;
246        }
247    }
248
249    (cursor, false)
250}
krabka_trace_context/sqlcommenter.rs

krabka_trace_context/
sqlcommenter.rs
