diff --git a/Cargo.lock b/Cargo.lock index 016c8552b9..ed278c1a7a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1675,6 +1675,7 @@ dependencies = [ "redis", "reqwest 0.12.28", "russh", + "rust-gaussdb", "rust_decimal", "rustls 0.23.40", "serde", @@ -1854,7 +1855,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -1987,7 +1988,7 @@ dependencies = [ "arrow", "cast", "comfy-table", - "fallible-iterator", + "fallible-iterator 0.3.0", "fallible-streaming-iterator", "hashlink", "libduckdb-sys", @@ -2188,7 +2189,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2213,6 +2214,12 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "fallible-iterator" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" + [[package]] name = "fallible-iterator" version = "0.3.0" @@ -3180,7 +3187,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.62.2", + "windows-core 0.61.2", ] [[package]] @@ -3929,6 +3936,16 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "md-5" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" +dependencies = [ + "cfg-if", + "digest 0.11.3", +] + [[package]] name = "md5" version = "0.7.0" @@ -4071,7 +4088,7 @@ dependencies = [ "hickory-resolver", "hmac 0.12.1", "macro_magic", - "md-5", + "md-5 0.10.6", "mongocrypt", "mongodb-internal-macros", "pbkdf2 0.12.2", @@ -4128,7 +4145,7 @@ dependencies = [ "png 0.18.1", "serde", "thiserror 2.0.18", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -4219,7 +4236,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -4847,7 +4864,7 @@ dependencies = [ "hmac 0.12.1", "hostname", "indexmap 2.14.0", - "md-5", + "md-5 0.10.6", "pbkdf2 0.12.2", "pkcs8 0.10.2", "rand 0.8.6", @@ -4880,7 +4897,7 @@ dependencies = [ "hmac 0.12.1", "hostname", "indexmap 2.14.0", - "md-5", + "md-5 0.10.6", "pbkdf2 0.12.2", "pkcs8 0.10.2", "rand 0.8.6", @@ -4915,7 +4932,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7d8fae84b431384b68627d0f9b3b1245fcf9f46f6c0e3dc902e9dce64edd1967" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.45.0", ] [[package]] @@ -5419,6 +5436,38 @@ dependencies = [ "rand 0.8.6", ] +[[package]] +name = "postgres-protocol" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56201207dac53e2f38e848e31b4b91616a6bb6e0c7205b77718994a7f49e70fc" +dependencies = [ + "base64 0.22.1", + "byteorder", + "bytes", + "fallible-iterator 0.2.0", + "hmac 0.13.0", + "md-5 0.11.0", + "memchr", + "rand 0.10.1", + "sha2 0.11.0", + "stringprep", +] + +[[package]] +name = "postgres-types" +version = "0.2.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8dc729a129e682e8d24170cd30ae1aa01b336b096cbb56df6d534ffec133d186" +dependencies = [ + "bytes", + "chrono", + "fallible-iterator 0.2.0", + "postgres-protocol", + "serde_core", + "serde_json", +] + [[package]] name = "potential_utf" version = "0.1.5" @@ -6189,6 +6238,30 @@ dependencies = [ "wasm-bindgen-futures", ] +[[package]] +name = "rust-gaussdb" +version = "0.1.0" +source = "git+https://github.com/t8y2/rust-gaussdb?branch=main#3cf091477233bda925dca690cef96edb65be5309" +dependencies = [ + "base64 0.22.1", + "byteorder", + "bytes", + "fallible-iterator 0.2.0", + "futures-channel", + "futures-util", + "hex", + "hmac 0.12.1", + "log", + "md-5 0.10.6", + "phf 0.11.3", + "postgres-protocol", + "postgres-types", + "rand 0.9.4", + "sha2 0.10.9", + "tokio", + "tokio-util", +] + [[package]] name = "rust_decimal" version = "1.41.0" @@ -6275,7 +6348,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -6376,7 +6449,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -7028,7 +7101,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -7225,7 +7298,7 @@ dependencies = [ "hmac 0.12.1", "itoa", "log", - "md-5", + "md-5 0.10.6", "memchr", "once_cell", "percent-encoding", @@ -7267,7 +7340,7 @@ dependencies = [ "home", "itoa", "log", - "md-5", + "md-5 0.10.6", "memchr", "once_cell", "rand 0.8.6", @@ -7953,7 +8026,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix 1.1.4", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -8452,7 +8525,7 @@ dependencies = [ "png 0.18.1", "serde", "thiserror 2.0.18", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -9064,7 +9137,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.48.0", ] [[package]] diff --git a/crates/dbx-core/Cargo.toml b/crates/dbx-core/Cargo.toml index 8eaecd0b61..cceef481dd 100644 --- a/crates/dbx-core/Cargo.toml +++ b/crates/dbx-core/Cargo.toml @@ -28,3 +28,4 @@ portpicker = "0.1.1" csv = "1" calamine = "0.30.1" odbc-api = "25" +rust-gaussdb = { git = "https://github.com/t8y2/rust-gaussdb", branch = "main" } diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 26c1a46221..45eeb7c6d1 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -29,7 +29,7 @@ pub enum PoolKind { Oracle(Arc>), Elasticsearch(db::elasticsearch_driver::EsClient), Dameng(Arc>), - Gaussdb(Arc>), + Gaussdb(Arc>), } pub struct AppState { @@ -188,7 +188,7 @@ impl AppState { &db_config.password, ) .await?; - PoolKind::Gaussdb(Arc::new(std::sync::Mutex::new(client))) + PoolKind::Gaussdb(Arc::new(tokio::sync::Mutex::new(client))) } }; diff --git a/crates/dbx-core/src/db/gaussdb_driver.rs b/crates/dbx-core/src/db/gaussdb_driver.rs index b0e6c36f37..bf833b51b8 100644 --- a/crates/dbx-core/src/db/gaussdb_driver.rs +++ b/crates/dbx-core/src/db/gaussdb_driver.rs @@ -1,103 +1,67 @@ use std::time::Instant; -use odbc_api::{buffers::TextRowSet, ConnectionOptions, Cursor, ResultSetMetadata}; - use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo}; use super::CONNECTION_TIMEOUT_SECS; pub struct GaussdbClient { - conn: odbc_api::Connection<'static>, + client: rust_gaussdb::Client, } unsafe impl Send for GaussdbClient {} impl GaussdbClient { - pub fn query_rows(&self, sql: &str) -> Result>, String> { - match self.conn.execute(sql, (), None).map_err(|e| e.to_string())? { - Some(cursor) => read_cursor(cursor), - None => Ok(vec![]), + async fn query_rows(&mut self, sql: &str) -> Result>>, String> { + let rows = self.client.query(sql, &[]).await.map_err(|e| e.to_string())?; + let mut result = Vec::new(); + for row in &rows { + let mut vals = Vec::new(); + for i in 0..row.columns().len() { + vals.push(row.try_get::(i)); + } + result.push(vals); } + Ok(result) } - fn query_single_column(&self, sql: &str) -> Result, String> { - Ok(self.query_rows(sql)?.into_iter().filter_map(|r| r.into_iter().next()).collect()) + async fn query_single_column(&mut self, sql: &str) -> Result, String> { + Ok(self.query_rows(sql).await?.into_iter().filter_map(|r| r.into_iter().next().flatten()).collect()) } } -fn read_cursor(cursor: impl Cursor) -> Result>, String> { - let mut cursor = cursor; - let col_count = cursor.num_result_cols().map_err(|e| e.to_string())? as u16; - let buffer = TextRowSet::for_cursor(1000, &mut cursor, Some(8192)).map_err(|e| e.to_string())?; - let mut row_cursor = cursor.bind_buffer(buffer).map_err(|e| e.to_string())?; - let mut rows = Vec::new(); - while let Some(batch) = row_cursor.fetch().map_err(|e| e.to_string())? { - for row_idx in 0..batch.num_rows() { - let vals: Vec = (0..col_count as usize) - .map(|col| { - batch.at(col, row_idx).and_then(|bytes| std::str::from_utf8(bytes).ok()).unwrap_or("").to_string() - }) - .collect(); - rows.push(vals); - } - } - Ok(rows) -} - pub async fn connect(host: &str, port: u16, database: &str, user: &str, pass: &str) -> Result { - let conn_str = format!( - "Driver={{GaussDBA}};Servername={host};Port={port};Database={db};UID={user};PWD={pass}", - host = host, - port = port, - db = database, - user = user, - pass = pass, - ); + let dsn = format!("host={host} port={port} user={user} password={pass} dbname={database}"); - let result = tokio::time::timeout( - std::time::Duration::from_secs(CONNECTION_TIMEOUT_SECS), - tokio::task::spawn_blocking(move || { - super::ODBC_ENV.connect_with_connection_string(&conn_str, ConnectionOptions::default()) - .map_err(|e| { - let msg = e.to_string(); - if msg.contains("Data source name not found") || msg.contains("Can't open lib") { - format!( - "GaussDB ODBC driver not found. Please install the GaussDB ODBC driver \ - and register it in odbcinst.ini (Linux/macOS) or the ODBC Data Source Administrator (Windows). \ - Original error: {msg}" - ) - } else { - format!("GaussDB connection failed: {msg}") - } - }) - .map(|conn| GaussdbClient { conn }) - }), - ) + let result = tokio::time::timeout(std::time::Duration::from_secs(CONNECTION_TIMEOUT_SECS), async { + rust_gaussdb::Client::connect(&dsn).await.map_err(|e| format!("GaussDB connection failed: {e}")) + }) .await - .map_err(|_| format!("GaussDB connection timed out ({CONNECTION_TIMEOUT_SECS}s)"))? - .map_err(|e| format!("GaussDB connection task failed: {e}"))?; + .map_err(|_| format!("GaussDB connection timed out ({CONNECTION_TIMEOUT_SECS}s)"))?; - result + result.map(|client| GaussdbClient { client }) } -pub fn list_databases(client: &GaussdbClient) -> Result, String> { - let rows = - client.query_single_column("SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname")?; +pub async fn list_databases(client: &mut GaussdbClient) -> Result, String> { + let rows = client + .query_single_column("SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname") + .await?; Ok(rows.into_iter().map(|name| DatabaseInfo { name }).collect()) } -pub fn list_schemas(client: &GaussdbClient) -> Result, String> { - client.query_single_column( - "SELECT DISTINCT nspname FROM pg_catalog.pg_namespace n \ - WHERE nspname NOT LIKE 'pg_%' \ - AND nspname NOT IN ('information_schema', 'cstore', 'snapshot', 'db4ai', 'dbe_perf', \ - 'dbe_pldebugger', 'dbe_pldeveloper', 'pkg_service', 'pkg_util', 'sqladvisor', 'blockchain') \ - AND EXISTS (SELECT 1 FROM pg_catalog.pg_class c WHERE c.relnamespace = n.oid AND c.relkind IN ('r', 'v', 'm')) \ - ORDER BY nspname", - ) +pub async fn list_schemas(client: &mut GaussdbClient) -> Result, String> { + client + .query_single_column( + "SELECT DISTINCT nspname FROM pg_catalog.pg_namespace n \ + WHERE nspname NOT LIKE 'pg_%' \ + AND nspname NOT IN ('information_schema', 'cstore', 'snapshot', 'db4ai', 'dbe_perf', \ + 'dbe_pldebugger', 'dbe_pldeveloper', 'pkg_service', 'pkg_util', 'sqladvisor', 'blockchain') \ + AND EXISTS (SELECT 1 FROM pg_catalog.pg_class c WHERE c.relnamespace = n.oid AND c.relkind IN ('r', 'v', 'm')) \ + ORDER BY nspname", + ) + .await } -pub fn list_tables(client: &GaussdbClient, schema: &str) -> Result, String> { +pub async fn list_tables(client: &mut GaussdbClient, schema: &str) -> Result, String> { let s = schema.replace('\'', "''"); let sql = format!( "SELECT c.relname, CASE c.relkind WHEN 'r' THEN 'TABLE' WHEN 'v' THEN 'VIEW' WHEN 'm' THEN 'VIEW' ELSE 'TABLE' END \ @@ -106,62 +70,66 @@ pub fn list_tables(client: &GaussdbClient, schema: &str) -> Result Result, String> { +pub async fn get_columns(client: &mut GaussdbClient, schema: &str, table: &str) -> Result, String> { let s = schema.replace('\'', "''"); let t = table.replace('\'', "''"); - let pk_rows = client.query_single_column(&format!( - "SELECT a.attname FROM pg_catalog.pg_index i \ - JOIN pg_catalog.pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) \ - WHERE i.indrelid = (SELECT oid FROM pg_catalog.pg_class WHERE relname = '{t}' \ - AND relnamespace = (SELECT oid FROM pg_catalog.pg_namespace WHERE nspname = '{s}')) \ - AND i.indisprimary" - ))?; + let pk_rows = client + .query_single_column(&format!( + "SELECT a.attname FROM pg_catalog.pg_index i \ + JOIN pg_catalog.pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) \ + WHERE i.indrelid = (SELECT oid FROM pg_catalog.pg_class WHERE relname = '{t}' \ + AND relnamespace = (SELECT oid FROM pg_catalog.pg_namespace WHERE nspname = '{s}')) \ + AND i.indisprimary" + )) + .await?; let pk_names: std::collections::HashSet = pk_rows.into_iter().collect(); - let col_rows = client.query_rows(&format!( - "SELECT a.attname, format_type(a.atttypid, a.atttypmod), \ - CASE WHEN a.attnotnull THEN 'NO' ELSE 'YES' END, \ - pg_catalog.pg_get_expr(d.adbin, d.adrelid), \ - CASE WHEN t.typname IN ('numeric', 'float4', 'float8') THEN COALESCE(((a.atttypmod - 4) >> 16) & 65535, -1) ELSE NULL END, \ - CASE WHEN t.typname = 'numeric' THEN COALESCE((a.atttypmod - 4) & 65535, -1) ELSE NULL END, \ - CASE WHEN a.atttypmod > 0 AND t.typname IN ('varchar', 'bpchar') THEN a.atttypmod - 4 ELSE NULL END \ - FROM pg_catalog.pg_attribute a \ - JOIN pg_catalog.pg_type t ON t.oid = a.atttypid \ - LEFT JOIN pg_catalog.pg_attrdef d ON d.adrelid = a.attrelid AND d.adnum = a.attnum \ - WHERE a.attrelid = (SELECT oid FROM pg_catalog.pg_class WHERE relname = '{t}' \ - AND relnamespace = (SELECT oid FROM pg_catalog.pg_namespace WHERE nspname = '{s}')) \ - AND a.attnum > 0 AND NOT a.attisdropped \ - ORDER BY a.attnum" - ))?; + let col_rows = client + .query_rows(&format!( + "SELECT a.attname, format_type(a.atttypid, a.atttypmod), \ + CASE WHEN a.attnotnull THEN 'NO' ELSE 'YES' END, \ + pg_catalog.pg_get_expr(d.adbin, d.adrelid), \ + CASE WHEN t.typname IN ('numeric', 'float4', 'float8') THEN COALESCE(((a.atttypmod - 4) >> 16) & 65535, -1) ELSE NULL END, \ + CASE WHEN t.typname = 'numeric' THEN COALESCE((a.atttypmod - 4) & 65535, -1) ELSE NULL END, \ + CASE WHEN a.atttypmod > 0 AND t.typname IN ('varchar', 'bpchar') THEN a.atttypmod - 4 ELSE NULL END \ + FROM pg_catalog.pg_attribute a \ + JOIN pg_catalog.pg_type t ON t.oid = a.atttypid \ + LEFT JOIN pg_catalog.pg_attrdef d ON d.adrelid = a.attrelid AND d.adnum = a.attnum \ + WHERE a.attrelid = (SELECT oid FROM pg_catalog.pg_class WHERE relname = '{t}' \ + AND relnamespace = (SELECT oid FROM pg_catalog.pg_namespace WHERE nspname = '{s}')) \ + AND a.attnum > 0 AND NOT a.attisdropped \ + ORDER BY a.attnum" + )) + .await?; Ok(col_rows .into_iter() .map(|r| { - let name = r.first().cloned().unwrap_or_default(); - let data_type = r.get(1).cloned().unwrap_or_default(); - let num_prec = r.get(4).and_then(|v| v.parse::().ok()); - let num_scale = r.get(5).and_then(|v| v.parse::().ok()); - let char_len = r.get(6).and_then(|v| v.parse::().ok()); + let name = r.first().and_then(|v| v.clone()).unwrap_or_default(); + let data_type = r.get(1).and_then(|v| v.clone()).unwrap_or_default(); + let num_prec = r.get(4).and_then(|v| v.as_ref()?.parse::().ok()); + let num_scale = r.get(5).and_then(|v| v.as_ref()?.parse::().ok()); + let char_len = r.get(6).and_then(|v| v.as_ref()?.parse::().ok()); ColumnInfo { is_primary_key: pk_names.contains(&name), name, data_type, - is_nullable: r.get(2).map(|v| v == "YES").unwrap_or(false), - column_default: r.get(3).filter(|v| !v.is_empty()).cloned(), + is_nullable: r.get(2).and_then(|v| v.as_ref()).map(|v| v == "YES").unwrap_or(false), + column_default: r.get(3).and_then(|v| v.clone()).filter(|v| !v.is_empty()), extra: None, comment: None, numeric_precision: num_prec, @@ -172,18 +140,18 @@ pub fn get_columns(client: &GaussdbClient, schema: &str, table: &str) -> Result< .collect()) } -pub fn list_indexes(client: &GaussdbClient, schema: &str, table: &str) -> Result, String> { +pub async fn list_indexes(client: &mut GaussdbClient, schema: &str, table: &str) -> Result, String> { let s = schema.replace('\'', "''"); let t = table.replace('\'', "''"); let sql = format!( "SELECT indexname, indexdef FROM pg_indexes WHERE schemaname = '{s}' AND tablename = '{t}' ORDER BY indexname" ); - let rows = client.query_rows(&sql)?; + let rows = client.query_rows(&sql).await?; Ok(rows .into_iter() .map(|r| { - let name = r.first().cloned().unwrap_or_default(); - let def = r.get(1).cloned().unwrap_or_default(); + let name = r.first().and_then(|v| v.clone()).unwrap_or_default(); + let def = r.get(1).and_then(|v| v.clone()).unwrap_or_default(); let is_unique = def.to_uppercase().contains("UNIQUE"); let is_primary = def.to_uppercase().contains("PRIMARY"); let columns = def @@ -205,7 +173,11 @@ pub fn list_indexes(client: &GaussdbClient, schema: &str, table: &str) -> Result .collect()) } -pub fn list_foreign_keys(client: &GaussdbClient, schema: &str, table: &str) -> Result, String> { +pub async fn list_foreign_keys( + client: &mut GaussdbClient, + schema: &str, + table: &str, +) -> Result, String> { let s = schema.replace('\'', "''"); let t = table.replace('\'', "''"); let sql = format!( @@ -219,48 +191,44 @@ pub fn list_foreign_keys(client: &GaussdbClient, schema: &str, table: &str) -> R WHERE con.contype = 'f' AND n.nspname = '{s}' AND cl.relname = '{t}' \ ORDER BY con.conname" ); - let rows = client.query_rows(&sql)?; + let rows = client.query_rows(&sql).await?; Ok(rows .into_iter() .map(|r| ForeignKeyInfo { - name: r.first().cloned().unwrap_or_default(), - column: r.get(1).cloned().unwrap_or_default(), - ref_table: r.get(2).cloned().unwrap_or_default(), - ref_column: r.get(3).cloned().unwrap_or_default(), + name: r.first().and_then(|v| v.clone()).unwrap_or_default(), + column: r.get(1).and_then(|v| v.clone()).unwrap_or_default(), + ref_table: r.get(2).and_then(|v| v.clone()).unwrap_or_default(), + ref_column: r.get(3).and_then(|v| v.clone()).unwrap_or_default(), }) .collect()) } -pub fn list_triggers(client: &GaussdbClient, schema: &str, table: &str) -> Result, String> { +pub async fn list_triggers(client: &mut GaussdbClient, schema: &str, table: &str) -> Result, String> { let s = schema.replace('\'', "''"); let t = table.replace('\'', "''"); let sql = format!( - "SELECT t.tgname, em.event, CASE WHEN t.tgtype & 2 = 2 THEN 'BEFORE' ELSE 'AFTER' END \ + "SELECT t.tgname, \ + CASE WHEN t.tgtype & 4 = 4 THEN 'INSERT' WHEN t.tgtype & 8 = 8 THEN 'DELETE' \ + WHEN t.tgtype & 16 = 16 THEN 'UPDATE' ELSE 'UNKNOWN' END, \ + CASE WHEN t.tgtype & 2 = 2 THEN 'BEFORE' ELSE 'AFTER' END \ FROM pg_catalog.pg_trigger t \ JOIN pg_catalog.pg_class c ON c.oid = t.tgrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ - CROSS JOIN LATERAL ( \ - SELECT CASE \ - WHEN t.tgtype & 4 = 4 THEN 'INSERT' \ - WHEN t.tgtype & 8 = 8 THEN 'DELETE' \ - WHEN t.tgtype & 16 = 16 THEN 'UPDATE' \ - ELSE 'UNKNOWN' END AS event \ - ) em \ WHERE NOT t.tgisinternal AND n.nspname = '{s}' AND c.relname = '{t}' \ ORDER BY t.tgname" ); - let rows = client.query_rows(&sql)?; + let rows = client.query_rows(&sql).await?; Ok(rows .into_iter() .map(|r| TriggerInfo { - name: r.first().cloned().unwrap_or_default(), - event: r.get(1).cloned().unwrap_or_default(), - timing: r.get(2).cloned().unwrap_or_default(), + name: r.first().and_then(|v| v.clone()).unwrap_or_default(), + event: r.get(1).and_then(|v| v.clone()).unwrap_or_default(), + timing: r.get(2).and_then(|v| v.clone()).unwrap_or_default(), }) .collect()) } -pub fn execute_query_sync(client: &GaussdbClient, sql: &str) -> Result { +pub async fn execute_query(client: &mut GaussdbClient, sql: &str) -> Result { let start = Instant::now(); let sql = sql.trim().trim_end_matches(';'); let trimmed = sql.to_uppercase(); @@ -268,75 +236,43 @@ pub fn execute_query_sync(client: &GaussdbClient, sql: &str) -> Result { - let col_count = cursor.num_result_cols().map_err(|e| e.to_string())? as u16; - let columns: Vec = (1..=col_count) - .map(|i| cursor.col_name(i).map_err(|e| e.to_string()).unwrap_or_else(|_| format!("col{i}"))) - .collect(); + let rows = client.client.query(sql, &[]).await.map_err(|e| e.to_string())?; - let buffer = TextRowSet::for_cursor(1000, &mut cursor, Some(8192)).map_err(|e| e.to_string())?; - let mut row_cursor = cursor.bind_buffer(buffer).map_err(|e| e.to_string())?; + let columns: Vec = if let Some(first) = rows.first() { + first.columns().iter().map(|c| c.name.clone()).collect() + } else { + Vec::new() + }; - let mut rows = Vec::new(); - while let Some(batch) = row_cursor.fetch().map_err(|e| e.to_string())? { - for row_idx in 0..batch.num_rows() { - let vals: Vec = (0..col_count as usize) - .map(|col| { - batch - .at(col, row_idx) - .and_then(|bytes| std::str::from_utf8(bytes).ok()) - .map(|s| serde_json::Value::String(s.to_string())) - .unwrap_or(serde_json::Value::Null) - }) - .collect(); - rows.push(vals); - if rows.len() >= crate::query::MAX_ROWS { - break; - } - } - if rows.len() >= crate::query::MAX_ROWS { - break; - } - } - - let truncated = rows.len() >= crate::query::MAX_ROWS; - Ok(QueryResult { - columns, - rows, - affected_rows: 0, - execution_time_ms: start.elapsed().as_millis(), - truncated, - }) + let mut result_rows = Vec::new(); + for row in &rows { + let vals: Vec = (0..row.columns().len()) + .map(|i| row.try_get::(i).map(serde_json::Value::String).unwrap_or(serde_json::Value::Null)) + .collect(); + result_rows.push(vals); + if result_rows.len() >= crate::query::MAX_ROWS { + break; } - None => Ok(QueryResult { - columns: vec![], - rows: vec![], - affected_rows: 0, - execution_time_ms: start.elapsed().as_millis(), - truncated: false, - }), } + + let truncated = result_rows.len() >= crate::query::MAX_ROWS; + Ok(QueryResult { + columns, + rows: result_rows, + affected_rows: 0, + execution_time_ms: start.elapsed().as_millis(), + truncated, + }) } else { - match client.conn.execute(sql, (), None) { - Ok(Some(_cursor)) => Ok(QueryResult { - columns: vec![], - rows: vec![], - affected_rows: 0, - execution_time_ms: start.elapsed().as_millis(), - truncated: false, - }), - Ok(None) => Ok(QueryResult { - columns: vec![], - rows: vec![], - affected_rows: 0, - execution_time_ms: start.elapsed().as_millis(), - truncated: false, - }), - Err(e) => Err(e.to_string()), - } + let affected = client.client.execute(sql).await.map_err(|e| e.to_string())?; + Ok(QueryResult { + columns: vec![], + rows: vec![], + affected_rows: affected, + execution_time_ms: start.elapsed().as_millis(), + truncated: false, + }) } } diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index f4d2156688..ece96748b3 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -233,11 +233,8 @@ pub async fn do_execute( let sql = sql.to_string(); drop(connections); wait_for_query(cancel_token, async move { - let task = tokio::task::spawn_blocking(move || { - let client = client.lock().map_err(|e| e.to_string())?; - db::gaussdb_driver::execute_query_sync(&client, &sql) - }); - task.await.map_err(|e| e.to_string())? + let mut client = client.lock().await; + db::gaussdb_driver::execute_query(&mut client, &sql).await }) .await .map(truncate_result) diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 8e0f6d16ff..591481ef2e 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -113,7 +113,7 @@ pub fn extract_dameng( pub fn extract_gaussdb( connections: &HashMap, key: &str, -) -> Option>> { +) -> Option>> { match connections.get(key)? { PoolKind::Gaussdb(client) => Some(client.clone()), _ => None, @@ -144,8 +144,8 @@ pub async fn list_databases_core(state: &AppState, connection_id: &str) -> Resul } if let Some(client) = extract_gaussdb(&connections, connection_id) { drop(connections); - let client = client.lock().map_err(|e| e.to_string())?; - return db::gaussdb_driver::list_databases(&client); + let mut client = client.lock().await; + return db::gaussdb_driver::list_databases(&mut client).await; } } @@ -183,8 +183,8 @@ pub async fn list_schemas_core(state: &AppState, connection_id: &str, database: } if let Some(client) = extract_gaussdb(&connections, &pool_key) { drop(connections); - let client = client.lock().map_err(|e| e.to_string())?; - return db::gaussdb_driver::list_schemas(&client); + let mut client = client.lock().await; + return db::gaussdb_driver::list_schemas(&mut client).await; } } @@ -233,8 +233,8 @@ pub async fn list_tables_core( } if let Some(client) = extract_gaussdb(&connections, &pool_key) { drop(connections); - let client = client.lock().map_err(|e| e.to_string())?; - return db::gaussdb_driver::list_tables(&client, schema); + let mut client = client.lock().await; + return db::gaussdb_driver::list_tables(&mut client, schema).await; } } @@ -286,8 +286,8 @@ pub async fn get_columns_core( } if let Some(client) = extract_gaussdb(&connections, &pool_key) { drop(connections); - let client = client.lock().map_err(|e| e.to_string())?; - return db::gaussdb_driver::get_columns(&client, schema, table); + let mut client = client.lock().await; + return db::gaussdb_driver::get_columns(&mut client, schema, table).await; } } @@ -330,8 +330,8 @@ pub async fn list_indexes_core( } if let Some(client) = extract_gaussdb(&connections, &pool_key) { drop(connections); - let client = client.lock().map_err(|e| e.to_string())?; - return db::gaussdb_driver::list_indexes(&client, schema, table); + let mut client = client.lock().await; + return db::gaussdb_driver::list_indexes(&mut client, schema, table).await; } } @@ -374,8 +374,8 @@ pub async fn list_foreign_keys_core( } if let Some(client) = extract_gaussdb(&connections, &pool_key) { drop(connections); - let client = client.lock().map_err(|e| e.to_string())?; - return db::gaussdb_driver::list_foreign_keys(&client, schema, table); + let mut client = client.lock().await; + return db::gaussdb_driver::list_foreign_keys(&mut client, schema, table).await; } } @@ -418,8 +418,8 @@ pub async fn list_triggers_core( } if let Some(client) = extract_gaussdb(&connections, &pool_key) { drop(connections); - let client = client.lock().map_err(|e| e.to_string())?; - return db::gaussdb_driver::list_triggers(&client, schema, table); + let mut client = client.lock().await; + return db::gaussdb_driver::list_triggers(&mut client, schema, table).await; } } @@ -488,8 +488,8 @@ pub async fn get_table_ddl_core( } if let Some(client) = extract_gaussdb(&connections, &pool_key) { drop(connections); - let client = client.lock().map_err(|e| e.to_string())?; - return build_gaussdb_ddl(&client, schema, table); + let mut client = client.lock().await; + return build_gaussdb_ddl(&mut client, schema, table).await; } } @@ -746,14 +746,14 @@ pub fn build_dameng_ddl(client: &db::dm_driver::DmClient, schema: &str, table: & Ok(ddl) } -pub fn build_gaussdb_ddl( - client: &db::gaussdb_driver::GaussdbClient, +pub async fn build_gaussdb_ddl( + client: &mut db::gaussdb_driver::GaussdbClient, schema: &str, table: &str, ) -> Result { - let columns = db::gaussdb_driver::get_columns(client, schema, table)?; - let indexes = db::gaussdb_driver::list_indexes(client, schema, table)?; - let fkeys = db::gaussdb_driver::list_foreign_keys(client, schema, table)?; + let columns = db::gaussdb_driver::get_columns(client, schema, table).await?; + let indexes = db::gaussdb_driver::list_indexes(client, schema, table).await?; + let fkeys = db::gaussdb_driver::list_foreign_keys(client, schema, table).await?; let mut ddl = format!("CREATE TABLE \"{schema}\".\"{table}\" (\n"); let col_lines: Vec = columns diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 8a0deca96f..4668328c11 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -217,7 +217,7 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi &config.password, ) .await?; - PoolKind::Gaussdb(std::sync::Arc::new(std::sync::Mutex::new(client))) + PoolKind::Gaussdb(std::sync::Arc::new(tokio::sync::Mutex::new(client))) } };