tank_core/
util.rs

1use crate::{AsValue, ColumnDef, DynQuery, TableRef, Value};
2use proc_macro2::TokenStream;
3use quote::{ToTokens, TokenStreamExt, quote};
4use rust_decimal::prelude::ToPrimitive;
5use serde_json::{Map, Number, Value as JsonValue};
6use std::{
7    borrow::Cow,
8    cmp::min,
9    collections::BTreeMap,
10    ffi::{CStr, CString},
11    ptr,
12};
13use syn::Path;
14
15#[derive(Clone)]
16/// Iterator adapter for two types.
17pub enum EitherIterator<A, B>
18where
19    A: Iterator,
20    B: Iterator<Item = A::Item>,
21{
22    Left(A),
23    Right(B),
24}
25
26impl<A, B> Iterator for EitherIterator<A, B>
27where
28    A: Iterator,
29    B: Iterator<Item = A::Item>,
30{
31    type Item = A::Item;
32    fn next(&mut self) -> Option<Self::Item> {
33        match self {
34            Self::Left(a) => a.next(),
35            Self::Right(b) => b.next(),
36        }
37    }
38}
39
40pub fn value_to_json(v: &Value) -> Option<JsonValue> {
41    Some(match v {
42        _ if v.is_null() => JsonValue::Null,
43        Value::Boolean(Some(v), ..) => JsonValue::Bool(*v),
44        Value::Int8(Some(v), ..) => JsonValue::Number(Number::from_i128(*v as _)?),
45        Value::Int16(Some(v), ..) => JsonValue::Number(Number::from_i128(*v as _)?),
46        Value::Int32(Some(v), ..) => JsonValue::Number(Number::from_i128(*v as _)?),
47        Value::Int64(Some(v), ..) => JsonValue::Number(Number::from_i128(*v as _)?),
48        Value::Int128(Some(v), ..) => JsonValue::Number(Number::from_i128(*v as _)?),
49        Value::UInt8(Some(v), ..) => JsonValue::Number(Number::from_u128(*v as _)?),
50        Value::UInt16(Some(v), ..) => JsonValue::Number(Number::from_u128(*v as _)?),
51        Value::UInt32(Some(v), ..) => JsonValue::Number(Number::from_u128(*v as _)?),
52        Value::UInt64(Some(v), ..) => JsonValue::Number(Number::from_u128(*v as _)?),
53        Value::UInt128(Some(v), ..) => JsonValue::Number(Number::from_u128(*v as _)?),
54        Value::Float32(Some(v), ..) => JsonValue::Number(Number::from_f64(*v as _)?),
55        Value::Float64(Some(v), ..) => JsonValue::Number(Number::from_f64(*v as _)?),
56        Value::Decimal(Some(v), ..) => JsonValue::Number(Number::from_f64(v.to_f64()?)?),
57        Value::Char(Some(v), ..) => JsonValue::String(v.to_string()),
58        Value::Varchar(Some(v), ..) => JsonValue::String(v.to_string()),
59        Value::Blob(Some(v), ..) => JsonValue::Array(
60            v.iter()
61                .map(|v| Number::from_u128(*v as _).map(JsonValue::Number))
62                .collect::<Option<_>>()?,
63        ),
64        v @ (Value::Date(Some(..), ..)
65        | Value::Time(Some(..), ..)
66        | Value::Timestamp(Some(..), ..)
67        | Value::TimestampWithTimezone(Some(..), ..)) => JsonValue::String(v.to_string()),
68        Value::Interval(Some(_v), ..) => {
69            return None;
70        }
71        Value::Uuid(Some(v), ..) => JsonValue::String(v.to_string()),
72        Value::Array(Some(v), ..) => {
73            JsonValue::Array(v.iter().map(value_to_json).collect::<Option<_>>()?)
74        }
75        Value::List(Some(v), ..) => {
76            JsonValue::Array(v.iter().map(value_to_json).collect::<Option<_>>()?)
77        }
78        Value::Map(Some(v), ..) => {
79            let mut map = Map::new();
80            for (k, v) in v.iter() {
81                let Ok(k) = String::try_from_value(k.clone()) else {
82                    return None;
83                };
84                let Some(v) = value_to_json(v) else {
85                    return None;
86                };
87                map.insert(k, v)?;
88            }
89            JsonValue::Object(map)
90        }
91        Value::Json(Some(v), ..) => v.clone(),
92        Value::Struct(Some(v), ..) => {
93            let mut map = Map::new();
94            for (k, v) in v.iter() {
95                let Some(v) = value_to_json(v) else {
96                    return None;
97                };
98                map.insert(k.clone(), v)?;
99            }
100            JsonValue::Object(map)
101        }
102        Value::Unknown(Some(v), ..) => JsonValue::String(v.clone()),
103        _ => {
104            return None;
105        }
106    })
107}
108
109/// Quote a `BTreeMap<K, V>` into tokens.
110pub fn quote_btree_map<K: ToTokens, V: ToTokens>(value: &BTreeMap<K, V>) -> TokenStream {
111    let mut tokens = TokenStream::new();
112    for (k, v) in value {
113        let ks = k.to_token_stream();
114        let vs = v.to_token_stream();
115        tokens.append_all(quote! {
116            (#ks, #vs),
117        });
118    }
119    quote! {
120        ::std::collections::BTreeMap::from([
121            #tokens
122        ])
123    }
124}
125
126/// Quote a `Cow<T>` preserving borrowed vs owned status for generated code.
127pub fn quote_cow<T: ToOwned + ToTokens + ?Sized>(value: &Cow<T>) -> TokenStream
128where
129    <T as ToOwned>::Owned: ToTokens,
130{
131    match value {
132        Cow::Borrowed(v) => quote! { ::std::borrow::Cow::Borrowed(#v) },
133        Cow::Owned(v) => quote! { ::std::borrow::Cow::Borrowed(#v) },
134    }
135}
136
137/// Quote an `Option<T>` into tokens.
138pub fn quote_option<T: ToTokens>(value: &Option<T>) -> TokenStream {
139    match value {
140        None => quote! { None },
141        Some(v) => quote! { Some(#v) },
142    }
143}
144
145/// Determine if the trailing segments of a `syn::Path` match the expected identifiers.
146pub fn matches_path(path: &Path, expect: &[&str]) -> bool {
147    let len = min(path.segments.len(), expect.len());
148    path.segments
149        .iter()
150        .rev()
151        .take(len)
152        .map(|v| &v.ident)
153        .eq(expect.iter().rev().take(len))
154}
155
156/// Write an iterator of items separated by a delimiter into a string.
157pub fn separated_by<T, F>(
158    out: &mut DynQuery,
159    values: impl IntoIterator<Item = T>,
160    mut f: F,
161    separator: &str,
162) where
163    F: FnMut(&mut DynQuery, T),
164{
165    let mut len = out.len();
166    for v in values {
167        if out.len() > len {
168            out.push_str(separator);
169        }
170        len = out.len();
171        f(out, v);
172    }
173}
174
175/// Write, escaping occurrences of `search` char with `replace` while copying into buffer.
176pub fn write_escaped(out: &mut DynQuery, value: &str, search: char, replace: &str) {
177    let mut position = 0;
178    for (i, c) in value.char_indices() {
179        if c == search {
180            out.push_str(&value[position..i]);
181            out.push_str(replace);
182            position = i + 1;
183        }
184    }
185    out.push_str(&value[position..]);
186}
187
188/// Convenience wrapper converting into a `CString`.
189pub fn as_c_string(str: impl Into<Vec<u8>>) -> CString {
190    CString::new(
191        str.into()
192            .into_iter()
193            .map(|b| if b == 0 { b'?' } else { b })
194            .collect::<Vec<u8>>(),
195    )
196    .unwrap_or_default()
197}
198
199pub fn error_message_from_ptr<'a>(ptr: &'a *const i8) -> Cow<'a, str> {
200    unsafe {
201        if *ptr != ptr::null() {
202            CStr::from_ptr(*ptr).to_string_lossy()
203        } else {
204            Cow::Borrowed("Unknown error: could not extract the error message")
205        }
206    }
207}
208
209/// Consume a prefix of `input` while the predicate returns true, returning that slice.
210pub fn consume_while<'s>(input: &mut &'s str, predicate: impl FnMut(&char) -> bool) -> &'s str {
211    let len = input.chars().take_while(predicate).count();
212    if len == 0 {
213        return "";
214    }
215    let result = &input[..len];
216    *input = &input[len..];
217    result
218}
219
220pub fn extract_number<'s, const SIGNED: bool>(input: &mut &'s str) -> &'s str {
221    let mut end = 0;
222    let mut chars = input.chars().peekable();
223    if SIGNED && matches!(chars.peek(), Some('+') | Some('-')) {
224        chars.next();
225        end += 1;
226    }
227    for _ in chars.take_while(char::is_ascii_digit) {
228        end += 1;
229    }
230    let result = &input[..end];
231    *input = &input[end..];
232    result
233}
234
235pub fn column_def(name: &str, table: &TableRef) -> Option<&'static ColumnDef> {
236    table.columns.iter().find(|v| v.name() == name)
237}
238
239#[macro_export]
240macro_rules! number_to_month {
241    ($month:expr, $throw:expr $(,)?) => {
242        match $month {
243            1 => Month::January,
244            2 => Month::February,
245            3 => Month::March,
246            4 => Month::April,
247            5 => Month::May,
248            6 => Month::June,
249            7 => Month::July,
250            8 => Month::August,
251            9 => Month::September,
252            10 => Month::October,
253            11 => Month::November,
254            12 => Month::December,
255            _ => $throw,
256        }
257    };
258}
259
260#[macro_export]
261macro_rules! month_to_number {
262    ($month:expr $(,)?) => {
263        match $month {
264            Month::January => 1,
265            Month::February => 2,
266            Month::March => 3,
267            Month::April => 4,
268            Month::May => 5,
269            Month::June => 6,
270            Month::July => 7,
271            Month::August => 8,
272            Month::September => 9,
273            Month::October => 10,
274            Month::November => 11,
275            Month::December => 12,
276        }
277    };
278}
279
280#[macro_export]
281/// Conditionally wrap a generated fragment in parentheses.
282macro_rules! possibly_parenthesized {
283    ($out:ident, $cond:expr, $v:expr) => {
284        if $cond {
285            $out.push('(');
286            $v;
287            $out.push(')');
288        } else {
289            $v;
290        }
291    };
292}
293
294pub const TRUNCATE_LONG_LIMIT: usize = {
295    #[cfg(debug_assertions)]
296    {
297        2000
298    }
299
300    #[cfg(not(debug_assertions))]
301    {
302        497
303    }
304};
305
306#[macro_export]
307/// Truncate long strings for logging and error messages purpose.
308///
309/// Returns a `format_args!` that yields at most 497 characters from the start
310/// of the input followed by `...` when truncation occurred. Minimal overhead.
311///
312/// If true is the second argument, it evaluates the first argument just once.
313///
314/// # Examples
315/// ```ignore
316/// use tank_core::truncate_long;
317/// let short = "SELECT 1";
318/// assert_eq!(format!("{}", truncate_long!(short)), "SELECT 1\n");
319/// let long = format!("SELECT {}", "X".repeat(600));
320/// let logged = format!("{}", truncate_long!(long));
321/// assert!(logged.starts_with("SELECT XXXXXX"));
322/// assert!(logged.ends_with("...\n"));
323/// ```
324macro_rules! truncate_long {
325    ($query:expr) => {
326        format_args!(
327            "{}{}",
328            &$query[..::std::cmp::min($query.len(), $crate::TRUNCATE_LONG_LIMIT)].trim(),
329            if $query.len() > $crate::TRUNCATE_LONG_LIMIT {
330                "...\n"
331            } else {
332                ""
333            },
334        )
335    };
336    ($query:expr,true) => {{
337        let query = $query;
338        format!(
339            "{}{}",
340            &query[..::std::cmp::min(query.len(), $crate::TRUNCATE_LONG_LIMIT)].trim(),
341            if query.len() > $crate::TRUNCATE_LONG_LIMIT {
342                "...\n"
343            } else {
344                ""
345            },
346        )
347    }};
348}
349
350/// Sends the value through the channel and logs in case of error.
351///
352/// Parameters:
353/// * `$tx`: sender channel
354/// * `$value`: value to be sent
355///
356/// *Example*:
357/// ```ignore
358/// send_value!(tx, Ok(QueryResult::Row(row)));
359/// ```
360
361#[macro_export]
362macro_rules! send_value {
363    ($tx:ident, $value:expr) => {{
364        if let Err(e) = $tx.send($value) {
365            log::error!("{e:#}");
366        }
367    }};
368}
369
370/// Incrementally accumulates tokens from a speculative parse stream until one
371/// of the supplied parsers succeeds.
372///
373/// Returns `(accumulated_tokens, (parser1_option, parser2_option, ...))` with
374/// exactly one `Some(T)`: the first successful parser.
375#[doc(hidden)]
376#[macro_export]
377macro_rules! take_until {
378    ($original:expr, $($parser:expr),+ $(,)?) => {{
379        let macro_local_input = $original.fork();
380        let mut macro_local_result = (
381            TokenStream::new(),
382            ($({
383                let _ = $parser;
384                None
385            }),+),
386        );
387        loop {
388            if macro_local_input.is_empty() {
389                break;
390            }
391            let mut parsed = false;
392            let produced = ($({
393                let attempt = macro_local_input.fork();
394                if let Ok(content) = ($parser)(&attempt) {
395                    macro_local_input.advance_to(&attempt);
396                    parsed = true;
397                    Some(content)
398                } else {
399                    None
400                }
401            }),+);
402            if parsed {
403                macro_local_result.1 = produced;
404                break;
405            }
406            macro_local_result.0.append(macro_local_input.parse::<TokenTree>()?);
407        }
408        $original.advance_to(&macro_local_input);
409        macro_local_result
410    }};
411}
412
413#[macro_export]
414/// Implement the `Executor` trait for a transaction wrapper type by
415/// delegating each operation to an underlying connection object.
416///
417/// This reduces boilerplate across driver implementations. The macro expands
418/// into an `impl Executor for $transaction<'c>` with forwarding methods for
419/// `prepare`, `run`, `fetch`, `execute`, and `append`.
420///
421/// Parameters:
422/// * `$driver`: concrete driver type.
423/// * `$transaction`: transaction wrapper type (generic over lifetime `'c`).
424/// * `$connection`: field name on the transaction pointing to the connection.
425///
426/// # Examples
427/// ```ignore
428/// use crate::{YourDBConnection, YourDBDriver};
429/// use tank_core::{Error, Result, Transaction, impl_executor_transaction};
430///
431/// pub struct YourDBTransaction<'c> {
432///     connection: &'c mut YourDBConnection,
433/// }
434///
435/// impl_executor_transaction!(YourDBDriver, YourDBTransaction<'c>, connection);
436///
437/// impl<'c> Transaction<'c> for YourDBTransaction<'c> { ... }
438/// ```
439macro_rules! impl_executor_transaction {
440    // Case 1: Lifetime is present (necessary for transactions)
441    ($driver:ty, $transaction:ident $(< $lt:lifetime >)?, $connection:ident) => {
442       impl $(<$lt>)? ::tank_core::Executor for $transaction $(<$lt>)? {
443            type Driver = $driver;
444
445            fn accepts_multiple_statements(&self) -> bool {
446                self.$connection.accepts_multiple_statements()
447            }
448
449            fn do_prepare(
450                &mut self,
451                sql: String,
452            ) -> impl Future<Output = ::tank_core::Result<::tank_core::Query<Self::Driver>>> + Send
453            {
454                self.$connection.do_prepare(sql)
455            }
456
457            fn run<'s>(
458                &'s mut self,
459                query: impl ::tank_core::AsQuery<Self::Driver> + 's,
460            ) -> impl ::tank_core::stream::Stream<
461                Item = ::tank_core::Result<::tank_core::QueryResult>,
462            > + Send {
463                self.$connection.run(query)
464            }
465
466            fn fetch<'s>(
467                &'s mut self,
468                query: impl ::tank_core::AsQuery<Self::Driver> + 's,
469            ) -> impl ::tank_core::stream::Stream<
470                Item = ::tank_core::Result<::tank_core::Row>,
471            > + Send
472            + 's {
473                self.$connection.fetch(query)
474            }
475
476            fn execute<'s>(
477                &'s mut self,
478                query: impl ::tank_core::AsQuery<Self::Driver> + 's,
479            ) -> impl Future<Output = ::tank_core::Result<::tank_core::RowsAffected>> + Send {
480                self.$connection.execute(query)
481            }
482
483            fn append<'a, E, It>(
484                &mut self,
485                entities: It,
486            ) -> impl Future<Output = ::tank_core::Result<::tank_core::RowsAffected>> + Send
487            where
488                E: ::tank_core::Entity + 'a,
489                It: IntoIterator<Item = &'a E> + Send,
490                <It as IntoIterator>::IntoIter: Send,
491            {
492                self.$connection.append(entities)
493            }
494        }
495    }
496}
497
498#[macro_export]
499macro_rules! current_timestamp_ms {
500    () => {{}};
501}
tank_core/util.rs

tank_core/
util.rs
