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
58impl 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#[throws(TrinoSourceError)]
84pub fn build_client_from_url(url: &url::Url) -> Client {
85 let username = match url.username() {
86 "" => "connectorx",
87 username => username,
88 };
89
90 let no_verify = url
91 .query_pairs()
92 .any(|(k, v)| k == "verify" && v == "false");
93
94 let mut builder = ClientBuilder::new(username, url.host().unwrap().to_owned())
95 .port(url.port().unwrap_or(8080))
96 .ssl(prusto::ssl::Ssl { root_cert: None })
97 .no_verify(no_verify)
98 .secure(url.scheme() == "trino+https")
99 .catalog(url.path_segments().unwrap().last().unwrap_or("hive"));
100
101 let mut session_props: HashMap<String, String> = HashMap::new();
102 let mut extra_creds: HashMap<String, String> = HashMap::new();
103
104 for (key, value) in url.query_pairs() {
105 match key.as_ref() {
106 "source" => builder = builder.source(value.as_ref()),
107 "schema" => builder = builder.schema(value.as_ref()),
108 "client_tags" => {
109 for tag in value.split(',') {
110 let tag = tag.trim();
111 if !tag.is_empty() {
112 builder = builder.client_tag(tag);
113 }
114 }
115 }
116 "client_info" => builder = builder.client_info(value.as_ref()),
117 "trace_token" => builder = builder.trace_token(value.as_ref()),
118 k if k.starts_with("session.") => {
119 if let Some(prop) = k.strip_prefix("session.").filter(|s| !s.is_empty()) {
120 session_props.insert(prop.to_string(), value.to_string());
121 }
122 }
123 k if k.starts_with("extra_credential.") => {
124 if let Some(cred) = k
125 .strip_prefix("extra_credential.")
126 .filter(|s| !s.is_empty())
127 {
128 extra_creds.insert(cred.to_string(), value.to_string());
129 }
130 }
131 _ => {}
132 }
133 }
134
135 if !session_props.is_empty() {
136 builder = builder.properties(session_props);
137 }
138 if !extra_creds.is_empty() {
139 builder = builder.extra_credentials(extra_creds);
140 }
141
142 let builder = match url.password() {
143 None => builder,
144 Some(password) => builder.auth(Auth::Basic(username.to_owned(), Some(password.to_owned()))),
145 };
146
147 builder.build().map_err(TrinoSourceError::PrustoError)?
148}
149
150impl TrinoSource {
151 #[throws(TrinoSourceError)]
152 pub fn new(rt: Arc<Runtime>, conn: &str) -> Self {
153 let decoded_conn = decode(conn)?.into_owned();
154
155 let url = decoded_conn
156 .parse::<url::Url>()
157 .map_err(TrinoSourceError::UrlParseError)?;
158
159 let username = match url.username() {
160 "" => "connectorx",
161 username => username,
162 };
163
164 let no_verify = url
165 .query_pairs()
166 .any(|(k, v)| k == "verify" && v == "false");
167
168 let builder = ClientBuilder::new(username, url.host().unwrap().to_owned())
169 .port(url.port().unwrap_or(8080))
170 .ssl(prusto::ssl::Ssl { root_cert: None })
171 .no_verify(no_verify)
172 .secure(url.scheme() == "trino+https")
173 .catalog(url.path_segments().unwrap().last().unwrap_or("hive"));
174
175 let builder = match url.password() {
176 None => builder,
177 Some(password) => {
178 builder.auth(Auth::Basic(username.to_owned(), Some(password.to_owned())))
179 }
180 };
181
182 let client = builder.build().map_err(TrinoSourceError::PrustoError)?;
183
184 Self {
185 client: Arc::new(client),
186 rt,
187 origin_query: None,
188 queries: vec![],
189 names: vec![],
190 schema: vec![],
191 }
192 }
193}
194
195impl Source for TrinoSource
196where
197 TrinoSourcePartition: SourcePartition<TypeSystem = TrinoTypeSystem, Error = TrinoSourceError>,
198{
199 const DATA_ORDERS: &'static [DataOrder] = &[DataOrder::RowMajor];
200 type TypeSystem = TrinoTypeSystem;
201 type Partition = TrinoSourcePartition;
202 type Error = TrinoSourceError;
203
204 #[throws(TrinoSourceError)]
205 fn set_data_order(&mut self, data_order: DataOrder) {
206 if !matches!(data_order, DataOrder::RowMajor) {
207 throw!(ConnectorXError::UnsupportedDataOrder(data_order));
208 }
209 }
210
211 fn set_queries<Q: ToString>(&mut self, queries: &[CXQuery<Q>]) {
212 self.queries = queries.iter().map(|q| q.map(Q::to_string)).collect();
213 }
214
215 fn set_origin_query(&mut self, query: Option<String>) {
216 self.origin_query = query;
217 }
218
219 #[throws(TrinoSourceError)]
220 fn fetch_metadata(&mut self) {
221 assert!(!self.queries.is_empty());
222
223 let first_query = &self.queries[0];
224 let cxq = limit0_query(first_query, &GenericDialect {})?;
225
226 let dataset: DataSet<Row> = self
227 .rt
228 .block_on(self.client.get_all::<Row>(cxq.to_string()))
229 .map_err(TrinoSourceError::PrustoError)?;
230
231 let schema = dataset.split().0;
232
233 for (name, t) in schema {
234 self.names.push(name.clone());
235 self.schema.push(TrinoTypeSystem::try_from(t.clone())?);
236 }
237 }
238
239 #[throws(TrinoSourceError)]
240 fn result_rows(&mut self) -> Option<usize> {
241 match &self.origin_query {
242 Some(q) => {
243 let cxq = CXQuery::Naked(q.clone());
244 let nrows = get_total_rows(self.rt.clone(), self.client.clone(), &cxq)?;
245 Some(nrows)
246 }
247 None => None,
248 }
249 }
250
251 fn names(&self) -> Vec<String> {
252 self.names.clone()
253 }
254
255 fn schema(&self) -> Vec<Self::TypeSystem> {
256 self.schema.clone()
257 }
258
259 #[throws(TrinoSourceError)]
260 fn partition(self) -> Vec<Self::Partition> {
261 let mut ret = vec![];
262
263 for query in self.queries {
264 ret.push(TrinoSourcePartition::new(
265 self.client.clone(),
266 query,
267 self.schema.clone(),
268 self.rt.clone(),
269 )?);
270 }
271 ret
272 }
273}
274
275pub struct TrinoSourcePartition {
276 client: Arc<Client>,
277 query: CXQuery<String>,
278 schema: Vec<TrinoTypeSystem>,
279 rt: Arc<Runtime>,
280 nrows: usize,
281}
282
283impl TrinoSourcePartition {
284 #[throws(TrinoSourceError)]
285 pub fn new(
286 client: Arc<Client>,
287 query: CXQuery<String>,
288 schema: Vec<TrinoTypeSystem>,
289 rt: Arc<Runtime>,
290 ) -> Self {
291 Self {
292 client,
293 query: query.clone(),
294 schema: schema.to_vec(),
295 rt,
296 nrows: 0,
297 }
298 }
299}
300
301impl SourcePartition for TrinoSourcePartition {
302 type TypeSystem = TrinoTypeSystem;
303 type Parser<'a> = TrinoSourcePartitionParser<'a>;
304 type Error = TrinoSourceError;
305
306 #[throws(TrinoSourceError)]
307 fn result_rows(&mut self) {
308 self.nrows = get_total_rows(self.rt.clone(), self.client.clone(), &self.query)?;
309 }
310
311 #[throws(TrinoSourceError)]
312 fn parser(&mut self) -> Self::Parser<'_> {
313 TrinoSourcePartitionParser::new(
314 self.rt.clone(),
315 self.client.clone(),
316 self.query.clone(),
317 &self.schema,
318 )?
319 }
320
321 fn nrows(&self) -> usize {
322 self.nrows
323 }
324
325 fn ncols(&self) -> usize {
326 self.schema.len()
327 }
328}
329
330pub struct TrinoSourcePartitionParser<'a> {
331 rt: Arc<Runtime>,
332 client: Arc<Client>,
333 next_uri: Option<String>,
334 rows: Vec<Row>,
335 ncols: usize,
336 current_col: usize,
337 current_row: usize,
338 _phantom: &'a PhantomData<DataSet<Row>>,
339}
340
341impl<'a> TrinoSourcePartitionParser<'a> {
342 #[throws(TrinoSourceError)]
343 pub fn new(
344 rt: Arc<Runtime>,
345 client: Arc<Client>,
346 query: CXQuery,
347 schema: &[TrinoTypeSystem],
348 ) -> Self {
349 let results = rt
350 .block_on(client.get::<Row>(query.to_string()))
351 .map_err(TrinoSourceError::PrustoError)?;
352
353 let rows = match results.data_set {
354 Some(x) => x.into_vec(),
355 _ => vec![],
356 };
357
358 Self {
359 rt,
360 client,
361 next_uri: results.next_uri,
362 rows,
363 ncols: schema.len(),
364 current_row: 0,
365 current_col: 0,
366 _phantom: &PhantomData,
367 }
368 }
369
370 #[throws(TrinoSourceError)]
371 fn next_loc(&mut self) -> (usize, usize) {
372 let ret = (self.current_row, self.current_col);
373 self.current_row += (self.current_col + 1) / self.ncols;
374 self.current_col = (self.current_col + 1) % self.ncols;
375 ret
376 }
377}
378
379impl<'a> PartitionParser<'a> for TrinoSourcePartitionParser<'a> {
380 type TypeSystem = TrinoTypeSystem;
381 type Error = TrinoSourceError;
382
383 #[throws(TrinoSourceError)]
384 fn fetch_next(&mut self) -> (usize, bool) {
385 assert!(self.current_col == 0);
386
387 match self.next_uri.clone() {
388 Some(uri) => {
389 let results = self
390 .rt
391 .block_on(self.client.get_next::<Row>(&uri))
392 .map_err(TrinoSourceError::PrustoError)?;
393
394 self.rows = match results.data_set {
395 Some(x) => x.into_vec(),
396 _ => vec![],
397 };
398
399 self.current_row = 0;
400 self.next_uri = results.next_uri;
401
402 (self.rows.len(), false)
403 }
404 None => return (self.rows.len(), true),
405 }
406 }
407}
408
409macro_rules! impl_produce_int {
410 ($($t: ty,)+) => {
411 $(
412 impl<'r, 'a> Produce<'r, $t> for TrinoSourcePartitionParser<'a> {
413 type Error = TrinoSourceError;
414
415 #[throws(TrinoSourceError)]
416 fn produce(&'r mut self) -> $t {
417 let (ridx, cidx) = self.next_loc()?;
418 let value = &self.rows[ridx].value()[cidx];
419
420 match value {
421 Value::Number(x) => {
422 if (x.is_i64()) {
423 <$t>::try_from(x.as_i64().unwrap()).map_err(|_| anyhow!("Trino cannot parse i64 at position: ({}, {}) {:?}", ridx, cidx, value))?
424 } else {
425 throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, x))
426 }
427 }
428 _ => throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, value))
429 }
430 }
431 }
432
433 impl<'r, 'a> Produce<'r, Option<$t>> for TrinoSourcePartitionParser<'a> {
434 type Error = TrinoSourceError;
435
436 #[throws(TrinoSourceError)]
437 fn produce(&'r mut self) -> Option<$t> {
438 let (ridx, cidx) = self.next_loc()?;
439 let value = &self.rows[ridx].value()[cidx];
440
441 match value {
442 Value::Null => None,
443 Value::Number(x) => {
444 if (x.is_i64()) {
445 Some(<$t>::try_from(x.as_i64().unwrap()).map_err(|_| anyhow!("Trino cannot parse i64 at position: ({}, {}) {:?}", ridx, cidx, value))?)
446 } else {
447 throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, x))
448 }
449 }
450 _ => throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, value))
451 }
452 }
453 }
454 )+
455 };
456}
457
458macro_rules! impl_produce_float {
459 ($($t: ty,)+) => {
460 $(
461 impl<'r, 'a> Produce<'r, $t> for TrinoSourcePartitionParser<'a> {
462 type Error = TrinoSourceError;
463
464 #[throws(TrinoSourceError)]
465 fn produce(&'r mut self) -> $t {
466 let (ridx, cidx) = self.next_loc()?;
467 let value = &self.rows[ridx].value()[cidx];
468
469 match value {
470 Value::Number(x) => {
471 if (x.is_f64()) {
472 x.as_f64().unwrap() as $t
473 } else {
474 throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, x))
475 }
476 }
477 Value::String(x) => x.parse::<$t>().map_err(|_| anyhow!("Trino cannot parse String at position: ({}, {}) {:?}", ridx, cidx, value))?,
478 _ => throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, value))
479 }
480 }
481 }
482
483 impl<'r, 'a> Produce<'r, Option<$t>> for TrinoSourcePartitionParser<'a> {
484 type Error = TrinoSourceError;
485
486 #[throws(TrinoSourceError)]
487 fn produce(&'r mut self) -> Option<$t> {
488 let (ridx, cidx) = self.next_loc()?;
489 let value = &self.rows[ridx].value()[cidx];
490
491 match value {
492 Value::Null => None,
493 Value::Number(x) => {
494 if (x.is_f64()) {
495 Some(x.as_f64().unwrap() as $t)
496 } else {
497 throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, x))
498 }
499 }
500 Value::String(x) => Some(x.parse::<$t>().map_err(|_| anyhow!("Trino cannot parse String at position: ({}, {}) {:?}", ridx, cidx, value))?),
501 _ => throw!(anyhow!("Trino cannot parse Number at position: ({}, {}) {:?}", ridx, cidx, value))
502 }
503 }
504 }
505 )+
506 };
507}
508
509macro_rules! impl_produce_text {
510 ($($t: ty,)+) => {
511 $(
512 impl<'r, 'a> Produce<'r, $t> for TrinoSourcePartitionParser<'a> {
513 type Error = TrinoSourceError;
514
515 #[throws(TrinoSourceError)]
516 fn produce(&'r mut self) -> $t {
517 let (ridx, cidx) = self.next_loc()?;
518 let value = &self.rows[ridx].value()[cidx];
519
520 match value {
521 Value::String(x) => {
522 x.parse().map_err(|_| anyhow!("Trino cannot parse String at position: ({}, {}): {:?}", ridx, cidx, value))?
523 }
524 Value::Array(_) | Value::Object(_) | Value::Number(_) | Value::Bool(_) => {
525 serde_json::to_string(value)
526 .unwrap_or_else(|_| value.to_string())
527 .parse()
528 .map_err(|_| anyhow!("Trino cannot convert complex value at ({}, {}): {:?}", ridx, cidx, value))?
529 }
530 _ => throw!(anyhow!("Trino unknown value at position: ({}, {}): {:?}", ridx, cidx, value))
531 }
532 }
533 }
534
535 impl<'r, 'a> Produce<'r, Option<$t>> for TrinoSourcePartitionParser<'a> {
536 type Error = TrinoSourceError;
537
538 #[throws(TrinoSourceError)]
539 fn produce(&'r mut self) -> Option<$t> {
540 let (ridx, cidx) = self.next_loc()?;
541 let value = &self.rows[ridx].value()[cidx];
542
543 match value {
544 Value::Null => None,
545 Value::String(x) => {
546 Some(x.parse().map_err(|_| anyhow!("Trino cannot parse String at position: ({}, {}): {:?}", ridx, cidx, value))?)
547 }
548 Value::Array(_) | Value::Object(_) | Value::Number(_) | Value::Bool(_) => {
549 Some(serde_json::to_string(value)
550 .unwrap_or_else(|_| value.to_string())
551 .parse()
552 .map_err(|_| anyhow!("Trino cannot convert complex value at ({}, {}): {:?}", ridx, cidx, value))?)
553 }
554 _ => throw!(anyhow!("Trino unknown value at position: ({}, {}): {:?}", ridx, cidx, value))
555 }
556 }
557 }
558 )+
559 };
560}
561
562macro_rules! impl_produce_timestamp {
563 ($($t: ty,)+) => {
564 $(
565 impl<'r, 'a> Produce<'r, $t> for TrinoSourcePartitionParser<'a> {
566 type Error = TrinoSourceError;
567
568 #[throws(TrinoSourceError)]
569 fn produce(&'r mut self) -> $t {
570 let (ridx, cidx) = self.next_loc()?;
571 let value = &self.rows[ridx].value()[cidx];
572
573 match value {
574 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))?,
575 _ => throw!(anyhow!("Trino unknown value at position: ({}, {}): {:?}", ridx, cidx, value))
576 }
577 }
578 }
579
580 impl<'r, 'a> Produce<'r, Option<$t>> for TrinoSourcePartitionParser<'a> {
581 type Error = TrinoSourceError;
582
583 #[throws(TrinoSourceError)]
584 fn produce(&'r mut self) -> Option<$t> {
585 let (ridx, cidx) = self.next_loc()?;
586 let value = &self.rows[ridx].value()[cidx];
587
588 match value {
589 Value::Null => None,
590 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))?),
591 _ => throw!(anyhow!("Trino unknown value at position: ({}, {}): {:?}", ridx, cidx, value))
592 }
593 }
594 }
595 )+
596 };
597}
598
599macro_rules! impl_produce_bool {
600 ($($t: ty,)+) => {
601 $(
602 impl<'r, 'a> Produce<'r, $t> for TrinoSourcePartitionParser<'a> {
603 type Error = TrinoSourceError;
604
605 #[throws(TrinoSourceError)]
606 fn produce(&'r mut self) -> $t {
607 let (ridx, cidx) = self.next_loc()?;
608 let value = &self.rows[ridx].value()[cidx];
609
610 match value {
611 Value::Bool(x) => *x,
612 _ => throw!(anyhow!("Trino unknown value at position: ({}, {}): {:?}", ridx, cidx, value))
613 }
614 }
615 }
616
617 impl<'r, 'a> Produce<'r, Option<$t>> for TrinoSourcePartitionParser<'a> {
618 type Error = TrinoSourceError;
619
620 #[throws(TrinoSourceError)]
621 fn produce(&'r mut self) -> Option<$t> {
622 let (ridx, cidx) = self.next_loc()?;
623 let value = &self.rows[ridx].value()[cidx];
624
625 match value {
626 Value::Null => None,
627 Value::Bool(x) => Some(*x),
628 _ => throw!(anyhow!("Trino unknown value at position: ({}, {}): {:?}", ridx, cidx, value))
629 }
630 }
631 }
632 )+
633 };
634}
635
636impl_produce_bool!(bool,);
637impl_produce_int!(i8, i16, i32, i64,);
638impl_produce_float!(f32, f64,);
639impl_produce_timestamp!(NaiveDateTime,);
640impl_produce_text!(String, char,);
641
642impl<'r, 'a> Produce<'r, NaiveTime> for TrinoSourcePartitionParser<'a> {
643 type Error = TrinoSourceError;
644
645 #[throws(TrinoSourceError)]
646 fn produce(&'r mut self) -> NaiveTime {
647 let (ridx, cidx) = self.next_loc()?;
648 let value = &self.rows[ridx].value()[cidx];
649
650 match value {
651 Value::String(x) => NaiveTime::parse_from_str(x, "%H:%M:%S%.f").map_err(|_| {
652 anyhow!(
653 "Trino cannot parse String at position: ({}, {}): {:?}",
654 ridx,
655 cidx,
656 value
657 )
658 })?,
659 _ => throw!(anyhow!(
660 "Trino unknown value at position: ({}, {}): {:?}",
661 ridx,
662 cidx,
663 value
664 )),
665 }
666 }
667}
668
669impl<'r, 'a> Produce<'r, Option<NaiveTime>> for TrinoSourcePartitionParser<'a> {
670 type Error = TrinoSourceError;
671
672 #[throws(TrinoSourceError)]
673 fn produce(&'r mut self) -> Option<NaiveTime> {
674 let (ridx, cidx) = self.next_loc()?;
675 let value = &self.rows[ridx].value()[cidx];
676
677 match value {
678 Value::Null => None,
679 Value::String(x) => {
680 Some(NaiveTime::parse_from_str(x, "%H:%M:%S%.f").map_err(|_| {
681 anyhow!(
682 "Trino cannot parse Time at position: ({}, {}): {:?}",
683 ridx,
684 cidx,
685 value
686 )
687 })?)
688 }
689 _ => throw!(anyhow!(
690 "Trino unknown value at position: ({}, {}): {:?}",
691 ridx,
692 cidx,
693 value
694 )),
695 }
696 }
697}
698
699impl<'r, 'a> Produce<'r, NaiveDate> for TrinoSourcePartitionParser<'a> {
700 type Error = TrinoSourceError;
701
702 #[throws(TrinoSourceError)]
703 fn produce(&'r mut self) -> NaiveDate {
704 let (ridx, cidx) = self.next_loc()?;
705 let value = &self.rows[ridx].value()[cidx];
706
707 match value {
708 Value::String(x) => NaiveDate::parse_from_str(x, "%Y-%m-%d").map_err(|_| {
709 anyhow!(
710 "Trino cannot parse Date at position: ({}, {}): {:?}",
711 ridx,
712 cidx,
713 value
714 )
715 })?,
716 _ => throw!(anyhow!(
717 "Trino unknown value at position: ({}, {}): {:?}",
718 ridx,
719 cidx,
720 value
721 )),
722 }
723 }
724}
725
726impl<'r, 'a> Produce<'r, Option<NaiveDate>> for TrinoSourcePartitionParser<'a> {
727 type Error = TrinoSourceError;
728
729 #[throws(TrinoSourceError)]
730 fn produce(&'r mut self) -> Option<NaiveDate> {
731 let (ridx, cidx) = self.next_loc()?;
732 let value = &self.rows[ridx].value()[cidx];
733
734 match value {
735 Value::Null => None,
736 Value::String(x) => Some(NaiveDate::parse_from_str(x, "%Y-%m-%d").map_err(|_| {
737 anyhow!(
738 "Trino cannot parse Date at position: ({}, {}): {:?}",
739 ridx,
740 cidx,
741 value
742 )
743 })?),
744 _ => throw!(anyhow!(
745 "Trino unknown value at position: ({}, {}): {:?}",
746 ridx,
747 cidx,
748 value
749 )),
750 }
751 }
752}
753
754#[cfg(test)]
755mod tests {
756 use super::*;
757
758 #[test]
759 fn test_new_with_all_params() {
760 let rt = Arc::new(Runtime::new().unwrap());
761 let conn = "trino+https://myuser:mypass@localhost:8443/mycatalog?\
762 source=connectorx-test&schema=myschema&\
763 client_tags=tag1,tag2,tag3&client_info=test-run&\
764 trace_token=abc123&\
765 session.query_max_execution_time=10m&\
766 extra_credential.token=secret123&verify=false";
767 assert!(TrinoSource::new(rt, conn).is_ok());
768 }
769
770 #[test]
771 fn test_new_minimal_url() {
772 let rt = Arc::new(Runtime::new().unwrap());
773 assert!(TrinoSource::new(rt, "trino://test@localhost:8080/memory").is_ok());
774 }
775
776 #[test]
777 fn test_new_ignores_empty_keys() {
778 let rt = Arc::new(Runtime::new().unwrap());
779 let conn = "trino://test@localhost:8080/memory?session.=val&extra_credential.=x";
780 assert!(TrinoSource::new(rt, conn).is_ok());
781 }
782}