connectorx/sources/mssql/
tiberius_impl.rs

1//! Tiberius-backed implementation of the MsSQL source.
2//!
3//! This is the original, default implementation (Phase 0/1 of the migration
4//! plan tracked against sfu-db/connector-x#942). Compiled only when the
5//! `src_mssql_tiberius` feature is active.
6
7use super::driver;
8use super::errors::MsSQLSourceError;
9use super::typesystem::{FloatN, IntN, MsSQLTypeSystem};
10use crate::constants::DB_BUFFER_SIZE;
11use crate::{
12    data_order::DataOrder,
13    errors::ConnectorXError,
14    sources::{PartitionParser, Produce, Source, SourcePartition},
15    sql::{count_query, CXQuery},
16    utils::DummyBox,
17};
18use anyhow::anyhow;
19use bb8_07::{Pool, PooledConnection};
20use bb8_tiberius::ConnectionManager;
21use chrono::{DateTime, Utc};
22use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
23use fehler::{throw, throws};
24use futures::StreamExt;
25use log::debug;
26use owning_ref::OwningHandle;
27use rust_decimal::Decimal;
28use sqlparser::dialect::MsSqlDialect;
29use std::collections::HashMap;
30use std::sync::Arc;
31use tiberius::{AuthMethod, Config, EncryptionLevel, QueryItem, QueryStream, Row};
32use tokio::runtime::{Handle, Runtime};
33use url::Url;
34use urlencoding::decode;
35use uuid_old::Uuid;
36
37type Conn<'a> = PooledConnection<'a, ConnectionManager>;
38pub struct MsSQLSource {
39    rt: Arc<Runtime>,
40    pool: Pool<ConnectionManager>,
41    origin_query: Option<String>,
42    queries: Vec<CXQuery<String>>,
43    names: Vec<String>,
44    schema: Vec<MsSQLTypeSystem>,
45}
46
47#[throws(MsSQLSourceError)]
48pub fn mssql_config(url: &Url) -> Config {
49    let mut config = Config::new();
50
51    let host = decode(url.host_str().unwrap_or("localhost"))?.into_owned();
52    let hosts: Vec<&str> = host.split('\\').collect();
53    match hosts.len() {
54        1 => config.host(host),
55        2 => {
56            // SQL Server support instance name: `server\instance:port`
57            config.host(hosts[0]);
58            config.instance_name(hosts[1]);
59        }
60        _ => throw!(anyhow!("MsSQL hostname parse error: {}", host)),
61    }
62    config.port(url.port().unwrap_or(1433));
63    // remove the leading "/"
64    config.database(decode(&url.path()[1..])?.into_owned());
65    // Using SQL Server authentication.
66    #[allow(unused)]
67    let params: HashMap<String, String> = url.query_pairs().into_owned().collect();
68    #[cfg(any(windows, feature = "integrated-auth-gssapi"))]
69    match params.get("trusted_connection") {
70        // pefer trusted_connection if set to true
71        Some(v) if v == "true" => {
72            debug!("mssql auth through trusted connection!");
73            config.authentication(AuthMethod::Integrated);
74        }
75        _ => {
76            debug!("mssql auth through sqlserver authentication");
77            config.authentication(AuthMethod::sql_server(
78                decode(url.username())?.into_owned(),
79                decode(url.password().unwrap_or(""))?.into_owned(),
80            ));
81        }
82    };
83    #[cfg(all(not(windows), not(feature = "integrated-auth-gssapi")))]
84    config.authentication(AuthMethod::sql_server(
85        decode(url.username())?.into_owned(),
86        decode(url.password().unwrap_or(""))?.into_owned(),
87    ));
88
89    match params.get("trust_server_certificate") {
90        Some(v) if v.to_lowercase() == "true" => config.trust_cert(),
91        _ => {}
92    };
93
94    if let Some(v) = params.get("trust_server_certificate_ca") {
95        config.trust_cert_ca(v);
96    }
97
98    match params.get("encrypt") {
99        Some(v) if v.to_lowercase() == "true" => config.encryption(EncryptionLevel::Required),
100        Some(v) if v.to_lowercase() == "false" => config.encryption(EncryptionLevel::Off),
101        _ => config.encryption(EncryptionLevel::NotSupported),
102    };
103
104    if let Some(appname) = params.get("appname") {
105        config.application_name(decode(appname)?.into_owned());
106    }
107
108    config
109}
110
111impl MsSQLSource {
112    #[throws(MsSQLSourceError)]
113    pub fn new(rt: Arc<Runtime>, conn: &str, nconn: usize) -> Self {
114        debug!("mssql source using driver: {:?}", driver::active_driver());
115        let url = Url::parse(conn)?;
116        let config = mssql_config(&url)?;
117        let manager = bb8_tiberius::ConnectionManager::new(config);
118        let pool = rt.block_on(Pool::builder().max_size(nconn as u32).build(manager))?;
119
120        Self {
121            rt,
122            pool,
123            origin_query: None,
124            queries: vec![],
125            names: vec![],
126            schema: vec![],
127        }
128    }
129}
130
131impl Source for MsSQLSource
132where
133    MsSQLSourcePartition: SourcePartition<TypeSystem = MsSQLTypeSystem, Error = MsSQLSourceError>,
134{
135    const DATA_ORDERS: &'static [DataOrder] = &[DataOrder::RowMajor];
136    type Partition = MsSQLSourcePartition;
137    type TypeSystem = MsSQLTypeSystem;
138    type Error = MsSQLSourceError;
139
140    #[throws(MsSQLSourceError)]
141    fn set_data_order(&mut self, data_order: DataOrder) {
142        if !matches!(data_order, DataOrder::RowMajor) {
143            throw!(ConnectorXError::UnsupportedDataOrder(data_order));
144        }
145    }
146
147    fn set_queries<Q: ToString>(&mut self, queries: &[CXQuery<Q>]) {
148        self.queries = queries.iter().map(|q| q.map(Q::to_string)).collect();
149    }
150
151    fn set_origin_query(&mut self, query: Option<String>) {
152        self.origin_query = query;
153    }
154
155    #[throws(MsSQLSourceError)]
156    fn fetch_metadata(&mut self) {
157        assert!(!self.queries.is_empty());
158
159        let mut conn = self.rt.block_on(self.pool.get())?;
160        let first_query = &self.queries[0];
161        let (names, types) = match self.rt.block_on(conn.query(first_query.as_str(), &[])) {
162            Ok(mut stream) => match self.rt.block_on(async { stream.columns().await }) {
163                Ok(Some(columns)) => columns
164                    .iter()
165                    .map(|col| {
166                        (
167                            col.name().to_string(),
168                            MsSQLTypeSystem::from(&col.column_type()),
169                        )
170                    })
171                    .unzip(),
172                Ok(None) => {
173                    throw!(anyhow!(
174                        "MsSQL returned no columns for query: {}",
175                        first_query
176                    ));
177                }
178                Err(e) => {
179                    throw!(anyhow!("Error fetching columns: {}", e));
180                }
181            },
182            Err(e) => {
183                debug!(
184                    "cannot get metadata for '{}', try next query: {}",
185                    first_query, e
186                );
187                throw!(e);
188            }
189        };
190
191        self.names = names;
192        self.schema = types;
193    }
194
195    #[throws(MsSQLSourceError)]
196    fn result_rows(&mut self) -> Option<usize> {
197        match &self.origin_query {
198            Some(q) => {
199                let cxq = CXQuery::Naked(q.clone());
200                let cquery = count_query(&cxq, &MsSqlDialect {})?;
201                let mut conn = self.rt.block_on(self.pool.get())?;
202
203                let stream = self.rt.block_on(conn.query(cquery.as_str(), &[]))?;
204                let row = self
205                    .rt
206                    .block_on(stream.into_row())?
207                    .ok_or_else(|| anyhow!("MsSQL failed to get the count of query: {}", q))?;
208
209                let row: i32 = row.get(0).ok_or(MsSQLSourceError::GetNRowsFailed)?; // the count in mssql is i32
210                Some(row as usize)
211            }
212            None => None,
213        }
214    }
215
216    fn names(&self) -> Vec<String> {
217        self.names.clone()
218    }
219
220    fn schema(&self) -> Vec<Self::TypeSystem> {
221        self.schema.clone()
222    }
223
224    #[throws(MsSQLSourceError)]
225    fn partition(self) -> Vec<Self::Partition> {
226        let mut ret = vec![];
227        for query in self.queries {
228            ret.push(MsSQLSourcePartition::new(
229                self.pool.clone(),
230                self.rt.clone(),
231                &query,
232                &self.schema,
233            ));
234        }
235        ret
236    }
237}
238
239pub struct MsSQLSourcePartition {
240    pool: Pool<ConnectionManager>,
241    rt: Arc<Runtime>,
242    query: CXQuery<String>,
243    schema: Vec<MsSQLTypeSystem>,
244    nrows: usize,
245    ncols: usize,
246}
247
248impl MsSQLSourcePartition {
249    pub fn new(
250        pool: Pool<ConnectionManager>,
251        handle: Arc<Runtime>,
252        query: &CXQuery<String>,
253        schema: &[MsSQLTypeSystem],
254    ) -> Self {
255        Self {
256            rt: handle,
257            pool,
258            query: query.clone(),
259            schema: schema.to_vec(),
260            nrows: 0,
261            ncols: schema.len(),
262        }
263    }
264}
265
266impl SourcePartition for MsSQLSourcePartition {
267    type TypeSystem = MsSQLTypeSystem;
268    type Parser<'a> = MsSQLSourceParser<'a>;
269    type Error = MsSQLSourceError;
270
271    #[throws(MsSQLSourceError)]
272    fn result_rows(&mut self) {
273        let cquery = count_query(&self.query, &MsSqlDialect {})?;
274        let mut conn = self.rt.block_on(self.pool.get())?;
275
276        let stream = self.rt.block_on(conn.query(cquery.as_str(), &[]))?;
277        let row = self
278            .rt
279            .block_on(stream.into_row())?
280            .ok_or_else(|| anyhow!("MsSQL failed to get the count of query: {}", self.query))?;
281
282        let row: i32 = row.get(0).ok_or(MsSQLSourceError::GetNRowsFailed)?; // the count in mssql is i32
283        self.nrows = row as usize;
284    }
285
286    #[throws(MsSQLSourceError)]
287    fn parser<'a>(&'a mut self) -> Self::Parser<'a> {
288        let conn = self.rt.block_on(self.pool.get())?;
289        let rows: OwningHandle<Box<Conn<'a>>, DummyBox<QueryStream<'a>>> =
290            OwningHandle::new_with_fn(Box::new(conn), |conn: *const Conn<'a>| unsafe {
291                let conn = &mut *(conn as *mut Conn<'a>);
292
293                DummyBox(
294                    self.rt
295                        .block_on(conn.query(self.query.as_str(), &[]))
296                        .unwrap(),
297                )
298            });
299
300        MsSQLSourceParser::new(self.rt.handle(), rows, &self.schema)
301    }
302
303    fn nrows(&self) -> usize {
304        self.nrows
305    }
306
307    fn ncols(&self) -> usize {
308        self.ncols
309    }
310}
311
312pub struct MsSQLSourceParser<'a> {
313    rt: &'a Handle,
314    iter: OwningHandle<Box<Conn<'a>>, DummyBox<QueryStream<'a>>>,
315    rowbuf: Vec<Row>,
316    ncols: usize,
317    current_col: usize,
318    current_row: usize,
319    is_finished: bool,
320}
321
322impl<'a> MsSQLSourceParser<'a> {
323    fn new(
324        rt: &'a Handle,
325        iter: OwningHandle<Box<Conn<'a>>, DummyBox<QueryStream<'a>>>,
326        schema: &[MsSQLTypeSystem],
327    ) -> Self {
328        Self {
329            rt,
330            iter,
331            rowbuf: Vec::with_capacity(DB_BUFFER_SIZE),
332            ncols: schema.len(),
333            current_row: 0,
334            current_col: 0,
335            is_finished: false,
336        }
337    }
338
339    #[throws(MsSQLSourceError)]
340    fn next_loc(&mut self) -> (usize, usize) {
341        let ret = (self.current_row, self.current_col);
342        self.current_row += (self.current_col + 1) / self.ncols;
343        self.current_col = (self.current_col + 1) % self.ncols;
344        ret
345    }
346}
347
348impl<'a> PartitionParser<'a> for MsSQLSourceParser<'a> {
349    type TypeSystem = MsSQLTypeSystem;
350    type Error = MsSQLSourceError;
351
352    #[throws(MsSQLSourceError)]
353    fn fetch_next(&mut self) -> (usize, bool) {
354        assert!(self.current_col == 0);
355        let remaining_rows = self.rowbuf.len() - self.current_row;
356        if remaining_rows > 0 {
357            return (remaining_rows, self.is_finished);
358        } else if self.is_finished {
359            return (0, self.is_finished);
360        }
361
362        if !self.rowbuf.is_empty() {
363            self.rowbuf.drain(..);
364        }
365
366        for _ in 0..DB_BUFFER_SIZE {
367            if let Some(item) = self.rt.block_on(self.iter.next()) {
368                match item.map_err(MsSQLSourceError::MsSQLError)? {
369                    QueryItem::Row(row) => self.rowbuf.push(row),
370                    _ => continue,
371                }
372            } else {
373                self.is_finished = true;
374                break;
375            }
376        }
377        self.current_row = 0;
378        self.current_col = 0;
379        (self.rowbuf.len(), self.is_finished)
380    }
381}
382
383macro_rules! impl_produce {
384    ($($t: ty,)+) => {
385        $(
386            impl<'r, 'a> Produce<'r, $t> for MsSQLSourceParser<'a> {
387                type Error = MsSQLSourceError;
388
389                #[throws(MsSQLSourceError)]
390                fn produce(&'r mut self) -> $t {
391                    let (ridx, cidx) = self.next_loc()?;
392                    let res = self.rowbuf[ridx].get(cidx).ok_or_else(|| anyhow!("MsSQL get None at position: ({}, {})", ridx, cidx))?;
393                    res
394                }
395            }
396
397            impl<'r, 'a> Produce<'r, Option<$t>> for MsSQLSourceParser<'a> {
398                type Error = MsSQLSourceError;
399
400                #[throws(MsSQLSourceError)]
401                fn produce(&'r mut self) -> Option<$t> {
402                    let (ridx, cidx) = self.next_loc()?;
403                    let res = self.rowbuf[ridx].get(cidx);
404                    res
405                }
406            }
407        )+
408    };
409}
410
411impl_produce!(
412    u8,
413    i16,
414    i32,
415    i64,
416    IntN,
417    f32,
418    f64,
419    FloatN,
420    bool,
421    &'r str,
422    &'r [u8],
423    Uuid,
424    Decimal,
425    NaiveDateTime,
426    NaiveDate,
427    NaiveTime,
428    DateTime<Utc>,
429);
connectorx/sources/mssql/tiberius_impl.rs

connectorx/sources/mssql/
tiberius_impl.rs
