Merge pull request #857 from kongscn/specify-schema-oracle

feat(oracle): support schema query parameter in connection URL

Author: wangxiaoying

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);

connectorx/src/sources/oracle/mod.rs

@@ -50,6 +50,7 @@ pub struct OracleSource {
     queries: Vec<CXQuery<String>>,
     names: Vec<String>,
     schema: Vec<OracleTypeSystem>,
+    current_schema: Option<String>,
 }
 
 #[throws(OracleSourceError)]
@@ -86,14 +87,25 @@ impl OracleSource {
             .max_size(nconn as u32)
             .build(manager)?;
 
+        let params: HashMap<String, String> = 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<OracleConn, OracleSourceError> {
+        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::<usize>(count_query(&cxq, &OracleDialect {})?.as_str(), &[])?;
@@ -192,8 +204,8 @@ where
     #[throws(OracleSourceError)]
     fn partition(self) -> Vec<Self::Partition> {
         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

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                           |
 |:-------------------------:|:---------------------------:|:----------------------------------:|