github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-29 16:18:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(oracle): speed up initial table loading

Browse Source
This commit is contained in:
t8y2
2026-05-10 20:47:54 +08:00
parent 61bd7480f4
commit d917842fca
7 changed files with 158 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/connection.rs
+46 -4
View File
@@ -1,5 +1,6 @@
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::sync::Mutex;
@@ -38,13 +39,53 @@ pub enum PoolKind {
MongoDb(mongodb::Client),
ClickHouse(db::clickhouse_driver::ChClient),
SqlServer(Arc<tokio::sync::Mutex<db::sqlserver::SqlServerClient>>),
Oracle(Arc<tokio::sync::Mutex<db::oracle_driver::OracleClient>>),
Oracle(Arc<OraclePool>),
Elasticsearch(db::elasticsearch_driver::EsClient),
Dameng(Arc<std::sync::Mutex<db::dm_driver::DmClient>>),
Gaussdb(Arc<tokio::sync::Mutex<db::gaussdb_driver::GaussdbClient>>),
ExternalDriver { driver_id: String, config: ConnectionConfig, session: Arc<PluginDriverSession> },
}
pub struct OraclePool {
clients: Vec<Arc<tokio::sync::Mutex<db::oracle_driver::OracleClient>>>,
next: AtomicUsize,
}
impl OraclePool {
pub fn new(clients: Vec<db::oracle_driver::OracleClient>) -> Self {
Self {
clients: clients.into_iter().map(|client| Arc::new(tokio::sync::Mutex::new(client))).collect(),
next: AtomicUsize::new(0),
}
}
pub fn client(&self) -> Arc<tokio::sync::Mutex<db::oracle_driver::OracleClient>> {
let index = self.next.fetch_add(1, Ordering::Relaxed) % self.clients.len();
self.clients[index].clone()
}
pub fn primary(&self) -> Arc<tokio::sync::Mutex<db::oracle_driver::OracleClient>> {
self.clients[0].clone()
}
}
async fn connect_oracle_pool(
host: &str,
port: u16,
service: &str,
user: &str,
pass: &str,
sysdba: bool,
) -> Result<OraclePool, String> {
let (first, second, third) = tokio::try_join!(
db::oracle_driver::connect(host, port, service, user, pass, sysdba),
db::oracle_driver::connect(host, port, service, user, pass, sysdba),
db::oracle_driver::connect(host, port, service, user, pass, sysdba),
)?;
let clients = vec![first, second, third];
Ok(OraclePool::new(clients))
}
pub struct AppState {
pub connections: Mutex<HashMap<String, PoolKind>>,
pub configs: Mutex<HashMap<String, ConnectionConfig>>,
@@ -133,7 +174,8 @@ impl AppState {
let conns = self.connections.lock().await;
if conns.contains_key(&pool_key) {
if let Some(PoolKind::Oracle(client)) = conns.get(&pool_key) {
if let Some(PoolKind::Oracle(pool)) = conns.get(&pool_key) {
let client = pool.primary();
let conn = client.lock().await;
if conn.is_closed() {
drop(conn);
@@ -205,7 +247,7 @@ impl AppState {
PoolKind::SqlServer(Arc::new(tokio::sync::Mutex::new(client)))
}
DatabaseType::Oracle => {
let client = db::oracle_driver::connect(
let pool = connect_oracle_pool(
&host,
port,
db_config.database.as_deref().unwrap_or("ORCL"),
@@ -214,7 +256,7 @@ impl AppState {
db_config.sysdba,
)
.await?;
PoolKind::Oracle(Arc::new(tokio::sync::Mutex::new(client)))
PoolKind::Oracle(Arc::new(pool))
}
DatabaseType::Elasticsearch => {
let client =
crates/dbx-core/src/db/oracle_driver.rs
+68 -1
View File
@@ -295,13 +295,15 @@ pub async fn list_triggers(conn: &OracleClient, schema: &str, table: &str) -> Re
pub async fn execute_query(conn: &OracleClient, sql: &str) -> Result<QueryResult, String> {
let start = Instant::now();
let sql = sql.trim().trim_end_matches(';');
let explicit_limit = explicit_select_row_limit(sql);
// Rewrite FETCH FIRST N ROWS ONLY → ROWNUM for Oracle 11g compatibility.
let sql = rewrite_fetch_first(sql);
if starts_with_executable_sql_keyword(sql.as_ref(), &["SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN"]) {
let capped_sql = cap_select_rows(sql.as_ref());
let result = conn.query_with_limit(capped_sql.as_ref(), &[], ORACLE_QUERY_LIMIT, 500).await.map_err(|e| {
let query_limit = explicit_limit.unwrap_or(ORACLE_QUERY_LIMIT).min(ORACLE_QUERY_LIMIT);
let result = conn.query_with_limit(capped_sql.as_ref(), &[], query_limit, 500).await.map_err(|e| {
log::error!("[oracle] execute_query SELECT failed: {e}");
e.to_string()
})?;
@@ -357,6 +359,55 @@ fn has_for_update_clause(sql: &str) -> bool {
sql.to_uppercase().contains(" FOR UPDATE")
}
fn explicit_select_row_limit(sql: &str) -> Option<usize> {
fetch_first_row_limit(sql).or_else(|| rownum_row_limit(sql))
}
fn fetch_first_row_limit(sql: &str) -> Option<usize> {
let upper = sql.to_uppercase();
let fetch_pos = upper.find("FETCH FIRST").or_else(|| upper.find("FETCH NEXT"))?;
let after_fetch = &upper[fetch_pos..];
let end = after_fetch.find("ROWS ONLY")?;
let keyword_len = if after_fetch.starts_with("FETCH FIRST") { 11 } else { 10 };
sql[fetch_pos + keyword_len..fetch_pos + end].trim().parse::<usize>().ok()
}
fn rownum_row_limit(sql: &str) -> Option<usize> {
let upper = sql.to_uppercase();
let mut rest = upper.as_str();
let mut best: Option<usize> = None;
while let Some(pos) = rest.find("ROWNUM") {
rest = &rest[pos + "ROWNUM".len()..];
let trimmed = rest.trim_start();
let value_start = if let Some(after) = trimmed.strip_prefix("<=") {
after.trim_start()
} else if let Some(after) = trimmed.strip_prefix('<') {
if let Some(n) = parse_leading_usize(after.trim_start()) {
let exclusive = n.saturating_sub(1);
best = Some(best.map_or(exclusive, |current| current.min(exclusive)));
}
continue;
} else {
continue;
};
if let Some(n) = parse_leading_usize(value_start) {
best = Some(best.map_or(n, |current| current.min(n)));
}
}
best
}
fn parse_leading_usize(value: &str) -> Option<usize> {
let digits: String = value.chars().take_while(|ch| ch.is_ascii_digit()).collect();
if digits.is_empty() {
return None;
}
digits.parse().ok()
}
fn rewrite_fetch_first(sql: &str) -> std::borrow::Cow<'_, str> {
let upper = sql.to_uppercase();
// Match: ... [OFFSET M ROWS] FETCH FIRST|NEXT N ROWS ONLY
@@ -420,4 +471,20 @@ mod tests {
"SELECT * FROM (SELECT * FROM users) WHERE ROWNUM <= 20"
);
}
#[test]
fn explicit_select_row_limit_reads_fetch_first() {
assert_eq!(explicit_select_row_limit("SELECT * FROM users FETCH FIRST 100 ROWS ONLY"), Some(100));
assert_eq!(explicit_select_row_limit("SELECT * FROM users OFFSET 20 ROWS FETCH NEXT 50 ROWS ONLY"), Some(50));
}
#[test]
fn explicit_select_row_limit_reads_rownum() {
assert_eq!(explicit_select_row_limit("SELECT * FROM users WHERE ROWNUM <= 100"), Some(100));
assert_eq!(explicit_select_row_limit("SELECT * FROM users WHERE ROWNUM < 101"), Some(100));
assert_eq!(
explicit_select_row_limit("SELECT * FROM (SELECT * FROM users WHERE ROWNUM <= 500) WHERE ROWNUM <= 100"),
Some(100)
);
}
}
crates/dbx-core/src/query.rs
+2 -2
View File
@@ -189,8 +189,8 @@ pub async fn do_execute(
};
wait_for_query(cancel_token, db::sqlserver::execute_query(&mut client, sql)).await.map(truncate_result)
}
PoolKind::Oracle(client) => {
let client = client.clone();
PoolKind::Oracle(pool) => {
let client = pool.client();
drop(connections);
let client = match cancel_token.as_ref() {
Some(token) => tokio::select! {
crates/dbx-core/src/schema.rs
+19 -14
View File
@@ -1,7 +1,7 @@
use std::collections::HashMap;
use std::sync::Arc;
use crate::connection::{AppState, MysqlMode, PoolKind};
use crate::connection::{AppState, MysqlMode, OraclePool, PoolKind};
use crate::db;
pub fn duckdb_query_tables(con: &duckdb::Connection) -> Result<Vec<db::TableInfo>, String> {
@@ -92,12 +92,9 @@ pub fn extract_clickhouse(
}
}
pub fn extract_oracle(
connections: &HashMap<String, PoolKind>,
key: &str,
) -> Option<Arc<tokio::sync::Mutex<db::oracle_driver::OracleClient>>> {
pub fn extract_oracle(connections: &HashMap<String, PoolKind>, key: &str) -> Option<Arc<OraclePool>> {
match connections.get(key)? {
PoolKind::Oracle(client) => Some(client.clone()),
PoolKind::Oracle(pool) => Some(pool.clone()),
_ => None,
}
}
@@ -142,8 +139,9 @@ pub async fn list_databases_core(state: &AppState, connection_id: &str) -> Resul
let mut client = client.lock().await;
return db::sqlserver::list_databases(&mut client).await;
}
if let Some(client) = extract_oracle(&connections, connection_id) {
if let Some(pool) = extract_oracle(&connections, connection_id) {
drop(connections);
let client = pool.client();
let client = client.lock().await;
return db::oracle_driver::list_databases(&*client).await;
}
@@ -195,8 +193,9 @@ pub async fn list_schemas_core(state: &AppState, connection_id: &str, database:
let mut client = client.lock().await;
return db::sqlserver::list_schemas(&mut client).await;
}
if let Some(client) = extract_oracle(&connections, &pool_key) {
if let Some(pool) = extract_oracle(&connections, &pool_key) {
drop(connections);
let client = pool.client();
let client = client.lock().await;
return db::oracle_driver::list_schemas(&*client).await;
}
@@ -256,8 +255,9 @@ pub async fn list_tables_core(
let mut client = client.lock().await;
return db::sqlserver::list_tables(&mut client, schema).await;
}
if let Some(client) = extract_oracle(&connections, &pool_key) {
if let Some(pool) = extract_oracle(&connections, &pool_key) {
drop(connections);
let client = pool.client();
let client = client.lock().await;
return db::oracle_driver::list_tables(&*client, schema).await;
}
@@ -331,8 +331,9 @@ pub async fn get_columns_core(
let mut client = client.lock().await;
return db::sqlserver::get_columns(&mut client, schema, table).await;
}
if let Some(client) = extract_oracle(&connections, &pool_key) {
if let Some(pool) = extract_oracle(&connections, &pool_key) {
drop(connections);
let client = pool.client();
let client = client.lock().await;
return db::oracle_driver::get_columns(&*client, schema, table).await;
}
@@ -381,8 +382,9 @@ pub async fn list_indexes_core(
let mut client = client.lock().await;
return db::sqlserver::list_indexes(&mut client, schema, table).await;
}
if let Some(client) = extract_oracle(&connections, &pool_key) {
if let Some(pool) = extract_oracle(&connections, &pool_key) {
drop(connections);
let client = pool.client();
let client = client.lock().await;
return db::oracle_driver::list_indexes(&*client, schema, table).await;
}
@@ -431,8 +433,9 @@ pub async fn list_foreign_keys_core(
let mut client = client.lock().await;
return db::sqlserver::list_foreign_keys(&mut client, schema, table).await;
}
if let Some(client) = extract_oracle(&connections, &pool_key) {
if let Some(pool) = extract_oracle(&connections, &pool_key) {
drop(connections);
let client = pool.client();
let client = client.lock().await;
return db::oracle_driver::list_foreign_keys(&*client, schema, table).await;
}
@@ -481,8 +484,9 @@ pub async fn list_triggers_core(
let mut client = client.lock().await;
return db::sqlserver::list_triggers(&mut client, schema, table).await;
}
if let Some(client) = extract_oracle(&connections, &pool_key) {
if let Some(pool) = extract_oracle(&connections, &pool_key) {
drop(connections);
let client = pool.client();
let client = client.lock().await;
return db::oracle_driver::list_triggers(&*client, schema, table).await;
}
@@ -557,8 +561,9 @@ pub async fn get_table_ddl_core(
let mut client = client.lock().await;
return build_sqlserver_ddl(&mut client, schema, table).await;
}
if let Some(client) = extract_oracle(&connections, &pool_key) {
if let Some(pool) = extract_oracle(&connections, &pool_key) {
drop(connections);
let client = pool.client();
let client = client.lock().await;
return build_oracle_ddl(&*client, schema, table).await;
}
crates/dbx-core/src/transfer.rs
+4 -4
View File
@@ -501,8 +501,8 @@ pub async fn execute_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Res
let mut client = client.lock().await;
db::sqlserver::execute_query(&mut client, sql).await
}
PoolKind::Oracle(client) => {
let client = client.clone();
PoolKind::Oracle(pool) => {
let client = pool.client();
drop(connections);
let client = client.lock().await;
db::oracle_driver::execute_query(&*client, sql).await
@@ -618,8 +618,8 @@ pub async fn get_columns_for_transfer(
let mut client = client.lock().await;
return db::sqlserver::get_columns(&mut client, &schema, &table).await;
}
if let Some(PoolKind::Oracle(client)) = connections.get(pool_key) {
let client = client.clone();
if let Some(PoolKind::Oracle(pool)) = connections.get(pool_key) {
let client = pool.client();
let schema = schema.to_string();
let table = table.to_string();
drop(connections);