Skip to main content

connectorx/sources/oracle/
mod.rs

1mod errors;
2mod typesystem;
3
4use std::collections::HashMap;
5
6pub use self::errors::OracleSourceError;
7pub use self::typesystem::OracleTypeSystem;
8use crate::constants::{DB_BUFFER_SIZE, ORACLE_ARRAY_SIZE};
9use crate::{
10    data_order::DataOrder,
11    errors::ConnectorXError,
12    sources::{PartitionParser, Produce, Source, SourcePartition},
13    sql::{count_query, limit0_query_oracle, CXQuery},
14    utils::DummyBox,
15};
16use chrono::{DateTime, NaiveDate, NaiveDateTime, Utc};
17use fehler::{throw, throws};
18use log::debug;
19use owning_ref::OwningHandle;
20use r2d2::{Pool, PooledConnection};
21use r2d2_oracle::oracle::ResultSet;
22use r2d2_oracle::{
23    oracle::{Connector, Row, Statement, Version},
24    OracleConnectionManager,
25};
26use rust_decimal::Decimal;
27use sqlparser::dialect::Dialect;
28use url::Url;
29use urlencoding::decode;
30
31type OracleManager = OracleConnectionManager;
32type OracleConn = PooledConnection<OracleManager>;
33
34#[derive(Debug)]
35pub struct OracleDialect {}
36
37// implementation copy from AnsiDialect
38impl Dialect for OracleDialect {
39    fn is_identifier_start(&self, ch: char) -> bool {
40        ch.is_ascii_lowercase() || ch.is_ascii_uppercase()
41    }
42
43    fn is_identifier_part(&self, ch: char) -> bool {
44        ch.is_ascii_lowercase() || ch.is_ascii_uppercase() || ch.is_ascii_digit() || ch == '_'
45    }
46}
47
48pub struct OracleSource {
49    pool: Pool<OracleManager>,
50    origin_query: Option<String>,
51    queries: Vec<CXQuery<String>>,
52    names: Vec<String>,
53    schema: Vec<OracleTypeSystem>,
54    current_schema: Option<String>,
55}
56
57#[throws(OracleSourceError)]
58pub fn connect_oracle(conn: &Url) -> Connector {
59    let user = decode(conn.username())?.into_owned();
60    let password = decode(conn.password().unwrap_or(""))?.into_owned();
61    let host = decode(conn.host_str().unwrap_or("localhost"))?.into_owned();
62
63    let params: HashMap<String, String> = conn.query_pairs().into_owned().collect();
64
65    let conn_str = if params.get("alias").is_some_and(|v| v == "true") {
66        host.clone()
67    } else {
68        let port = conn.port().unwrap_or(1521);
69        let path = decode(conn.path())?.into_owned();
70        format!("//{}:{}{}", host, port, path)
71    };
72
73    let mut connector = oracle::Connector::new(user.as_str(), password.as_str(), conn_str.as_str());
74    if user.is_empty() && password.is_empty() {
75        debug!("No username or password provided, assuming system auth.");
76        connector.external_auth(true);
77    }
78    connector
79}
80
81impl OracleSource {
82    #[throws(OracleSourceError)]
83    pub fn new(conn: &str, nconn: usize) -> Self {
84        let conn = Url::parse(conn)?;
85        let connector = connect_oracle(&conn)?;
86        let manager = OracleConnectionManager::from_connector(connector);
87        let pool = r2d2::Pool::builder()
88            .max_size(nconn as u32)
89            .build(manager)?;
90
91        let params: HashMap<String, String> = conn.query_pairs().into_owned().collect();
92        let current_schema = params.get("schema").cloned();
93
94        Self {
95            pool,
96            origin_query: None,
97            queries: vec![],
98            names: vec![],
99            schema: vec![],
100            current_schema,
101        }
102    }
103    pub fn get_conn(&self) -> Result<OracleConn, OracleSourceError> {
104        let conn = self.pool.get()?;
105        if let Some(schema) = &self.current_schema {
106            conn.set_current_schema(schema)?;
107        }
108        Ok(conn)
109    }
110}
111
112/// Oldest Instant Client known to read native `BOOLEAN` columns correctly. On the first execution
113/// of a statement on a connection, Instant Client 23.4 can return `TRUE` for a `FALSE` value that
114/// directly follows another `FALSE`, without raising an error. Clients older than 23 report
115/// `BOOLEAN` columns as `NUMBER`, so they never reach this check.
116const MIN_BOOLEAN_CLIENT_VERSION: Version = Version::new(23, 5, 0, 0, 0);
117
118fn client_reads_boolean_correctly(client: &Version) -> bool {
119    *client >= MIN_BOOLEAN_CLIENT_VERSION
120}
121
122/// Refuses to read a native `BOOLEAN` column through an Instant Client that may return wrong
123/// values for it.
124#[throws(OracleSourceError)]
125fn check_boolean_client(names: &[String], schema: &[OracleTypeSystem]) {
126    if let Some(i) = schema
127        .iter()
128        .position(|ty| matches!(ty, OracleTypeSystem::Boolean(_)))
129    {
130        let client = Version::client()?;
131        if !client_reads_boolean_correctly(&client) {
132            throw!(OracleSourceError::UnreliableBooleanClient {
133                column: names[i].clone(),
134                client: client.to_string(),
135            });
136        }
137    }
138}
139
140impl Source for OracleSource
141where
142    OracleSourcePartition:
143        SourcePartition<TypeSystem = OracleTypeSystem, Error = OracleSourceError>,
144{
145    const DATA_ORDERS: &'static [DataOrder] = &[DataOrder::RowMajor];
146    type Partition = OracleSourcePartition;
147    type TypeSystem = OracleTypeSystem;
148    type Error = OracleSourceError;
149
150    #[throws(OracleSourceError)]
151    fn set_data_order(&mut self, data_order: DataOrder) {
152        if !matches!(data_order, DataOrder::RowMajor) {
153            throw!(ConnectorXError::UnsupportedDataOrder(data_order));
154        }
155    }
156
157    fn set_queries<Q: ToString>(&mut self, queries: &[CXQuery<Q>]) {
158        self.queries = queries.iter().map(|q| q.map(Q::to_string)).collect();
159    }
160
161    fn set_origin_query(&mut self, query: Option<String>) {
162        self.origin_query = query;
163    }
164
165    #[throws(OracleSourceError)]
166    fn fetch_metadata(&mut self) {
167        assert!(!self.queries.is_empty());
168
169        let conn = self.get_conn()?;
170        for (i, query) in self.queries.iter().enumerate() {
171            // assuming all the partition queries yield same schema
172            // without rownum = 1, derived type might be wrong
173            // example: select avg(test_int), test_char from test_table group by test_char
174            // -> (NumInt, Char) instead of (NumtFloat, Char)
175            match conn.query(limit0_query_oracle(query)?.as_str(), &[]) {
176                Ok(rows) => {
177                    let (names, types): (Vec<String>, Vec<OracleTypeSystem>) = rows
178                        .column_info()
179                        .iter()
180                        .map(|col| {
181                            (
182                                col.name().to_string(),
183                                OracleTypeSystem::from(col.oracle_type()),
184                            )
185                        })
186                        .unzip();
187                    check_boolean_client(&names, &types)?;
188                    self.names = names;
189                    self.schema = types;
190                    return;
191                }
192                Err(e) if i == self.queries.len() - 1 => {
193                    // tried the last query but still get an error
194                    debug!("cannot get metadata for '{}': {}", query, e);
195                    throw!(e);
196                }
197                Err(_) => {}
198            }
199        }
200        // tried all queries but all get empty result set
201        let iter = conn.query(self.queries[0].as_str(), &[])?;
202        let (names, types) = iter
203            .column_info()
204            .iter()
205            .map(|col| (col.name().to_string(), OracleTypeSystem::VarChar(false)))
206            .unzip();
207        self.names = names;
208        self.schema = types;
209    }
210
211    #[throws(OracleSourceError)]
212    fn result_rows(&mut self) -> Option<usize> {
213        match &self.origin_query {
214            Some(q) => {
215                let cxq = CXQuery::Naked(q.clone());
216                let conn = self.get_conn()?;
217
218                let nrows = conn
219                    .query_row_as::<usize>(count_query(&cxq, &OracleDialect {})?.as_str(), &[])?;
220                Some(nrows)
221            }
222            None => None,
223        }
224    }
225
226    fn names(&self) -> Vec<String> {
227        self.names.clone()
228    }
229
230    fn schema(&self) -> Vec<Self::TypeSystem> {
231        self.schema.clone()
232    }
233
234    #[throws(OracleSourceError)]
235    fn partition(self) -> Vec<Self::Partition> {
236        let mut ret = vec![];
237        for query in &self.queries {
238            let conn = self.get_conn()?;
239            ret.push(OracleSourcePartition::new(conn, query, &self.schema));
240        }
241        ret
242    }
243}
244
245pub struct OracleSourcePartition {
246    conn: OracleConn,
247    query: CXQuery<String>,
248    schema: Vec<OracleTypeSystem>,
249    nrows: usize,
250    ncols: usize,
251}
252
253impl OracleSourcePartition {
254    pub fn new(conn: OracleConn, query: &CXQuery<String>, schema: &[OracleTypeSystem]) -> Self {
255        Self {
256            conn,
257            query: query.clone(),
258            schema: schema.to_vec(),
259            nrows: 0,
260            ncols: schema.len(),
261        }
262    }
263}
264
265impl SourcePartition for OracleSourcePartition {
266    type TypeSystem = OracleTypeSystem;
267    type Parser<'a> = OracleTextSourceParser<'a>;
268    type Error = OracleSourceError;
269
270    #[throws(OracleSourceError)]
271    fn result_rows(&mut self) {
272        self.nrows = self
273            .conn
274            .query_row_as::<usize>(count_query(&self.query, &OracleDialect {})?.as_str(), &[])?;
275    }
276
277    #[throws(OracleSourceError)]
278    fn parser(&mut self) -> Self::Parser<'_> {
279        let query = self.query.clone();
280
281        // let iter = self.conn.query(query.as_str(), &[])?;
282        OracleTextSourceParser::new(&self.conn, query.as_str(), &self.schema)?
283    }
284
285    fn nrows(&self) -> usize {
286        self.nrows
287    }
288
289    fn ncols(&self) -> usize {
290        self.ncols
291    }
292}
293
294unsafe impl<'a> Send for OracleTextSourceParser<'a> {}
295
296pub struct OracleTextSourceParser<'a> {
297    rows: OwningHandle<Box<Statement>, DummyBox<ResultSet<'a, Row>>>,
298    rowbuf: Vec<Row>,
299    ncols: usize,
300    current_col: usize,
301    current_row: usize,
302    is_finished: bool,
303}
304
305impl<'a> OracleTextSourceParser<'a> {
306    #[throws(OracleSourceError)]
307    pub fn new(conn: &'a OracleConn, query: &str, schema: &[OracleTypeSystem]) -> Self {
308        let stmt = conn
309            .statement(query)
310            .prefetch_rows(ORACLE_ARRAY_SIZE)
311            .fetch_array_size(ORACLE_ARRAY_SIZE)
312            .build()?;
313        let rows: OwningHandle<Box<Statement>, DummyBox<ResultSet<'a, Row>>> =
314            OwningHandle::new_with_fn(Box::new(stmt), |stmt: *const Statement| unsafe {
315                DummyBox((*(stmt as *mut Statement)).query(&[]).unwrap())
316            });
317
318        Self {
319            rows,
320            rowbuf: Vec::with_capacity(DB_BUFFER_SIZE),
321            ncols: schema.len(),
322            current_row: 0,
323            current_col: 0,
324            is_finished: false,
325        }
326    }
327
328    #[throws(OracleSourceError)]
329    fn next_loc(&mut self) -> (usize, usize) {
330        let ret = (self.current_row, self.current_col);
331        self.current_row += (self.current_col + 1) / self.ncols;
332        self.current_col = (self.current_col + 1) % self.ncols;
333        ret
334    }
335}
336
337impl<'a> PartitionParser<'a> for OracleTextSourceParser<'a> {
338    type TypeSystem = OracleTypeSystem;
339    type Error = OracleSourceError;
340
341    #[throws(OracleSourceError)]
342    fn fetch_next(&mut self) -> (usize, bool) {
343        assert!(self.current_col == 0);
344        let remaining_rows = self.rowbuf.len() - self.current_row;
345        if remaining_rows > 0 {
346            return (remaining_rows, self.is_finished);
347        } else if self.is_finished {
348            return (0, self.is_finished);
349        }
350
351        if !self.rowbuf.is_empty() {
352            self.rowbuf.drain(..);
353        }
354        for _ in 0..DB_BUFFER_SIZE {
355            if let Some(item) = (*self.rows).next() {
356                self.rowbuf.push(item?);
357            } else {
358                self.is_finished = true;
359                break;
360            }
361        }
362        self.current_row = 0;
363        self.current_col = 0;
364        (self.rowbuf.len(), self.is_finished)
365    }
366}
367
368macro_rules! impl_produce_text {
369    ($($t: ty,)+) => {
370        $(
371            impl<'r, 'a> Produce<'r, $t> for OracleTextSourceParser<'a> {
372                type Error = OracleSourceError;
373
374                #[throws(OracleSourceError)]
375                fn produce(&'r mut self) -> $t {
376                    let (ridx, cidx) = self.next_loc()?;
377                    let res = self.rowbuf[ridx].get(cidx)?;
378                    res
379                }
380            }
381
382            impl<'r, 'a> Produce<'r, Option<$t>> for OracleTextSourceParser<'a> {
383                type Error = OracleSourceError;
384
385                #[throws(OracleSourceError)]
386                fn produce(&'r mut self) -> Option<$t> {
387                    let (ridx, cidx) = self.next_loc()?;
388                    let res = self.rowbuf[ridx].get(cidx)?;
389                    res
390                }
391            }
392        )+
393    };
394}
395
396impl_produce_text!(
397    i64,
398    bool,
399    f64,
400    String,
401    NaiveDate,
402    NaiveDateTime,
403    DateTime<Utc>,
404    Vec<u8>,
405);
406
407// Manual implementation for Decimal since Oracle doesn't support it directly via FromSql
408impl<'r, 'a> Produce<'r, Decimal> for OracleTextSourceParser<'a> {
409    type Error = OracleSourceError;
410
411    #[throws(OracleSourceError)]
412    fn produce(&'r mut self) -> Decimal {
413        let (ridx, cidx) = self.next_loc()?;
414        let s: String = self.rowbuf[ridx].get(cidx)?;
415        let res = s.parse::<Decimal>()?;
416        res
417    }
418}
419
420impl<'r, 'a> Produce<'r, Option<Decimal>> for OracleTextSourceParser<'a> {
421    type Error = OracleSourceError;
422
423    #[throws(OracleSourceError)]
424    fn produce(&'r mut self) -> Option<Decimal> {
425        let (ridx, cidx) = self.next_loc()?;
426        let s: Option<String> = self.rowbuf[ridx].get(cidx)?;
427        match s {
428            Some(val) => Some(val.parse::<Decimal>()?),
429            None => None,
430        }
431    }
432}
433
434#[cfg(test)]
435mod tests {
436    use super::*;
437
438    #[test]
439    fn boolean_client_threshold() {
440        // versions as reported by `Version::client()` for the Instant Client zips
441        assert!(!client_reads_boolean_correctly(&Version::new(
442            23, 4, 0, 24, 5
443        )));
444        assert!(client_reads_boolean_correctly(&Version::new(
445            23, 5, 0, 24, 7
446        )));
447        assert!(client_reads_boolean_correctly(&Version::new(
448            23, 26, 3, 0, 0
449        )));
450    }
451
452    #[test]
453    fn check_boolean_client_skips_queries_without_boolean() {
454        // no Boolean column -> the client version is never looked up
455        let names = vec!["A".to_string()];
456        let schema = vec![OracleTypeSystem::NumInt(true)];
457        assert!(check_boolean_client(&names, &schema).is_ok());
458    }
459}