diff --git a/connectorx/src/partition.rs b/connectorx/src/partition.rs index 3b7493c8c2..566596eb7e 100644 --- a/connectorx/src/partition.rs +++ b/connectorx/src/partition.rs @@ -7,7 +7,7 @@ use crate::sources::mssql::{mssql_config, FloatN, IntN, MsSQLTypeSystem}; #[cfg(feature = "src_mysql")] use crate::sources::mysql::{MySQLSourceError, MySQLTypeSystem}; #[cfg(feature = "src_oracle")] -use crate::sources::oracle::{connect_oracle, OracleDialect}; +use crate::sources::oracle::{connect_oracle, OracleDialect, OracleSource}; #[cfg(feature = "src_postgres")] use crate::sources::postgres::{rewrite_tls_args, PostgresTypeSystem}; #[cfg(feature = "src_trino")] @@ -451,8 +451,8 @@ fn mssql_get_partition_range(conn: &Url, query: &str, col: &str) -> (i64, i64) { #[cfg(feature = "src_oracle")] #[throws(ConnectorXOutError)] fn oracle_get_partition_range(conn: &Url, query: &str, col: &str) -> (i64, i64) { - let connector = connect_oracle(conn)?; - let conn = connector.connect()?; + let source = OracleSource::new(conn.as_str(), 1)?; + let conn = source.get_conn()?; let range_query = get_partition_range_query(query, col, &OracleDialect {})?; let row = conn.query_row(range_query.as_str(), &[])?; let min_v: i64 = row.get(0).unwrap_or(0); diff --git a/connectorx/src/sources/oracle/mod.rs b/connectorx/src/sources/oracle/mod.rs index deb9b7b7fe..2d01484875 100644 --- a/connectorx/src/sources/oracle/mod.rs +++ b/connectorx/src/sources/oracle/mod.rs @@ -50,6 +50,7 @@ pub struct OracleSource { queries: Vec>, names: Vec, schema: Vec, + current_schema: Option, } #[throws(OracleSourceError)] @@ -86,14 +87,25 @@ impl OracleSource { .max_size(nconn as u32) .build(manager)?; + let params: HashMap = conn.query_pairs().into_owned().collect(); + let current_schema = params.get("schema").cloned(); + Self { pool, origin_query: None, queries: vec![], names: vec![], schema: vec![], + current_schema, } } + pub fn get_conn(&self) -> Result { + let conn = self.pool.get()?; + if let Some(schema) = &self.current_schema { + conn.set_current_schema(schema)?; + } + Ok(conn) + } } impl Source for OracleSource @@ -125,7 +137,7 @@ where fn fetch_metadata(&mut self) { assert!(!self.queries.is_empty()); - let conn = self.pool.get()?; + let conn = self.get_conn()?; for (i, query) in self.queries.iter().enumerate() { // assuming all the partition queries yield same schema // without rownum = 1, derived type might be wrong @@ -171,7 +183,7 @@ where match &self.origin_query { Some(q) => { let cxq = CXQuery::Naked(q.clone()); - let conn = self.pool.get()?; + let conn = self.get_conn()?; let nrows = conn .query_row_as::(count_query(&cxq, &OracleDialect {})?.as_str(), &[])?; @@ -192,8 +204,8 @@ where #[throws(OracleSourceError)] fn partition(self) -> Vec { let mut ret = vec![]; - for query in self.queries { - let conn = self.pool.get()?; + for query in &self.queries { + let conn = self.get_conn()?; ret.push(OracleSourcePartition::new(conn, &query, &self.schema)); } ret diff --git a/docs/databases/oracle.md b/docs/databases/oracle.md index bc6ce2b388..631f1ef1e6 100644 --- a/docs/databases/oracle.md +++ b/docs/databases/oracle.md @@ -21,6 +21,18 @@ query = 'SELECT * FROM table' # query string cx.read_sql(conn, query) # read data from Oracle ``` +### Specifying Default Schema + +With version>=0.4.5 you can specify a default schema by adding the `schema` query parameter to the connection URL. This automatically sets the current schema for all connections, eliminating the need to prefix table names with the schema in your queries. + +```py +import connectorx as cx +# Specify schema in the connection URL +conn = 'oracle://username:password@server:port/database?schema=MY_SCHEMA' +query = 'SELECT * FROM table' # No need to use MY_SCHEMA.table +cx.read_sql(conn, query) +``` + ### Oracle-Pandas Type Mapping | Oracle Type | Pandas Type | Comment | |:-------------------------:|:---------------------------:|:----------------------------------:|