connectorx/
source_router.rs

1use crate::constants::CONNECTORX_PROTOCOL;
2use crate::errors::{ConnectorXError, Result};
3use crate::utils::remove_query_params;
4use anyhow::anyhow;
5use fehler::throws;
6#[cfg(feature = "src_postgres")]
7use redshift_iam::redshift_to_postgres;
8use std::convert::TryFrom;
9use url::Url;
10
11#[derive(Debug, Clone)]
12pub enum SourceType {
13    Postgres,
14    SQLite,
15    MySQL,
16    MsSQL,
17    Oracle,
18    BigQuery,
19    DuckDB,
20    Trino,
21    ClickHouse,
22    Unknown,
23}
24
25#[derive(Debug, Clone)]
26pub struct SourceConn {
27    pub ty: SourceType,
28    pub conn: Url,
29    pub proto: String,
30}
31
32impl TryFrom<&str> for SourceConn {
33    type Error = ConnectorXError;
34
35    fn try_from(conn: &str) -> Result<SourceConn> {
36        let old_url = Url::parse(conn).map_err(|e| anyhow!("parse error: {}", e))?;
37
38        // parse connectorx protocol
39        let proto = match old_url.query_pairs().find(|p| p.0 == CONNECTORX_PROTOCOL) {
40            Some((_, proto)) => proto.to_owned().to_string(),
41            None => "binary".to_string(),
42        };
43
44        // create url by removing connectorx protocol
45        let url = remove_query_params(&old_url, &[CONNECTORX_PROTOCOL]);
46
47        // users from sqlalchemy may set engine in connection url (e.g. mssql+pymssql://...)
48        // only for compatablility, we don't use the same engine
49        match url.scheme().split('+').collect::<Vec<&str>>()[0] {
50            "postgres" | "postgresql" => Ok(SourceConn::new(SourceType::Postgres, url, proto)),
51            #[cfg(feature = "src_postgres")]
52            "redshift-iam" => Ok(SourceConn::new(
53                SourceType::Postgres,
54                redshift_to_postgres(url),
55                "cursor".to_string(),
56            )),
57            "sqlite" => Ok(SourceConn::new(SourceType::SQLite, url, proto)),
58            "mysql" => Ok(SourceConn::new(SourceType::MySQL, url, proto)),
59            "mssql" => Ok(SourceConn::new(SourceType::MsSQL, url, proto)),
60            "oracle" => Ok(SourceConn::new(SourceType::Oracle, url, proto)),
61            "bigquery" => Ok(SourceConn::new(SourceType::BigQuery, url, proto)),
62            "duckdb" => Ok(SourceConn::new(SourceType::DuckDB, url, proto)),
63            "trino" => Ok(SourceConn::new(SourceType::Trino, url, proto)),
64            "clickhouse" => Ok(SourceConn::new(SourceType::ClickHouse, url, proto)),
65            _ => Ok(SourceConn::new(SourceType::Unknown, url, proto)),
66        }
67    }
68}
69
70impl SourceConn {
71    pub fn new(ty: SourceType, conn: Url, proto: String) -> Self {
72        Self { ty, conn, proto }
73    }
74    pub fn set_protocol(&mut self, protocol: &str) {
75        self.proto = protocol.to_string();
76    }
77}
78
79#[throws(ConnectorXError)]
80pub fn parse_source(conn: &str, protocol: Option<&str>) -> SourceConn {
81    let mut source_conn = SourceConn::try_from(conn)?;
82    match protocol {
83        Some(p) => source_conn.set_protocol(p),
84        None => {}
85    }
86    source_conn
87}
88
89#[cfg(test)]
90mod tests {
91    use super::SourceConn;
92    use std::convert::TryFrom;
93
94    /// Removing the connectorx protocol must not re-encode the remaining parameters:
95    /// sources that percent-decode the query would otherwise receive a literal `+`
96    /// wherever the caller wrote a space.
97    #[test]
98    fn keeps_remaining_query_params_verbatim() {
99        let source_conn = SourceConn::try_from(
100            "postgresql://u:p@host:5432/db?options=-c%20statement_timeout%3D1s&cxprotocol=cursor",
101        )
102        .unwrap();
103
104        assert_eq!(
105            source_conn.conn.query(),
106            Some("options=-c%20statement_timeout%3D1s")
107        );
108        assert_eq!(source_conn.proto, "cursor");
109    }
110
111    /// The query is left alone even when there is no protocol parameter to remove.
112    #[test]
113    fn leaves_the_query_alone_when_there_is_nothing_to_remove() {
114        let source_conn = SourceConn::try_from("mysql://host:3306/db?a=x%20y&b=%2Fz").unwrap();
115
116        assert_eq!(source_conn.conn.query(), Some("a=x%20y&b=%2Fz"));
117        assert_eq!(source_conn.proto, "binary");
118    }
119}
connectorx/source_router.rs

connectorx/
source_router.rs
