connectorx/sources/oracle/
mod.rs1mod 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
37impl 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
112const 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#[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 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 debug!("cannot get metadata for '{}': {}", query, e);
195 throw!(e);
196 }
197 Err(_) => {}
198 }
199 }
200 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 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
407impl<'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 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 let names = vec!["A".to_string()];
456 let schema = vec![OracleTypeSystem::NumInt(true)];
457 assert!(check_boolean_client(&names, &schema).is_ok());
458 }
459}