connectorx/sources/trino/
mod.rs

1use std::{collections::HashMap, marker::PhantomData, sync::Arc};
2
3use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
4use fehler::{throw, throws};
5use prusto::{auth::Auth, Client, ClientBuilder, DataSet, Presto, Row};
6use serde_json::Value;
7use sqlparser::dialect::{Dialect, GenericDialect};
8use std::convert::TryFrom;
9use tokio::runtime::Runtime;
10
11use crate::{
12    data_order::DataOrder,
13    errors::ConnectorXError,
14    sources::Produce,
15    sql::{count_query, limit0_query, CXQuery},
16};
17
18pub use self::{errors::TrinoSourceError, typesystem::TrinoTypeSystem};
19use urlencoding::decode;
20
21use super::{PartitionParser, Source, SourcePartition};
22
23use anyhow::anyhow;
24
25pub mod errors;
26pub mod typesystem;
27
28#[throws(TrinoSourceError)]
29fn get_total_rows(rt: Arc<Runtime>, client: Arc<Client>, query: &CXQuery<String>) -> usize {
30    let cquery = count_query(query, &TrinoDialect {})?;
31
32    let row = rt
33        .block_on(client.get_all::<Row>(cquery.to_string()))
34        .map_err(TrinoSourceError::PrustoError)?
35        .split()
36        .1[0]
37        .clone();
38
39    let value = row
40        .value()
41        .first()
42        .ok_or_else(|| anyhow!("Trino count dataset is empty"))?;
43
44    value
45        .as_i64()
46        .ok_or_else(|| anyhow!("Trino cannot parse i64"))? as usize
47}
48
49#[derive(Presto, Debug)]
50pub struct TrinoPartitionQueryResult {
51    pub _col0: i64,
52    pub _col1: i64,
53}
54
55#[derive(Debug)]
56pub struct TrinoDialect {}
57
58// implementation copy from AnsiDialect
59impl Dialect for TrinoDialect {
60    fn is_identifier_start(&self, ch: char) -> bool {
61        ch.is_ascii_lowercase() || ch.is_ascii_uppercase()
62    }
63
64    fn is_identifier_part(&self, ch: char) -> bool {
65        ch.is_ascii_lowercase() || ch.is_ascii_uppercase() || ch.is_ascii_digit() || ch == '_'
66    }
67}
68
69pub struct TrinoSource {
70    client: Arc<Client>,
71    rt: Arc<Runtime>,
72    origin_query: Option<String>,
73    queries: Vec<CXQuery<String>>,
74    names: Vec<String>,
75    schema: Vec<TrinoTypeSystem>,
76}
77
78/// Extract the credentials from a Trino connection URL.
79///
80/// `Url::username` and `Url::password` hand back the components still
81/// percent-encoded, so they have to be decoded before they go into an
82/// `Authorization` header: a password of `p@ss` reaches the server as
83/// `p%40ss` otherwise, and Trino answers `401 Invalid credentials`.
84#[throws(TrinoSourceError)]
85fn credentials_from_url(url: &url::Url) -> (String, Option<String>) {
86    let username = match url.username() {
87        "" => "connectorx".to_owned(),
88        username => decode(username)?.into_owned(),
89    };
90
91    let password = match url.password() {
92        None => None,
93        Some(password) => Some(decode(password)?.into_owned()),
94    };
95
96    (username, password)
97}
98
99#[throws(TrinoSourceError)]
100fn catalog_from_url(url: &url::Url) -> String {
101    match url.path_segments().and_then(|mut s| s.next_back()) {
102        Some(segment) => decode(segment)?.into_owned(),
103        None => "hive".to_owned(),
104    }
105}
106
107fn host_from_url(url: &url::Url) -> &str {
108    url.host_str().unwrap_or("localhost")
109}
110
111/// Build a prusto Client from a Trino connection URL, parsing all supported query parameters.
112///
113/// Supported URL params:
114///   source, schema, client_tags (comma-separated), client_info, trace_token,
115///   session.<key>=<value>, extra_credential.<key>=<value>, verify=false
116#[throws(TrinoSourceError)]
117pub fn build_client_from_url(url: &url::Url) -> Client {
118    let (username, password) = credentials_from_url(url)?;
119
120    let no_verify = url
121        .query_pairs()
122        .any(|(k, v)| k == "verify" && v == "false");
123
124    let catalog = catalog_from_url(url)?;
125
126    let mut builder = ClientBuilder::new(&username, host_from_url(url).to_owned())
127        .port(url.port().unwrap_or(8080))
128        .ssl(prusto::ssl::Ssl { root_cert: None })
129        .no_verify(no_verify)
130        .secure(url.scheme() == "trino+https")
131        .catalog(&catalog);
132
133    let mut session_props: HashMap<String, String> = HashMap::new();
134    let mut extra_creds: HashMap<String, String> = HashMap::new();
135
136    for (key, value) in url.query_pairs() {
137        match key.as_ref() {
138            "source" => builder = builder.source(value.as_ref()),
139            "schema" => builder = builder.schema(value.as_ref()),
140            "client_tags" => {
141                for tag in value.split(',') {
142                    let tag = tag.trim();
143                    if !tag.is_empty() {
144                        builder = builder.client_tag(tag);
145                    }
146                }
147            }
148            "client_info" => builder = builder.client_info(value.as_ref()),
149            "trace_token" => builder = builder.trace_token(value.as_ref()),
150            k if k.starts_with("session.") => {
151                if let Some(prop) = k.strip_prefix("session.").filter(|s| !s.is_empty()) {
152                    session_props.insert(prop.to_string(), value.to_string());
153                }
154            }
155            k if k.starts_with("extra_credential.") => {
156                if let Some(cred) = k
157                    .strip_prefix("extra_credential.")
158                    .filter(|s| !s.is_empty())
159                {
160                    extra_creds.insert(cred.to_string(), value.to_string());
161                }
162            }
163            _ => {}
164        }
165    }
166
167    if !session_props.is_empty() {
168        builder = builder.properties(session_props);
169    }
170    if !extra_creds.is_empty() {
171        builder = builder.extra_credentials(extra_creds);
172    }
173
174    let builder = match password {
175        None => builder,
176        Some(password) => builder.auth(Auth::Basic(username, Some(password))),
177    };
178
179    builder.build().map_err(TrinoSourceError::PrustoError)?
180}
181
182impl TrinoSource {
183    #[throws(TrinoSourceError)]
184    pub fn new(rt: Arc<Runtime>, conn: &str) -> Self {
185        // The connection string must be parsed as-is: percent-decoding it first
186        // would feed reserved characters from the password back into the URL
187        // grammar, so `p@ss` would either re-encode to `p%40ss` or, for `/`,
188        // `#` and `?`, silently truncate the authority.
189        let url = conn
190            .parse::<url::Url>()
191            .map_err(TrinoSourceError::UrlParseError)?;
192
193        let client = build_client_from_url(&url)?;
194
195        Self {
196            client: Arc::new(client),
197            rt,
198            origin_query: None,
199            queries: vec![],
200            names: vec![],
201            schema: vec![],
202        }
203    }
204}
205
206impl Source for TrinoSource
207where
208    TrinoSourcePartition: SourcePartition<TypeSystem = TrinoTypeSystem, Error = TrinoSourceError>,
209{
210    const DATA_ORDERS: &'static [DataOrder] = &[DataOrder::RowMajor];
211    type TypeSystem = TrinoTypeSystem;
212    type Partition = TrinoSourcePartition;
213    type Error = TrinoSourceError;
214
215    #[throws(TrinoSourceError)]
216    fn set_data_order(&mut self, data_order: DataOrder) {
217        if !matches!(data_order, DataOrder::RowMajor) {
218            throw!(ConnectorXError::UnsupportedDataOrder(data_order));
219        }
220    }
221
222    fn set_queries<Q: ToString>(&mut self, queries: &[CXQuery<Q>]) {
223        self.queries = queries.iter().map(|q| q.map(Q::to_string)).collect();
224    }
225
226    fn set_origin_query(&mut self, query: Option<String>) {
227        self.origin_query = query;
228    }
229
230    #[throws(TrinoSourceError)]
231    fn fetch_metadata(&mut self) {
232        assert!(!self.queries.is_empty());
233
234        let first_query = &self.queries[0];
235        let cxq = limit0_query(first_query, &GenericDialect {})?;
236
237        let dataset: DataSet<Row> = self
238            .rt
239            .block_on(self.client.get_all::<Row>(cxq.to_string()))
240            .map_err(TrinoSourceError::PrustoError)?;
241
242        let schema = dataset.split().0;
243
244        for (name, t) in schema {
245            self.names.push(name.clone());
246            self.schema.push(TrinoTypeSystem::try_from(t.clone())?);
247        }
248    }
249
250    #[throws(TrinoSourceError)]
251    fn result_rows(&mut self) -> Option<usize> {
252        match &self.origin_query {
253            Some(q) => {
254                let cxq = CXQuery::Naked(q.clone());
255                let nrows = get_total_rows(self.rt.clone(), self.client.clone(), &cxq)?;
256                Some(nrows)
257            }
258            None => None,
259        }
260    }
261
262    fn names(&self) -> Vec<String> {
263        self.names.clone()
264    }
265
266    fn schema(&self) -> Vec<Self::TypeSystem> {
267        self.schema.clone()
268    }
269
270    #[throws(TrinoSourceError)]
271    fn partition(self) -> Vec<Self::Partition> {
272        let mut ret = vec![];
273
274        for query in self.queries {
275            ret.push(TrinoSourcePartition::new(
276                self.client.clone(),
277                query,
278                self.schema.clone(),
279                self.rt.clone(),
280            )?);
281        }
282        ret
283    }
284}
285
286pub struct TrinoSourcePartition {
287    client: Arc<Client>,
288    query: CXQuery<String>,
289    schema: Vec<TrinoTypeSystem>,
290    rt: Arc<Runtime>,
291    nrows: usize,
292}
293
294impl TrinoSourcePartition {
295    #[throws(TrinoSourceError)]
296    pub fn new(
297        client: Arc<Client>,
298        query: CXQuery<String>,
299        schema: Vec<TrinoTypeSystem>,
300        rt: Arc<Runtime>,
301    ) -> Self {
302        Self {
303            client,
304            query: query.clone(),
305            schema: schema.to_vec(),
306            rt,
307            nrows: 0,
308        }
309    }
310}
311
312impl SourcePartition for TrinoSourcePartition {
313    type TypeSystem = TrinoTypeSystem;
314    type Parser<'a> = TrinoSourcePartitionParser<'a>;
315    type Error = TrinoSourceError;
316
317    #[throws(TrinoSourceError)]
318    fn result_rows(&mut self) {
319        self.nrows = get_total_rows(self.rt.clone(), self.client.clone(), &self.query)?;
320    }
321
322    #[throws(TrinoSourceError)]
323    fn parser(&mut self) -> Self::Parser<'_> {
324        TrinoSourcePartitionParser::new(
325            self.rt.clone(),
326            self.client.clone(),
327            self.query.clone(),
328            &self.schema,
329        )?
330    }
331
332    fn nrows(&self) -> usize {
333        self.nrows
334    }
335
336    fn ncols(&self) -> usize {
337        self.schema.len()
338    }
339}
340
341pub struct TrinoSourcePartitionParser<'a> {
342    rt: Arc<Runtime>,
343    client: Arc<Client>,
344    next_uri: Option<String>,
345    rows: Vec<Row>,
346    ncols: usize,
347    current_col: usize,
348    current_row: usize,
349    _phantom: &'a PhantomData<DataSet<Row>>,
350}
351
352impl<'a> TrinoSourcePartitionParser<'a> {
353    #[throws(TrinoSourceError)]
354    pub fn new(
355        rt: Arc<Runtime>,
356        client: Arc<Client>,
357        query: CXQuery,
358        schema: &[TrinoTypeSystem],
359    ) -> Self {
360        let results = rt
361            .block_on(client.get::<Row>(query.to_string()))
362            .map_err(TrinoSourceError::PrustoError)?;
363
364        let rows = match results.data_set {
365            Some(x) => x.into_vec(),
366            _ => vec![],
367        };
368
369        Self {
370            rt,
371            client,
372            next_uri: results.next_uri,
373            rows,
374            ncols: schema.len(),
375            current_row: 0,
376            current_col: 0,
377            _phantom: &PhantomData,
378        }
379    }
380
381    #[throws(TrinoSourceError)]
382    fn next_loc(&mut self) -> (usize, usize) {
383        let ret = (self.current_row, self.current_col);
384        self.current_row += (self.current_col + 1) / self.ncols;
385        self.current_col = (self.current_col + 1) % self.ncols;
386        ret
387    }
388}
389
390impl<'a> PartitionParser<'a> for TrinoSourcePartitionParser<'a> {
391    type TypeSystem = TrinoTypeSystem;
392    type Error = TrinoSourceError;
393
394    #[throws(TrinoSourceError)]
395    fn fetch_next(&mut self) -> (usize, bool) {
396        assert!(self.current_col == 0);
397
398        match self.next_uri.clone() {
399            Some(uri) => {
400                let results = self
401                    .rt
402                    .block_on(self.client.get_next::<Row>(&uri))
403                    .map_err(TrinoSourceError::PrustoError)?;
404
405                self.rows = match results.data_set {
406                    Some(x) => x.into_vec(),
407                    _ => vec![],
408                };
409
410                self.current_row = 0;
411                self.next_uri = results.next_uri;
412
413                (self.rows.len(), false)
414            }
415            None => return (self.rows.len(), true),
416        }
417    }
418}
419
420macro_rules! impl_produce_int {
421    ($($t: ty,)+) => {
422        $(
423            impl<'r, 'a> Produce<'r, $t> for TrinoSourcePartitionParser<'a> {
424                type Error = TrinoSourceError;
425
426                #[throws(TrinoSourceError)]
427                fn produce(&'r mut self) -> $t {
428                    let (ridx, cidx) = self.next_loc()?;
429                    let value = &self.rows[ridx].value()[cidx];
430
431                    match value {
432                        Value::Number(x) => {
433                            if (x.is_i64()) {
434                                <$t>::try_from(x.as_i64().unwrap()).map_err(|_| anyhow!("Trino cannot parse i64 at position: ({}, {}) {:?}", ridx, cidx, value))?
435                            } else {
436                                throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, x))
437                            }
438                        }
439                        _ => throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, value))
440                    }
441                }
442            }
443
444            impl<'r, 'a> Produce<'r, Option<$t>> for TrinoSourcePartitionParser<'a> {
445                type Error = TrinoSourceError;
446
447                #[throws(TrinoSourceError)]
448                fn produce(&'r mut self) -> Option<$t> {
449                    let (ridx, cidx) = self.next_loc()?;
450                    let value = &self.rows[ridx].value()[cidx];
451
452                    match value {
453                        Value::Null => None,
454                        Value::Number(x) => {
455                            if (x.is_i64()) {
456                                Some(<$t>::try_from(x.as_i64().unwrap()).map_err(|_| anyhow!("Trino cannot parse i64 at position: ({}, {}) {:?}", ridx, cidx, value))?)
457                            } else {
458                                throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, x))
459                            }
460                        }
461                        _ => throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, value))
462                    }
463                }
464            }
465        )+
466    };
467}
468
469macro_rules! impl_produce_float {
470    ($($t: ty,)+) => {
471        $(
472            impl<'r, 'a> Produce<'r, $t> for TrinoSourcePartitionParser<'a> {
473                type Error = TrinoSourceError;
474
475                #[throws(TrinoSourceError)]
476                fn produce(&'r mut self) -> $t {
477                    let (ridx, cidx) = self.next_loc()?;
478                    let value = &self.rows[ridx].value()[cidx];
479
480                    match value {
481                        Value::Number(x) => {
482                            if (x.is_f64()) {
483                                x.as_f64().unwrap() as $t
484                            } else {
485                                throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, x))
486                            }
487                        }
488                        Value::String(x) => x.parse::<$t>().map_err(|_| anyhow!("Trino cannot parse String at position: ({}, {}) {:?}", ridx, cidx, value))?,
489                        _ => throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, value))
490                    }
491                }
492            }
493
494            impl<'r, 'a> Produce<'r, Option<$t>> for TrinoSourcePartitionParser<'a> {
495                type Error = TrinoSourceError;
496
497                #[throws(TrinoSourceError)]
498                fn produce(&'r mut self) -> Option<$t> {
499                    let (ridx, cidx) = self.next_loc()?;
500                    let value = &self.rows[ridx].value()[cidx];
501
502                    match value {
503                        Value::Null => None,
504                        Value::Number(x) => {
505                            if (x.is_f64()) {
506                                Some(x.as_f64().unwrap() as $t)
507                            } else {
508                                throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, x))
509                            }
510                        }
511                        Value::String(x) => Some(x.parse::<$t>().map_err(|_| anyhow!("Trino cannot parse String at position: ({}, {}) {:?}", ridx, cidx, value))?),
512                        _ => throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, value))
513                    }
514                }
515            }
516        )+
517    };
518}
519
520macro_rules! impl_produce_text {
521    ($($t: ty,)+) => {
522        $(
523            impl<'r, 'a> Produce<'r, $t> for TrinoSourcePartitionParser<'a> {
524                type Error = TrinoSourceError;
525
526                #[throws(TrinoSourceError)]
527                fn produce(&'r mut self) -> $t {
528                    let (ridx, cidx) = self.next_loc()?;
529                    let value = &self.rows[ridx].value()[cidx];
530
531                    match value {
532                        Value::String(x) => {
533                            x.parse().map_err(|_| anyhow!("Trino cannot parse String at position: ({}, {}): {:?}", ridx, cidx, value))?
534                        }
535                        Value::Array(_) | Value::Object(_) | Value::Number(_) | Value::Bool(_) => {
536                            serde_json::to_string(value)
537                                .unwrap_or_else(|_| value.to_string())
538                                .parse()
539                                .map_err(|_| anyhow!("Trino cannot convert complex value at ({}, {}): {:?}", ridx, cidx, value))?
540                        }
541                        _ => throw!(anyhow!("Trino unknown value at position: ({}, {}): {:?}", ridx, cidx, value))
542                    }
543                }
544            }
545
546            impl<'r, 'a> Produce<'r, Option<$t>> for TrinoSourcePartitionParser<'a> {
547                type Error = TrinoSourceError;
548
549                #[throws(TrinoSourceError)]
550                fn produce(&'r mut self) -> Option<$t> {
551                    let (ridx, cidx) = self.next_loc()?;
552                    let value = &self.rows[ridx].value()[cidx];
553
554                    match value {
555                        Value::Null => None,
556                        Value::String(x) => {
557                            Some(x.parse().map_err(|_| anyhow!("Trino cannot parse String at position: ({}, {}): {:?}", ridx, cidx, value))?)
558                        }
559                        Value::Array(_) | Value::Object(_) | Value::Number(_) | Value::Bool(_) => {
560                            Some(serde_json::to_string(value)
561                                .unwrap_or_else(|_| value.to_string())
562                                .parse()
563                                .map_err(|_| anyhow!("Trino cannot convert complex value at ({}, {}): {:?}", ridx, cidx, value))?)
564                        }
565                    }
566                }
567            }
568        )+
569    };
570}
571
572macro_rules! impl_produce_timestamp {
573    ($($t: ty,)+) => {
574        $(
575            impl<'r, 'a> Produce<'r, $t> for TrinoSourcePartitionParser<'a> {
576                type Error = TrinoSourceError;
577
578                #[throws(TrinoSourceError)]
579                fn produce(&'r mut self) -> $t {
580                    let (ridx, cidx) = self.next_loc()?;
581                    let value = &self.rows[ridx].value()[cidx];
582
583                    match value {
584                        Value::String(x) => NaiveDateTime::parse_from_str(x, "%Y-%m-%d %H:%M:%S%.f").map_err(|_| anyhow!("Trino cannot parse String at position: ({}, {}): {:?}", ridx, cidx, value))?,
585                        _ => throw!(anyhow!("Trino unknown value at position: ({}, {}): {:?}", ridx, cidx, value))
586                    }
587                }
588            }
589
590            impl<'r, 'a> Produce<'r, Option<$t>> for TrinoSourcePartitionParser<'a> {
591                type Error = TrinoSourceError;
592
593                #[throws(TrinoSourceError)]
594                fn produce(&'r mut self) -> Option<$t> {
595                    let (ridx, cidx) = self.next_loc()?;
596                    let value = &self.rows[ridx].value()[cidx];
597
598                    match value {
599                        Value::Null => None,
600                        Value::String(x) => Some(NaiveDateTime::parse_from_str(x, "%Y-%m-%d %H:%M:%S%.f").map_err(|_| anyhow!("Trino cannot parse String at position: ({}, {}): {:?}", ridx, cidx, value))?),
601                        _ => throw!(anyhow!("Trino unknown value at position: ({}, {}): {:?}", ridx, cidx, value))
602                    }
603                }
604            }
605        )+
606    };
607}
608
609macro_rules! impl_produce_bool {
610    ($($t: ty,)+) => {
611        $(
612            impl<'r, 'a> Produce<'r, $t> for TrinoSourcePartitionParser<'a> {
613                type Error = TrinoSourceError;
614
615                #[throws(TrinoSourceError)]
616                fn produce(&'r mut self) -> $t {
617                    let (ridx, cidx) = self.next_loc()?;
618                    let value = &self.rows[ridx].value()[cidx];
619
620                    match value {
621                        Value::Bool(x) => *x,
622                        _ => throw!(anyhow!("Trino unknown value at position: ({}, {}): {:?}", ridx, cidx, value))
623                    }
624                }
625            }
626
627            impl<'r, 'a> Produce<'r, Option<$t>> for TrinoSourcePartitionParser<'a> {
628                type Error = TrinoSourceError;
629
630                #[throws(TrinoSourceError)]
631                fn produce(&'r mut self) -> Option<$t> {
632                    let (ridx, cidx) = self.next_loc()?;
633                    let value = &self.rows[ridx].value()[cidx];
634
635                    match value {
636                        Value::Null => None,
637                        Value::Bool(x) => Some(*x),
638                        _ => throw!(anyhow!("Trino unknown value at position: ({}, {}): {:?}", ridx, cidx, value))
639                    }
640                }
641            }
642        )+
643    };
644}
645
646impl_produce_bool!(bool,);
647impl_produce_int!(i8, i16, i32, i64,);
648impl_produce_float!(f32, f64,);
649impl_produce_timestamp!(NaiveDateTime,);
650impl_produce_text!(String, char,);
651
652impl<'r, 'a> Produce<'r, NaiveTime> for TrinoSourcePartitionParser<'a> {
653    type Error = TrinoSourceError;
654
655    #[throws(TrinoSourceError)]
656    fn produce(&'r mut self) -> NaiveTime {
657        let (ridx, cidx) = self.next_loc()?;
658        let value = &self.rows[ridx].value()[cidx];
659
660        match value {
661            Value::String(x) => NaiveTime::parse_from_str(x, "%H:%M:%S%.f").map_err(|_| {
662                anyhow!(
663                    "Trino cannot parse String at position: ({}, {}): {:?}",
664                    ridx,
665                    cidx,
666                    value
667                )
668            })?,
669            _ => throw!(anyhow!(
670                "Trino unknown value at position: ({}, {}): {:?}",
671                ridx,
672                cidx,
673                value
674            )),
675        }
676    }
677}
678
679impl<'r, 'a> Produce<'r, Option<NaiveTime>> for TrinoSourcePartitionParser<'a> {
680    type Error = TrinoSourceError;
681
682    #[throws(TrinoSourceError)]
683    fn produce(&'r mut self) -> Option<NaiveTime> {
684        let (ridx, cidx) = self.next_loc()?;
685        let value = &self.rows[ridx].value()[cidx];
686
687        match value {
688            Value::Null => None,
689            Value::String(x) => {
690                Some(NaiveTime::parse_from_str(x, "%H:%M:%S%.f").map_err(|_| {
691                    anyhow!(
692                        "Trino cannot parse Time at position: ({}, {}): {:?}",
693                        ridx,
694                        cidx,
695                        value
696                    )
697                })?)
698            }
699            _ => throw!(anyhow!(
700                "Trino unknown value at position: ({}, {}): {:?}",
701                ridx,
702                cidx,
703                value
704            )),
705        }
706    }
707}
708
709impl<'r, 'a> Produce<'r, NaiveDate> for TrinoSourcePartitionParser<'a> {
710    type Error = TrinoSourceError;
711
712    #[throws(TrinoSourceError)]
713    fn produce(&'r mut self) -> NaiveDate {
714        let (ridx, cidx) = self.next_loc()?;
715        let value = &self.rows[ridx].value()[cidx];
716
717        match value {
718            Value::String(x) => NaiveDate::parse_from_str(x, "%Y-%m-%d").map_err(|_| {
719                anyhow!(
720                    "Trino cannot parse Date at position: ({}, {}): {:?}",
721                    ridx,
722                    cidx,
723                    value
724                )
725            })?,
726            _ => throw!(anyhow!(
727                "Trino unknown value at position: ({}, {}): {:?}",
728                ridx,
729                cidx,
730                value
731            )),
732        }
733    }
734}
735
736impl<'r, 'a> Produce<'r, Option<NaiveDate>> for TrinoSourcePartitionParser<'a> {
737    type Error = TrinoSourceError;
738
739    #[throws(TrinoSourceError)]
740    fn produce(&'r mut self) -> Option<NaiveDate> {
741        let (ridx, cidx) = self.next_loc()?;
742        let value = &self.rows[ridx].value()[cidx];
743
744        match value {
745            Value::Null => None,
746            Value::String(x) => Some(NaiveDate::parse_from_str(x, "%Y-%m-%d").map_err(|_| {
747                anyhow!(
748                    "Trino cannot parse Date at position: ({}, {}): {:?}",
749                    ridx,
750                    cidx,
751                    value
752                )
753            })?),
754            _ => throw!(anyhow!(
755                "Trino unknown value at position: ({}, {}): {:?}",
756                ridx,
757                cidx,
758                value
759            )),
760        }
761    }
762}
763
764#[cfg(test)]
765mod tests {
766    use super::*;
767
768    #[test]
769    fn test_new_with_all_params() {
770        let rt = Arc::new(Runtime::new().unwrap());
771        let conn = "trino+https://myuser:mypass@localhost:8443/mycatalog?\
772            source=connectorx-test&schema=myschema&\
773            client_tags=tag1,tag2,tag3&client_info=test-run&\
774            trace_token=abc123&\
775            session.query_max_execution_time=10m&\
776            extra_credential.token=secret123&verify=false";
777        assert!(TrinoSource::new(rt, conn).is_ok());
778    }
779
780    #[test]
781    fn test_new_minimal_url() {
782        let rt = Arc::new(Runtime::new().unwrap());
783        assert!(TrinoSource::new(rt, "trino://test@localhost:8080/memory").is_ok());
784    }
785
786    #[test]
787    fn catalog_is_percent_decoded() {
788        let url = "trino://test@localhost:8080/my%5Fcatalog"
789            .parse::<url::Url>()
790            .unwrap();
791
792        assert_eq!(catalog_from_url(&url).unwrap(), "my_catalog");
793    }
794
795    #[test]
796    fn catalog_defaults_to_hive_without_path_segments() {
797        let url = "trino:memory".parse::<url::Url>().unwrap();
798
799        assert_eq!(catalog_from_url(&url).unwrap(), "hive");
800    }
801
802    #[test]
803    fn hostless_url_defaults_to_localhost() {
804        let url = "trino:///memory".parse::<url::Url>().unwrap();
805
806        assert_eq!(host_from_url(&url), "localhost");
807    }
808
809    #[test]
810    fn test_new_ignores_empty_keys() {
811        let rt = Arc::new(Runtime::new().unwrap());
812        let conn = "trino://test@localhost:8080/memory?session.=val&extra_credential.=x";
813        assert!(TrinoSource::new(rt, conn).is_ok());
814    }
815
816    /// Reserved characters in the password must survive the round trip through
817    /// the URL: percent-encoded on the way in, decoded again for the
818    /// `Authorization` header.
819    #[test]
820    fn credentials_are_percent_decoded() {
821        for (encoded, expected) in [
822            ("p%40ss", "p@ss"),
823            ("pa%3Ass", "pa:ss"),
824            ("pa%2Fss", "pa/ss"),
825            ("pa%23ss", "pa#ss"),
826            ("pa%3Fss", "pa?ss"),
827            ("pa%20ss", "pa ss"),
828            ("plain", "plain"),
829        ] {
830            let url = format!("trino+https://me:{encoded}@localhost:8443/hive")
831                .parse::<url::Url>()
832                .unwrap();
833            let (username, password) = credentials_from_url(&url).unwrap();
834
835            assert_eq!(username, "me");
836            assert_eq!(password.as_deref(), Some(expected), "for {encoded}");
837        }
838    }
839
840    /// A percent-encoded username is decoded too, since it also travels in the
841    /// `Authorization` header and in `X-Trino-User`.
842    #[test]
843    fn username_is_percent_decoded() {
844        let url = "trino+https://svc%2Dname:pw@localhost:8443/hive"
845            .parse::<url::Url>()
846            .unwrap();
847        let (username, _) = credentials_from_url(&url).unwrap();
848
849        assert_eq!(username, "svc-name");
850    }
851
852    /// No userinfo at all still yields the default Trino user and no auth.
853    #[test]
854    fn credentials_default_without_userinfo() {
855        let url = "trino://localhost:8080/memory".parse::<url::Url>().unwrap();
856        let (username, password) = credentials_from_url(&url).unwrap();
857
858        assert_eq!(username, "connectorx");
859        assert_eq!(password, None);
860    }
861}
connectorx/sources/trino/mod.rs

connectorx/sources/trino/
mod.rs
