connectorx 0.4.6

Load data from databases to dataframes, the fastest way.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

use crate::constants::CONNECTORX_PROTOCOL;
use crate::errors::{ConnectorXError, Result};
use crate::utils::remove_query_params;
use anyhow::anyhow;
use fehler::throws;
#[cfg(feature = "src_postgres")]
use redshift_iam::redshift_to_postgres;
use std::convert::TryFrom;
use url::Url;

#[derive(Debug, Clone)]
pub enum SourceType {
    Postgres,
    SQLite,
    MySQL,
    MsSQL,
    Oracle,
    BigQuery,
    DuckDB,
    Trino,
    ClickHouse,
    Unknown,
}

#[derive(Debug, Clone)]
pub struct SourceConn {
    pub ty: SourceType,
    pub conn: Url,
    pub proto: String,
}

impl TryFrom<&str> for SourceConn {
    type Error = ConnectorXError;

    fn try_from(conn: &str) -> Result<SourceConn> {
        let old_url = Url::parse(conn).map_err(|e| anyhow!("parse error: {}", e))?;

        // parse connectorx protocol
        let proto = match old_url.query_pairs().find(|p| p.0 == CONNECTORX_PROTOCOL) {
            Some((_, proto)) => proto.to_owned().to_string(),
            None => "binary".to_string(),
        };

        // create url by removing connectorx protocol
        let url = remove_query_params(&old_url, &[CONNECTORX_PROTOCOL]);

        // users from sqlalchemy may set engine in connection url (e.g. mssql+pymssql://...)
        // only for compatablility, we don't use the same engine
        match url.scheme().split('+').collect::<Vec<&str>>()[0] {
            "postgres" | "postgresql" => Ok(SourceConn::new(SourceType::Postgres, url, proto)),
            #[cfg(feature = "src_postgres")]
            "redshift-iam" => Ok(SourceConn::new(
                SourceType::Postgres,
                redshift_to_postgres(url),
                "cursor".to_string(),
            )),
            "sqlite" => Ok(SourceConn::new(SourceType::SQLite, url, proto)),
            "mysql" => Ok(SourceConn::new(SourceType::MySQL, url, proto)),
            "mssql" => Ok(SourceConn::new(SourceType::MsSQL, url, proto)),
            "oracle" => Ok(SourceConn::new(SourceType::Oracle, url, proto)),
            "bigquery" => Ok(SourceConn::new(SourceType::BigQuery, url, proto)),
            "duckdb" => Ok(SourceConn::new(SourceType::DuckDB, url, proto)),
            "trino" => Ok(SourceConn::new(SourceType::Trino, url, proto)),
            "clickhouse" => Ok(SourceConn::new(SourceType::ClickHouse, url, proto)),
            _ => Ok(SourceConn::new(SourceType::Unknown, url, proto)),
        }
    }
}

impl SourceConn {
    pub fn new(ty: SourceType, conn: Url, proto: String) -> Self {
        Self { ty, conn, proto }
    }
    pub fn set_protocol(&mut self, protocol: &str) {
        self.proto = protocol.to_string();
    }
}

#[throws(ConnectorXError)]
pub fn parse_source(conn: &str, protocol: Option<&str>) -> SourceConn {
    let mut source_conn = SourceConn::try_from(conn)?;
    match protocol {
        Some(p) => source_conn.set_protocol(p),
        None => {}
    }
    source_conn
}

#[cfg(test)]
mod tests {
    use super::SourceConn;
    use std::convert::TryFrom;

    /// Removing the connectorx protocol must not re-encode the remaining parameters:
    /// sources that percent-decode the query would otherwise receive a literal `+`
    /// wherever the caller wrote a space.
    #[test]
    fn keeps_remaining_query_params_verbatim() {
        let source_conn = SourceConn::try_from(
            "postgresql://u:p@host:5432/db?options=-c%20statement_timeout%3D1s&cxprotocol=cursor",
        )
        .unwrap();

        assert_eq!(
            source_conn.conn.query(),
            Some("options=-c%20statement_timeout%3D1s")
        );
        assert_eq!(source_conn.proto, "cursor");
    }

    /// The query is left alone even when there is no protocol parameter to remove.
    #[test]
    fn leaves_the_query_alone_when_there_is_nothing_to_remove() {
        let source_conn = SourceConn::try_from("mysql://host:3306/db?a=x%20y&b=%2Fz").unwrap();

        assert_eq!(source_conn.conn.query(), Some("a=x%20y&b=%2Fz"));
        assert_eq!(source_conn.proto, "binary");
    }
}